make drain option configuable

k0sproject · Jun 2, 2024 · f41d0d4 · f41d0d4
1 parent 5de84fb
commit f41d0d4
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 8 deletions.
diff --git a/phase/reset_controllers.go b/phase/reset_controllers.go
@@ -57,7 +57,10 @@ func (p *ResetControllers) Run() error {
 				Metadata: cluster.HostMetadata{
 					Hostname: h.Metadata.Hostname,
 				},
-			}); err != nil {
+			},
+				"120",
+				"5m",
+			); err != nil {
 				log.Warnf("%s: failed to drain node: %s", h, err.Error())
 			}
 		}

diff --git a/phase/reset_workers.go b/phase/reset_workers.go
@@ -56,7 +56,7 @@ func (p *ResetWorkers) Run() error {
 				Metadata: cluster.HostMetadata{
 					Hostname: h.Metadata.Hostname,
 				},
-			}); err != nil {
+			}, "120", "5m"); err != nil {
 				log.Warnf("%s: failed to drain node: %s", h, err.Error())
 			}
 		}
@@ -99,7 +99,7 @@ func (p *ResetWorkers) Run() error {
 			log.Warnf("%s: failed to remove existing configuration %s: %s", h, h.Configurer.K0sConfigPath(), dErr)
 		}
 		log.Debugf("%s: removing config completed", h)
-		
+
 		if len(h.Environment) > 0 {
 			if err := h.Configurer.CleanupServiceEnvironment(h, h.K0sServiceName()); err != nil {
 				log.Warnf("%s: failed to clean up service environment: %s", h, err.Error())

diff --git a/phase/upgrade_workers.go b/phase/upgrade_workers.go
@@ -136,7 +136,21 @@ func (p *UpgradeWorkers) drainWorker(h *cluster.Host) error {
 		return nil
 	}
 	log.Debugf("%s: drain", h)
-	if err := p.leader.DrainNode(h); err != nil {
+	upgradeSettings := p.Config.Spec.UpgradeSettings
+	if upgradeSettings == nil {
+		upgradeSettings = &cluster.UpgradeSettings{}
+	}
+
+	gracePeriod := upgradeSettings.DrainGracePeriod
+	if gracePeriod == "" {
+		gracePeriod = "30"
+	}
+	timeout := upgradeSettings.DrainTimeout
+	if timeout == "" {
+		timeout = "5m"
+	}
+
+	if err := p.leader.DrainNode(h, gracePeriod, timeout); err != nil {
 		return fmt.Errorf("drain node: %w", err)
 	}
 	return nil

diff --git a/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/host.go b/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/host.go
@@ -417,8 +417,8 @@ func (h *Host) K0sDataDir() string {
 }
 
 // DrainNode drains the given node
-func (h *Host) DrainNode(node *Host) error {
-	return h.Exec(h.Configurer.KubectlCmdf(h, h.K0sDataDir(), "drain --grace-period=120 --force --timeout=5m --ignore-daemonsets --delete-emptydir-data %s", node.Metadata.Hostname), exec.Sudo(h))
+func (h *Host) DrainNode(node *Host, gracePeriod string, timeout string) error {
+	return h.Exec(h.Configurer.KubectlCmdf(h, h.K0sDataDir(), "drain --grace-period=%s --force --timeout=%s --ignore-daemonsets --delete-emptydir-data %s", gracePeriod, timeout, node.Metadata.Hostname), exec.Sudo(h))
 }
 
 // CordonNode marks the node unschedulable

diff --git a/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/spec.go b/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/spec.go
@@ -9,17 +9,24 @@ import (
 
 // Spec defines cluster config spec section
 type Spec struct {
-	Hosts Hosts `yaml:"hosts,omitempty"`
-	K0s   *K0s  `yaml:"k0s,omitempty"`
+	Hosts           Hosts            `yaml:"hosts,omitempty"`
+	K0s             *K0s             `yaml:"k0s,omitempty"`
+	UpgradeSettings *UpgradeSettings `yaml:"upgrade,omitempty"`
 
 	k0sLeader *Host
 }
 
+type UpgradeSettings struct {
+	DrainGracePeriod string `yaml:"drainGracePeriod,omitempty"`
+	DrainTimeout     string `yaml:"drainTimeout,omitempty"`
+}
+
 // UnmarshalYAML sets in some sane defaults when unmarshaling the data from yaml
 func (s *Spec) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	type spec Spec
 	ys := (*spec)(s)
 	ys.K0s = &K0s{}
+	ys.UpgradeSettings = &UpgradeSettings{}
 
 	if err := unmarshal(ys); err != nil {
 		return err