From 5c54e6c449e50a2d18a6f32f987cbfa95968cbdf Mon Sep 17 00:00:00 2001
From: "sunyakun.king" <sunyakun.king@bytedance.com>
Date: Thu, 13 Jul 2023 00:39:36 +0800
Subject: [PATCH] fix(provider): update join k3d function and rollback logic

---
 pkg/cluster/base.go      | 14 +++++++++++---
 pkg/providers/k3d/k3d.go | 27 +++++++++++++++++----------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/pkg/cluster/base.go b/pkg/cluster/base.go
index 6cb5dfc6..523155b2 100644
--- a/pkg/cluster/base.go
+++ b/pkg/cluster/base.go
@@ -1142,10 +1142,18 @@ func (p *ProviderBase) RollbackCluster(rollbackInstance func(ids []string) error
 		if err := rollbackInstance(ids); err != nil {
 			return err
 		}
-		// remove context.
-		if err := common.FileManager.ClearCfgByContext(p.ContextName); err != nil {
-			logrus.Errorf("failed to remove cluster context %s from kube config", p.ContextName)
+
+		state, err := common.DefaultDB.GetCluster(p.Name, p.Provider)
+		if err != nil {
+			return err
 		}
+		if state == nil || state.Status != common.StatusRunning {
+			// remove context.
+			if err := common.FileManager.ClearCfgByContext(p.ContextName); err != nil {
+				logrus.Errorf("failed to remove cluster context %s from kube config", p.ContextName)
+			}
+		}
+
 		p.Logger.Infof("[%s] successfully executed rollback logic", p.Provider)
 	}
 
diff --git a/pkg/providers/k3d/k3d.go b/pkg/providers/k3d/k3d.go
index 79b3e321..03cc507b 100644
--- a/pkg/providers/k3d/k3d.go
+++ b/pkg/providers/k3d/k3d.go
@@ -368,6 +368,13 @@ func (p *K3d) createK3d(ssh *types.SSH) (*types.Cluster, error) {
 		return nil, err
 	}
 
+	p.M.Range(func(key, val interface{}) bool {
+		v := val.(types.Node)
+		v.RollBack = true
+		p.M.Store(v.InstanceID, v)
+		return true
+	})
+
 	c := &types.Cluster{
 		Metadata: p.Metadata,
 		Options:  p.Options,
@@ -415,13 +422,13 @@ func (p *K3d) joinK3d(ssh *types.SSH) (*types.Cluster, error) {
 	sort.Ints(serverIndexes)
 	sort.Ints(agentIndexes)
 
+	index := 0
+	if len(serverIndexes) > 0 {
+		index = serverIndexes[len(serverIndexes)-1] + 1
+	}
 	for i := 0; i < masterNum; i++ {
-		index := 0
-		if len(serverIndexes)-1 > 0 {
-			index = serverIndexes[len(serverIndexes)-1]
-		}
 		node := &k3d.Node{
-			Name:  fmt.Sprintf("%s-%s-%s-%d", k3d.DefaultObjectNamePrefix, p.Name, "server", index+1+i),
+			Name:  fmt.Sprintf("%s-%s-%s-%d", k3d.DefaultObjectNamePrefix, p.Name, "server", index+i),
 			Role:  k3d.ServerRole,
 			Image: p.Image,
 			K3sNodeLabels: map[string]string{
@@ -435,13 +442,13 @@ func (p *K3d) joinK3d(ssh *types.SSH) (*types.Cluster, error) {
 		nodes = append(nodes, node)
 	}
 
+	index = 0
+	if len(agentIndexes) > 0 {
+		index = agentIndexes[len(agentIndexes)-1] + 1
+	}
 	for i := 0; i < workerNum; i++ {
-		index := 0
-		if len(agentIndexes)-1 > 0 {
-			index = agentIndexes[len(agentIndexes)-1]
-		}
 		node := &k3d.Node{
-			Name:  fmt.Sprintf("%s-%s-%s-%d", k3d.DefaultObjectNamePrefix, p.Name, "agent", index+1+i),
+			Name:  fmt.Sprintf("%s-%s-%s-%d", k3d.DefaultObjectNamePrefix, p.Name, "agent", index+i),
 			Role:  k3d.AgentRole,
 			Image: p.Image,
 			K3sNodeLabels: map[string]string{