feat: node patch improvements (#33)

Author: anjmao

actions/actions.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"runtime/debug"
 	"sync"
 	"time"
 
@@ -176,6 +177,13 @@ func (s *service) startProcessing(actionID string) bool {
 func (s *service) handleAction(ctx context.Context, action *castai.ClusterAction) (err error) {
 	data := action.Data()
 	actionType := reflect.TypeOf(data)
+
+	defer func() {
+		if rerr := recover(); rerr != nil {
+			err = fmt.Errorf("panic: handling action %s: %s: %s", actionType, rerr, string(debug.Stack()))
+		}
+	}()
+
 	s.log.Infof("handling action, id=%s, type=%s", action.ID, actionType)
 	handler, ok := s.actionHandlers[actionType]
 	if !ok {

actions/actions_test.go

@@ -140,14 +140,46 @@ func TestActions(t *testing.T) {
 		}()
 		svc.Run(ctx)
 	})
+
+	t.Run("ack with error when action handler panic occurred", func(t *testing.T) {
+		r := require.New(t)
+
+		apiActions := []*castai.ClusterAction{
+			{
+				ID:        "a1",
+				CreatedAt: time.Now(),
+				ActionPatchNode: &castai.ActionPatchNode{
+					NodeName: "n1",
+				},
+			},
+		}
+		client := mock.NewMockAPIClient(apiActions)
+		handler := &mockAgentActionHandler{panicErr: errors.New("ups")}
+		svc := newTestService(handler, client)
+		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+		defer func() {
+			cancel()
+			svc.startedActionsWg.Wait()
+
+			r.Empty(client.Actions)
+			r.Len(client.Acks, 1)
+			r.Equal("a1", client.Acks[0].ActionID)
+			r.Contains(*client.Acks[0].Err, "panic: handling action *castai.ActionPatchNode: ups: goroutine")
+		}()
+		svc.Run(ctx)
+	})
 }
 
 type mockAgentActionHandler struct {
 	err         error
+	panicErr    error
 	handleDelay time.Duration
 }
 
 func (m *mockAgentActionHandler) Handle(ctx context.Context, data interface{}) error {
 	time.Sleep(m.handleDelay)
+	if m.panicErr != nil {
+		panic(m.panicErr)
+	}
 	return m.err
 }

actions/patch_node_handler.go

@@ -2,6 +2,7 @@ package actions
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/sirupsen/logrus"
@@ -30,6 +31,21 @@ func (h *patchNodeHandler) Handle(ctx context.Context, data interface{}) error {
 	if !ok {
 		return fmt.Errorf("unexpected type %T for delete patch handler", data)
 	}
+	for k := range req.Labels {
+		if k == "" {
+			return errors.New("labels contain entry with empty key")
+		}
+	}
+	for k := range req.Annotations {
+		if k == "" {
+			return errors.New("annotations contain entry with empty key")
+		}
+	}
+	for _, t := range req.Taints {
+		if t.Key == "" {
+			return errors.New("taint contain entry with empty key")
+		}
+	}
 
 	log := h.log.WithField("node_name", req.NodeName)
 
@@ -42,34 +58,59 @@ func (h *patchNodeHandler) Handle(ctx context.Context, data interface{}) error {
 		return err
 	}
 
-	log.Infof("patching node, labels=%v, taints=%v", req.Labels, req.Taints)
+	log.Infof("patching node, labels=%v, taints=%v, annotations=%v", req.Labels, req.Taints, req.Annotations)
 
 	return patchNode(ctx, h.clientset, node, func(n *v1.Node) error {
-		if n.Labels == nil {
-			n.Labels = map[string]string{}
-		}
+		node.Labels = patchNodeMapField(node.Labels, req.Labels)
+		node.Annotations = patchNodeMapField(node.Annotations, req.Annotations)
+		node.Spec.Taints = patchTaints(node.Spec.Taints, req.Taints)
+		return nil
+	})
+}
+
+func patchNodeMapField(values map[string]string, patch map[string]string) map[string]string {
+	if values == nil {
+		values = map[string]string{}
+	}
 
-		for key, val := range req.Labels {
-			n.Labels[key] = val
+	for k, v := range patch {
+		if k[0] == '-' {
+			delete(values, k[1:])
+		} else {
+			values[k] = v
 		}
+	}
+	return values
+}
 
-		seen := make(map[string]struct{}, len(n.Spec.Taints))
-		for _, taint := range n.Spec.Taints {
-			seen[taint.Key] = struct{}{}
+func patchTaints(taints []v1.Taint, patch []castai.NodeTaint) []v1.Taint {
+	for _, v := range patch {
+		taint := &v1.Taint{Key: v.Key, Value: v.Value, Effect: v1.TaintEffect(v.Effect)}
+		if v.Key[0] == '-' {
+			taint.Key = taint.Key[1:]
+			taints = deleteTaint(taints, taint)
+		} else if _, found := findTaint(taints, taint); !found {
+			taints = append(taints, *taint)
 		}
+	}
+	return taints
+}
 
-		for _, newTaint := range req.Taints {
-			if _, ok := seen[newTaint.Key]; ok {
-				continue
-			}
-			taint := v1.Taint{
-				Key:    newTaint.Key,
-				Value:  newTaint.Value,
-				Effect: v1.TaintEffect(newTaint.Effect),
-			}
-			n.Spec.Taints = append(n.Spec.Taints, taint)
+func findTaint(taints []v1.Taint, t *v1.Taint) (v1.Taint, bool) {
+	for _, taint := range taints {
+		if taint.MatchTaint(t) {
+			return taint, true
 		}
+	}
+	return v1.Taint{}, false
+}
 
-		return nil
-	})
+func deleteTaint(taints []v1.Taint, t *v1.Taint) []v1.Taint {
+	var res []v1.Taint
+	for _, taint := range taints {
+		if !taint.MatchTaint(t) {
+			res = append(res, taint)
+		}
+	}
+	return res
 }

actions/patch_node_handler_test.go

@@ -24,6 +24,26 @@ func TestPatchNodeHandler(t *testing.T) {
 		node := &v1.Node{
 			ObjectMeta: metav1.ObjectMeta{
 				Name: nodeName,
+				Labels: map[string]string{
+					"l1": "v1",
+				},
+				Annotations: map[string]string{
+					"a1": "v1",
+				},
+			},
+			Spec: v1.NodeSpec{
+				Taints: []v1.Taint{
+					{
+						Key:    "t1",
+						Value:  "v1",
+						Effect: v1.TaintEffectNoSchedule,
+					},
+					{
+						Key:    "t2",
+						Value:  "v2",
+						Effect: v1.TaintEffectNoSchedule,
+					},
+				},
 			},
 		}
 		clientset := fake.NewSimpleClientset(node)
@@ -36,13 +56,23 @@ func TestPatchNodeHandler(t *testing.T) {
 		req := &castai.ActionPatchNode{
 			NodeName: "node1",
 			Labels: map[string]string{
-				"label": "ok",
+				"-l1": "",
+				"l2":  "v2",
+			},
+			Annotations: map[string]string{
+				"-a1": "",
+				"a2":  "",
 			},
 			Taints: []castai.NodeTaint{
 				{
-					Key:    "taint",
-					Value:  "ok2",
-					Effect: "NoSchedule",
+					Key:    "t3",
+					Value:  "t3",
+					Effect: string(v1.TaintEffectNoSchedule),
+				},
+				{
+					Key:    "-t2",
+					Value:  "",
+					Effect: string(v1.TaintEffectNoSchedule),
 				},
 			},
 		}
@@ -52,8 +82,22 @@ func TestPatchNodeHandler(t *testing.T) {
 
 		n, err := clientset.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
 		r.NoError(err)
-		r.Equal("ok", n.Labels["label"])
-		r.Equal([]v1.Taint{{Key: "taint", Value: "ok2", Effect: "NoSchedule", TimeAdded: (*metav1.Time)(nil)}}, n.Spec.Taints)
+
+		expectedLabels := map[string]string{
+			"l2": "v2",
+		}
+		r.Equal(expectedLabels, n.Labels)
+
+		expectedAnnotations := map[string]string{
+			"a2": "",
+		}
+		r.Equal(expectedAnnotations, n.Annotations)
+
+		expectedTaints := []v1.Taint{
+			{Key: "t1", Value: "v1", Effect: "NoSchedule", TimeAdded: (*metav1.Time)(nil)},
+			{Key: "t3", Value: "t3", Effect: "NoSchedule", TimeAdded: (*metav1.Time)(nil)},
+		}
+		r.Equal(expectedTaints, n.Spec.Taints)
 	})
 
 	t.Run("skip patch when node not found", func(t *testing.T) {

castai/types.go

@@ -89,9 +89,10 @@ type ActionApproveCSR struct {
 }
 
 type ActionPatchNode struct {
-	NodeName string            `json:"nodeName"`
-	Labels   map[string]string `json:"labels"`
-	Taints   []NodeTaint       `json:"taints"`
+	NodeName    string            `json:"nodeName"`
+	Labels      map[string]string `json:"labels"`
+	Taints      []NodeTaint       `json:"taints"`
+	Annotations map[string]string `json:"annotations"`
 }
 
 type NodeTaint struct {