pkg/config: Add config option for drain timeout

Add the option to configure the drain timeout. Move kubeconfig to new kubernetes section in the configuration. Signed-off-by: Heathcliff <[email protected]>
heathcliff26 · Nov 19, 2024 · 1db89df · 1db89df
1 parent 9f81e96
commit 1db89df
Show file tree

Hide file tree

Showing 9 changed files with 78 additions and 24 deletions.
diff --git a/examples/example-config.yaml b/examples/example-config.yaml
@@ -2,9 +2,13 @@
 # The level of loggin output
 logLevel: info
 
-# The path to the kubeconfig file. Used for draining nodes and the kubernetes storage option if used.
-# Falls back to in-cluster serviceaccount if left empty (recommended)
-kubeconfig: ""
+kubernetes:
+  # The path to the kubeconfig file. Used for draining nodes and the kubernetes storage option if used.
+  # Falls back to in-cluster serviceaccount if left empty (recommended)
+  kubeconfig: ""
+  # The timeout for drain node in seconds.
+  # Defaults to 5m.
+  drainTimeoutSeconds: 300
 
 server:
   # The listen address of the server in the form of <ip>:<port>

diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
@@ -60,7 +60,7 @@ func run(cmd *cobra.Command, configPath string, env bool) {
 		exitError(cmd, fmt.Errorf("failed to load configuration: %w", err))
 	}
 
-	k8s, err := k8s.NewClient(cfg.Kubeconfig)
+	k8s, err := k8s.NewClient(cfg.KubernetesConfig)
 	if err != nil {
 		exitError(cmd, fmt.Errorf("failed to create kubernetes client: %w", err))
 	}

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -5,6 +5,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/heathcliff26/fleetlock/pkg/k8s"
 	lockmanager "github.com/heathcliff26/fleetlock/pkg/lock-manager"
 	"github.com/heathcliff26/fleetlock/pkg/server"
 	"gopkg.in/yaml.v3"
@@ -27,11 +28,11 @@ func init() {
 }
 
 type Config struct {
-	LogLevel   string                    `yaml:"logLevel,omitempty"`
-	Kubeconfig string                    `yaml:"kubeconfig,omitempty"`
-	Server     *server.ServerConfig      `yaml:"server,omitempty"`
-	Storage    lockmanager.StorageConfig `yaml:"storage,omitempty"`
-	Groups     lockmanager.Groups        `yaml:"groups,omitempty"`
+	LogLevel         string                    `yaml:"logLevel,omitempty"`
+	KubernetesConfig k8s.Config                `yaml:"kubernetes,omitempty"`
+	Server           *server.ServerConfig      `yaml:"server,omitempty"`
+	Storage          lockmanager.StorageConfig `yaml:"storage,omitempty"`
+	Groups           lockmanager.Groups        `yaml:"groups,omitempty"`
 }
 
 // Parse a given string and set the resulting log level
