...
1
16
17 package utils
18
19 import (
20 "context"
21 "fmt"
22 "strings"
23 "time"
24
25 "k8s.io/api/core/v1"
26 apierrors "k8s.io/apimachinery/pkg/api/errors"
27 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
28 "k8s.io/apimachinery/pkg/types"
29 clientset "k8s.io/client-go/kubernetes"
30 "k8s.io/klog/v2"
31 )
32
33 const (
34 retries = 5
35 )
36
37 func AddLabelsToNode(c clientset.Interface, nodeName string, labels map[string]string) error {
38 tokens := make([]string, 0, len(labels))
39 for k, v := range labels {
40 tokens = append(tokens, "\""+k+"\":\""+v+"\"")
41 }
42 labelString := "{" + strings.Join(tokens, ",") + "}"
43 patch := fmt.Sprintf(`{"metadata":{"labels":%v}}`, labelString)
44 var err error
45 for attempt := 0; attempt < retries; attempt++ {
46 _, err = c.CoreV1().Nodes().Patch(context.TODO(), nodeName, types.MergePatchType, []byte(patch), metav1.PatchOptions{})
47 if err != nil {
48 if !apierrors.IsConflict(err) {
49 return err
50 }
51 } else {
52 break
53 }
54 time.Sleep(100 * time.Millisecond)
55 }
56 return err
57 }
58
59
60
61 func RemoveLabelOffNode(c clientset.Interface, nodeName string, labelKeys []string) error {
62 var node *v1.Node
63 var err error
64 for attempt := 0; attempt < retries; attempt++ {
65 node, err = c.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
66 if err != nil {
67 return err
68 }
69 if node.Labels == nil {
70 return nil
71 }
72 for _, labelKey := range labelKeys {
73 if node.Labels == nil || len(node.Labels[labelKey]) == 0 {
74 break
75 }
76 delete(node.Labels, labelKey)
77 }
78 _, err = c.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
79 if err != nil {
80 if !apierrors.IsConflict(err) {
81 return err
82 } else {
83 klog.V(2).Infof("Conflict when trying to remove a labels %v from %v", labelKeys, nodeName)
84 }
85 } else {
86 break
87 }
88 time.Sleep(100 * time.Millisecond)
89 }
90 return err
91 }
92
93
94
95 func VerifyLabelsRemoved(c clientset.Interface, nodeName string, labelKeys []string) error {
96 node, err := c.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
97 if err != nil {
98 return err
99 }
100 for _, labelKey := range labelKeys {
101 if node.Labels != nil && len(node.Labels[labelKey]) != 0 {
102 return fmt.Errorf("Failed removing label " + labelKey + " of the node " + nodeName)
103 }
104 }
105 return nil
106 }
107
View as plain text