From 768204ea7ea8bb8b320380ece80eb6887ff001c9 Mon Sep 17 00:00:00 2001
From: Keiji Yoshimi <walf443@gmail.com>
Date: Sat, 1 Jun 2024 18:54:34 +0900
Subject: [PATCH] make drain option configuable

---
 phase/reset_controllers.go                       |  5 ++++-
 phase/reset_workers.go                           |  4 ++--
 phase/upgrade_workers.go                         | 16 +++++++++++++++-
 .../k0sctl.k0sproject.io/v1beta1/cluster/host.go |  4 ++--
 .../k0sctl.k0sproject.io/v1beta1/cluster/spec.go | 11 +++++++++--
 5 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/phase/reset_controllers.go b/phase/reset_controllers.go
index 5b33f718..0c6f6591 100644
--- 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
index 1007fea7..db4cb3e7 100644
--- 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
index ddf8a02d..8330bfe7 100644
--- a/phase/upgrade_workers.go
+++ b/phase/upgrade_workers.go
@@ -131,7 +131,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 = "120"
+	}
+	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
index 321321f4..c5f200d2 100644
--- 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
index 02a6dc49..82f50566 100644
--- 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