@@ -53,9 +54,10 @@ func setLogLevel(level string) error {
 
 func DefaultConfig() *Config {
 	return &Config{
-		LogLevel: DEFAULT_LOG_LEVEL,
-		Server:   server.NewDefaultServerConfig(),
-		Storage:  lockmanager.NewDefaultStorageConfig(),
+		LogLevel:         DEFAULT_LOG_LEVEL,
+		KubernetesConfig: k8s.NewDefaultConfig(),
+		Server:           server.NewDefaultServerConfig(),
+		Storage:          lockmanager.NewDefaultStorageConfig(),
 	}
 }
 
@@ -100,7 +102,7 @@ func (c *Config) Defaults() {
 		c.Groups = lockmanager.NewDefaultGroups()
 	}
 
-	c.Storage.Kubernetes.Kubeconfig = c.Kubeconfig
+	c.Storage.Kubernetes.Kubeconfig = c.KubernetesConfig.Kubeconfig
 }
 
 func (c *Config) Validate() error {

diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
@@ -5,6 +5,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/heathcliff26/fleetlock/pkg/k8s"
 	lockmanager "github.com/heathcliff26/fleetlock/pkg/lock-manager"
 	"github.com/heathcliff26/fleetlock/pkg/lock-manager/storage/kubernetes"
 	"github.com/heathcliff26/fleetlock/pkg/lock-manager/storage/sql"
@@ -15,7 +16,8 @@ import (
 
 func TestValidConfigs(t *testing.T) {
 	c1 := &Config{
-		LogLevel: "debug",
+		LogLevel:         "debug",
+		KubernetesConfig: k8s.NewDefaultConfig(),
 		Server: &server.ServerConfig{
 			Listen: "127.0.0.1:8443",
 			SSL: server.SSLConfig{
@@ -60,7 +62,8 @@ func TestValidConfigs(t *testing.T) {
 			Kubeconfig: "some-path",
 		},
 	}
-	cfgKubernetes.Kubeconfig = "some-path"
+	cfgKubernetes.KubernetesConfig.Kubeconfig = "some-path"
+	cfgKubernetes.KubernetesConfig.DrainTimeoutSeconds = 60
 
 	tMatrix := []struct {
 		Name, Path string
@@ -87,7 +90,7 @@ func TestValidConfigs(t *testing.T) {
 			Result: defaultConfig,
 		},
 		{
-			Name:   "ValidConfig3",
+			Name:   "ValidConfigKubernetes",
 			Path:   "testdata/valid-config-kubernetes.yaml",
 			Result: cfgKubernetes,
 		},

diff --git a/pkg/config/testdata/valid-config-kubernetes.yaml b/pkg/config/testdata/valid-config-kubernetes.yaml
@@ -3,4 +3,6 @@ storage:
   type: kubernetes
   kubernetes:
     namespace: "test"
-kubeconfig: "some-path"
+kubernetes:
+  kubeconfig: "some-path"
+  drainTimeoutSeconds: 60
diff --git a/pkg/k8s/client.go b/pkg/k8s/client.go
@@ -25,8 +25,8 @@ type Client struct {
 }
 
 // Create a new kubernetes client, defaults to in-cluster if no kubeconfig is provided
-func NewClient(kubeconfig string) (*Client, error) {
-	client, err := utils.CreateNewClientset(kubeconfig)
+func NewClient(config Config) (*Client, error) {
+	client, err := utils.CreateNewClientset(config.Kubeconfig)
 	if err == rest.ErrNotInCluster {
 		return nil, nil
 	} else if err != nil {
@@ -38,10 +38,14 @@ func NewClient(kubeconfig string) (*Client, error) {
 		return nil, err
 	}
 
+	if config.DrainTimeoutSeconds < 1 {
+		return nil, NewErrorDrainTimeoutSecondsInvalid()
+	}
+
 	return &Client{
 		client:              client,
 		namespace:           ns,
-		drainTimeoutSeconds: 300,
+		drainTimeoutSeconds: config.DrainTimeoutSeconds,
 	}, nil
 }
 

diff --git a/pkg/k8s/client_test.go b/pkg/k8s/client_test.go
@@ -57,22 +57,39 @@ func initTestCluster(client *fake.Clientset) {
 
 func TestNewClient(t *testing.T) {
 	t.Run("NotInCluster", func(t *testing.T) {
-		c, err := NewClient("")
+		c, err := NewClient(NewDefaultConfig())
 		assert.Nil(t, c, "Should not return a client")
 		assert.Nil(t, err, "Should not return an error if not in cluster and no kubeconfig provided")
 	})
 	t.Run("KubeconfigNotFound", func(t *testing.T) {
-		c, err := NewClient("not-a-file")
+		cfg := NewDefaultConfig()
+		cfg.Kubeconfig = "not-a-file"
+
+		c, err := NewClient(cfg)
 		assert.Nil(t, c, "Should not return a client")
 		assert.Error(t, err, "Should return an error if it can't find a kubeconfig")
 	})
-	t.Run("Kubeconfig", func(t *testing.T) {
-		c, err := NewClient("testdata/kubeconfig")
+	t.Run("InvalidDrainTimeout", func(t *testing.T) {
+		cfg := NewDefaultConfig()
+		cfg.Kubeconfig = "testdata/kubeconfig"
+		cfg.DrainTimeoutSeconds = 0
+
+		c, err := NewClient(cfg)
+		assert.Nil(t, c, "Should not return a client")
+		assert.Error(t, err, "Should return an error")
+	})
+	t.Run("Success", func(t *testing.T) {
+		cfg := NewDefaultConfig()
+		cfg.Kubeconfig = "testdata/kubeconfig"
+		cfg.DrainTimeoutSeconds = 5
+
+		c, err := NewClient(cfg)
 		assert.Nil(t, err, "Should not return an error")
 		if !assert.NotNil(t, c, "Should return a client") {
 			t.FailNow()
 		}
 		assert.Equal(t, "fleetlock", c.namespace)
+		assert.Equal(t, cfg.DrainTimeoutSeconds, c.drainTimeoutSeconds)
 	})
 }
 

diff --git a/pkg/k8s/config.go b/pkg/k8s/config.go
@@ -0,0 +1,12 @@
+package k8s
+
+type Config struct {
+	Kubeconfig          string `yaml:"kubeconfig,omitempty"`
+	DrainTimeoutSeconds int32  `yaml:"drainTimeoutSeconds,omitempty"`
+}
+
+func NewDefaultConfig() Config {
+	return Config{
+		DrainTimeoutSeconds: 300,
+	}
+}
diff --git a/pkg/k8s/errors.go b/pkg/k8s/errors.go
@@ -39,3 +39,13 @@ func NewErrorLeaseNil() error {
 func (e ErrorLeaseNil) Error() string {
 	return "Tried changing lease, but lease status on cluster is unknown"
 }
+
+type ErrorDrainTimeoutSecondsInvalid struct{}
+
+func NewErrorDrainTimeoutSecondsInvalid() error {
+	return ErrorDrainTimeoutSecondsInvalid{}
+}
+
+func (e ErrorDrainTimeoutSecondsInvalid) Error() string {
+	return "drainTimeoutSeconds value needs to be greater than 0"
+}