From e6efa9062ff3d34874b65a83602736e4cb4a0bdd Mon Sep 17 00:00:00 2001
From: zeaslity <ice@qq.com>
Date: Tue, 4 Jun 2024 17:46:25 +0800
Subject: [PATCH] [agent-deploy] - node cordon

---
 agent-operator/K8sOperator.go      | 96 ++++++++++++++++++++++++++++++
 agent-operator/K8sOperator_test.go | 32 +++++++++-
 2 files changed, 126 insertions(+), 2 deletions(-)

diff --git a/agent-operator/K8sOperator.go b/agent-operator/K8sOperator.go
index dde0ae7..79df48e 100644
--- a/agent-operator/K8sOperator.go
+++ b/agent-operator/K8sOperator.go
@@ -6,6 +6,7 @@ import (
 	v1 "k8s.io/api/apps/v1"
 	autoscalingv1 "k8s.io/api/autoscaling/v1"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/api/policy/v1beta1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/kubernetes/scheme"
@@ -909,3 +910,98 @@ func (op *CmiiK8sOperator) NodeExists(cmiiEnv, nodeName string) (node *CmiiNodeI
 
 	return &convert
 }
+
+func (op *CmiiK8sOperator) NodeCordon(cmiiEnv, nodeName string) (ok bool) {
+	op.changeOperatorEnv(cmiiEnv)
+	client := op.CurrentClient
+
+	nodeInfo, err := client.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+	if err != nil {
+		log.ErrorF("[NodeExists] - [%s] [%s] not exists !", cmiiEnv, nodeName)
+		return false
+	}
+
+	if nodeInfo.Spec.Unschedulable {
+		return true
+	}
+
+	// 设置 Unschedulable 字段为 true
+	nodeInfo.Spec.Unschedulable = true
+
+	// 更新节点对象
+	_, err = client.CoreV1().Nodes().Update(context.TODO(), nodeInfo, metav1.UpdateOptions{})
+	if err != nil {
+		log.ErrorF("NodeCordon] - [%s] [%s] cordon error ! %s", cmiiEnv, nodeName, err.Error())
+		return false
+	}
+
+	// 删除节点上的全部Pod
+
+	return true
+}
+
+func (op *CmiiK8sOperator) NodeUnCordon(cmiiEnv, nodeName string) (ok bool) {
+	op.changeOperatorEnv(cmiiEnv)
+	client := op.CurrentClient
+
+	nodeInfo, err := client.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+	if err != nil {
+		log.ErrorF("[NodeExists] - [%s] [%s] not exists !", cmiiEnv, nodeName)
+		return false
+	}
+
+	if !nodeInfo.Spec.Unschedulable {
+		return true
+	}
+	// 设置 Unschedulable 字段为 true
+	nodeInfo.Spec.Unschedulable = false
+
+	// 更新节点对象
+	_, err = client.CoreV1().Nodes().Update(context.TODO(), nodeInfo, metav1.UpdateOptions{})
+	if err != nil {
+		log.ErrorF("NodeCordon] - [%s] [%s] cordon error ! %s", cmiiEnv, nodeName, err.Error())
+		return false
+	}
+
+	// 删除节点上的全部Pod
+
+	return true
+}
+
+func (op *CmiiK8sOperator) PodEvictFromNode(cmiiEnv, nodeName string) (podList []*CmiiPodInterface) {
+
+	op.changeOperatorEnv(cmiiEnv)
+	client := op.CurrentClient
+
+	// 获取节点上的所有 Pod
+	pods, err := client.CoreV1().Pods("").List(context.TODO(), metav1.ListOptions{
+		FieldSelector: "spec.nodeName=" + nodeName,
+	})
+	if err != nil {
+		return nil
+	}
+
+	gracePeriodSeconds := int64(0)
+	// 驱逐每个 Pod
+	for _, pod := range pods.Items {
+		err = client.CoreV1().Pods(pod.Namespace).Evict(context.TODO(), &v1beta1.Eviction{
+			TypeMeta: metav1.TypeMeta{},
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      pod.Name,
+				Namespace: pod.Namespace,
+			},
+			DeleteOptions: &metav1.DeleteOptions{
+				GracePeriodSeconds: &gracePeriodSeconds,
+			},
+		})
+		if err != nil {
+			log.WarnF("PodEvictFromNode] - 驱逐 Pod [%s] 失败: %s", pod.Name, err.Error())
+			continue
+		}
+		podInterface := CmiiPodInterface{}
+		convert := podInterface.Convert(pod)
+		podList = append(podList, &convert)
+	}
+
+	return podList
+}
diff --git a/agent-operator/K8sOperator_test.go b/agent-operator/K8sOperator_test.go
index b2ee624..1a3d019 100644
--- a/agent-operator/K8sOperator_test.go
+++ b/agent-operator/K8sOperator_test.go
@@ -212,8 +212,8 @@ func TestCmiiK8sOperator_PodFizz2(t *testing.T) {
 
 func TestCmiiK8sOperator_PodByNodeName(t *testing.T) {
 
-	cmiiEnv := "uat"
-	nodeName := "test-03.ecs.io"
+	cmiiEnv := devOperation
+	nodeName := "dev-04.ecs.io"
 
 	exists := DefaultCmiiOperator.PodByNodeName(cmiiEnv, nodeName)
 	exists = FilterAllCmiiPodSoft(exists)
@@ -227,6 +227,34 @@ func TestCmiiK8sOperator_PodByNodeName(t *testing.T) {
 	}
 }
 
+func TestCmiiK8sOperator_NodeCordon(t *testing.T) {
+	cmiiEnv := devOperation
+	nodeName := "dev-06.ecs.io"
+
+	cordon := DefaultCmiiOperator.NodeCordon(cmiiEnv, nodeName)
+	assert.Equal(t, cordon, true, "cordon node failed !")
+}
+
+func TestCmiiK8sOperator_NodeUnCordon(t *testing.T) {
+	cmiiEnv := devOperation
+	nodeName := "dev-06.ecs.io"
+
+	uncordon := DefaultCmiiOperator.NodeUnCordon(cmiiEnv, nodeName)
+	assert.Equal(t, uncordon, true, "uncordon node failed !")
+}
+
+func TestCmiiK8sOperator_PodEvictFromNode(t *testing.T) {
+	cmiiEnv := devOperation
+	nodeName := "dev-04.ecs.io"
+
+	podList := DefaultCmiiOperator.PodEvictFromNode(cmiiEnv, nodeName)
+
+	for _, podInterface := range podList {
+		utils.BeautifulPrint(podInterface)
+	}
+
+}
+
 func TestCmiiK8sOperator_PodExec(t *testing.T) {
 
 	podList := DefaultCmiiOperator.PodByAppName(devFlight, "cmii-uav-gateway")