From 7bd1440710278166715532143ccfc1048a096995 Mon Sep 17 00:00:00 2001 From: Preetha Appan Date: Wed, 9 Jan 2019 12:57:56 -0600 Subject: [PATCH] REfactor statedb factory config to set it directly in client config --- client/client.go | 12 ++++++++---- client/client_test.go | 19 ++++++++++--------- client/config/config.go | 4 ++++ client/testing.go | 13 +------------ command/agent/agent.go | 7 +++++-- 5 files changed, 28 insertions(+), 27 deletions(-) diff --git a/client/client.go b/client/client.go index f6f537009..464663edd 100644 --- a/client/client.go +++ b/client/client.go @@ -256,7 +256,7 @@ var ( ) // NewClient is used to create a new client from the given configuration -func NewClient(cfg *config.Config, consulCatalog consul.CatalogAPI, consulService consulApi.ConsulServiceAPI, stateDBFunc state.NewStateDBFunc) (*Client, error) { +func NewClient(cfg *config.Config, consulCatalog consul.CatalogAPI, consulService consulApi.ConsulServiceAPI) (*Client, error) { // Create the tls wrapper var tlsWrap tlsutil.RegionWrapper if cfg.TLSConfig.EnableRPC { @@ -270,6 +270,10 @@ func NewClient(cfg *config.Config, consulCatalog consul.CatalogAPI, consulServic } } + if cfg.StateDBFactory == nil { + cfg.StateDBFactory = state.GetStateDBFactory(cfg.DevMode) + } + // Create the logger logger := cfg.Logger.ResetNamed("client") @@ -303,7 +307,7 @@ func NewClient(cfg *config.Config, consulCatalog consul.CatalogAPI, consulServic c.servers = servers.New(c.logger, c.shutdownCh, c) // Initialize the client - if err := c.init(stateDBFunc); err != nil { + if err := c.init(); err != nil { return nil, fmt.Errorf("failed to initialize client: %v", err) } @@ -454,7 +458,7 @@ func (c *Client) Ready() <-chan struct{} { // init is used to initialize the client and perform any setup // needed before we begin starting its various components. -func (c *Client) init(statedbFunc state.NewStateDBFunc) error { +func (c *Client) init() error { // Ensure the state dir exists if we have one if c.config.StateDir != "" { if err := os.MkdirAll(c.config.StateDir, 0700); err != nil { @@ -478,7 +482,7 @@ func (c *Client) init(statedbFunc state.NewStateDBFunc) error { c.logger.Info("using state directory", "state_dir", c.config.StateDir) // Open the state database - db, err := statedbFunc(c.logger, c.config.StateDir) + db, err := c.config.StateDBFactory(c.logger, c.config.StateDir) if err != nil { return fmt.Errorf("failed to open state database: %v", err) } diff --git a/client/client_test.go b/client/client_test.go index 5e62c9070..76d473e98 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -9,7 +9,7 @@ import ( "testing" "time" - memdb "github.com/hashicorp/go-memdb" + "github.com/hashicorp/go-memdb" "github.com/hashicorp/nomad/client/config" consulApi "github.com/hashicorp/nomad/client/consul" "github.com/hashicorp/nomad/client/fingerprint" @@ -27,7 +27,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/hashicorp/go-hclog" - "github.com/hashicorp/nomad/client/state" cstate "github.com/hashicorp/nomad/client/state" ctestutil "github.com/hashicorp/nomad/client/testutil" "github.com/stretchr/testify/require" @@ -603,8 +602,8 @@ func TestClient_SaveRestoreState(t *testing.T) { c1.config.Logger = logger catalog := consul.NewMockCatalog(logger) mockService := consulApi.NewMockConsulServiceClient(t, logger) - statedbFunc := cstate.GetStateDBFactory(c1.config.DevMode) - c2, err := NewClient(c1.config, catalog, mockService, statedbFunc) + + c2, err := NewClient(c1.config, catalog, mockService) if err != nil { t.Fatalf("err: %v", err) } @@ -702,11 +701,12 @@ func TestClient_RestoreError(t *testing.T) { mockService := consulApi.NewMockConsulServiceClient(t, logger) // This stateDB returns errors for all methods called by restore - statedbFunc := func(hclog.Logger, string) (cstate.StateDB, error) { + stateDBFunc := func(hclog.Logger, string) (cstate.StateDB, error) { return &cstate.ErrDB{[]*structs.Allocation{alloc1}}, nil } + c1.config.StateDBFactory = stateDBFunc - c2, err := NewClient(c1.config, catalog, mockService, statedbFunc) + c2, err := NewClient(c1.config, catalog, mockService) require.Nil(err) defer c2.Shutdown() @@ -736,12 +736,13 @@ func TestClient_Init(t *testing.T) { client := &Client{ config: &config.Config{ - AllocDir: allocDir, + AllocDir: allocDir, + StateDBFactory: cstate.GetStateDBFactory(true), }, logger: testlog.HCLogger(t), } - stateDBFunc := state.GetStateDBFactory(true) - if err := client.init(stateDBFunc); err != nil { + + if err := client.init(); err != nil { t.Fatalf("err: %s", err) } diff --git a/client/config/config.go b/client/config/config.go index 13581e63b..eb8408e78 100644 --- a/client/config/config.go +++ b/client/config/config.go @@ -9,6 +9,7 @@ import ( "time" log "github.com/hashicorp/go-hclog" + "github.com/hashicorp/nomad/client/state" "github.com/hashicorp/nomad/helper" "github.com/hashicorp/nomad/nomad/structs" "github.com/hashicorp/nomad/nomad/structs/config" @@ -214,6 +215,9 @@ type Config struct { // PluginSingletonLoader is a plugin loader that will returns singleton // instances of the plugins. PluginSingletonLoader loader.PluginCatalog + + // StateDBFactory is used to override stateDB implementations, + StateDBFactory state.NewStateDBFunc } func (c *Config) Copy() *Config { diff --git a/client/testing.go b/client/testing.go index 0f23e7850..1a4ab5d13 100644 --- a/client/testing.go +++ b/client/testing.go @@ -7,7 +7,6 @@ import ( "github.com/hashicorp/nomad/client/config" consulApi "github.com/hashicorp/nomad/client/consul" "github.com/hashicorp/nomad/client/fingerprint" - "github.com/hashicorp/nomad/client/state" "github.com/hashicorp/nomad/command/agent/consul" "github.com/hashicorp/nomad/helper/testlog" "github.com/hashicorp/nomad/plugins/shared/catalog" @@ -22,13 +21,6 @@ import ( // and removed in the returned cleanup function. If they are overridden in the // callback then the caller still must run the returned cleanup func. func TestClient(t testing.T, cb func(c *config.Config)) (*Client, func() error) { - return TestClientWithCustomStateDB(t, cb, nil) -} - -// TestClientWithCustomStateDB creates an in-memory client for testing purposes -// where the state DB factory can be overridden. It is used in tests that -// simulate state restore failures -func TestClientWithCustomStateDB(t testing.T, cb func(c *config.Config), stateDBFunc state.NewStateDBFunc) (*Client, func() error) { conf, cleanup := config.TestClientConfig(t) // Tighten the fingerprinter timeouts (must be done in client package @@ -54,10 +46,7 @@ func TestClientWithCustomStateDB(t testing.T, cb func(c *config.Config), stateDB } catalog := consul.NewMockCatalog(logger) mockService := consulApi.NewMockConsulServiceClient(t, logger) - if stateDBFunc == nil { - stateDBFunc = state.GetStateDBFactory(conf.DevMode) - } - client, err := NewClient(conf, catalog, mockService, stateDBFunc) + client, err := NewClient(conf, catalog, mockService) if err != nil { cleanup() t.Fatalf("err: %v", err) diff --git a/command/agent/agent.go b/command/agent/agent.go index 6c860c1f0..52e835e43 100644 --- a/command/agent/agent.go +++ b/command/agent/agent.go @@ -731,8 +731,11 @@ func (a *Agent) setupClient() error { return err } } - statedbFactory := state.GetStateDBFactory(conf.DevMode) - client, err := client.NewClient(conf, a.consulCatalog, a.consulService, statedbFactory) + if conf.StateDBFactory == nil { + conf.StateDBFactory = state.GetStateDBFactory(conf.DevMode) + } + + client, err := client.NewClient(conf, a.consulCatalog, a.consulService) if err != nil { return fmt.Errorf("client setup failed: %v", err) }