From dccde10368ff3437a6d1967f0eee714ea3060dbd Mon Sep 17 00:00:00 2001
From: Lyndon-Li <lyonghui@vmware.com>
Date: Mon, 18 Mar 2024 10:22:35 +0800
Subject: [PATCH] issue-7036: make affinity as list and take 1st one

Signed-off-by: Lyndon-Li <lyonghui@vmware.com>
---
 pkg/cmd/cli/nodeagent/server.go  |   4 +-
 pkg/exposer/csi_snapshot_test.go | 113 +++++++++++++++++++++++++++++++
 pkg/nodeagent/node_agent.go      |   2 +-
 3 files changed, 116 insertions(+), 3 deletions(-)

diff --git a/pkg/cmd/cli/nodeagent/server.go b/pkg/cmd/cli/nodeagent/server.go
index 95d47ca8e..c28a7fbc2 100644
--- a/pkg/cmd/cli/nodeagent/server.go
+++ b/pkg/cmd/cli/nodeagent/server.go
@@ -286,8 +286,8 @@ func (s *nodeAgentServer) run() {
 	}
 
 	var loadAffinity *nodeagent.LoadAffinity
-	if s.dataPathConfigs != nil {
-		loadAffinity = s.dataPathConfigs.LoadAffinity
+	if s.dataPathConfigs != nil && len(s.dataPathConfigs.LoadAffinity) > 0 {
+		loadAffinity = s.dataPathConfigs.LoadAffinity[0]
 	}
 	dataUploadReconciler := controller.NewDataUploadReconciler(s.mgr.GetClient(), s.kubeClient, s.csiSnapshotClient.SnapshotV1(), s.dataPathMgr, loadAffinity, repoEnsurer, clock.RealClock{}, credentialGetter, s.nodeName, s.fileSystem, s.config.dataMoverPrepareTimeout, s.logger, s.metrics)
 	s.attemptDataUploadResume(dataUploadReconciler)
diff --git a/pkg/exposer/csi_snapshot_test.go b/pkg/exposer/csi_snapshot_test.go
index a5fa48976..5d54b8661 100644
--- a/pkg/exposer/csi_snapshot_test.go
+++ b/pkg/exposer/csi_snapshot_test.go
@@ -18,6 +18,7 @@ package exposer
 
 import (
 	"context"
+	"reflect"
 	"testing"
 	"time"
 
@@ -35,6 +36,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 
 	velerov1 "github.com/vmware-tanzu/velero/pkg/apis/velero/v1"
+	"github.com/vmware-tanzu/velero/pkg/nodeagent"
 	velerotest "github.com/vmware-tanzu/velero/pkg/test"
 	"github.com/vmware-tanzu/velero/pkg/util/boolptr"
 
@@ -718,3 +720,114 @@ func TestPeekExpose(t *testing.T) {
 		})
 	}
 }
+
+func TestToSystemAffinity(t *testing.T) {
+	tests := []struct {
+		name         string
+		loadAffinity *nodeagent.LoadAffinity
+		expected     *corev1.Affinity
+	}{
+		{
+			name: "loadAffinity is nil",
+		},
+		{
+			name:         "loadAffinity is empty",
+			loadAffinity: &nodeagent.LoadAffinity{},
+		},
+		{
+			name: "with match label",
+			loadAffinity: &nodeagent.LoadAffinity{
+				NodeSelector: metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						"key-1": "value-1",
+						"key-2": "value-2",
+					},
+				},
+			},
+			expected: &corev1.Affinity{
+				NodeAffinity: &corev1.NodeAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+						NodeSelectorTerms: []corev1.NodeSelectorTerm{
+							{
+								MatchExpressions: []corev1.NodeSelectorRequirement{
+									{
+										Key:      "key-1",
+										Values:   []string{"value-1"},
+										Operator: corev1.NodeSelectorOpIn,
+									},
+									{
+										Key:      "key-2",
+										Values:   []string{"value-2"},
+										Operator: corev1.NodeSelectorOpIn,
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "with match expression",
+			loadAffinity: &nodeagent.LoadAffinity{
+				NodeSelector: metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						"key-1": "value-1",
+						"key-2": "value-2",
+					},
+					MatchExpressions: []metav1.LabelSelectorRequirement{
+						{
+							Key:      "key-3",
+							Values:   []string{"value-3-1", "value-3-2"},
+							Operator: metav1.LabelSelectorOpNotIn,
+						},
+						{
+							Key:      "key-4",
+							Values:   []string{"value-4-1", "value-4-2", "value-4-3"},
+							Operator: metav1.LabelSelectorOpDoesNotExist,
+						},
+					},
+				},
+			},
+			expected: &corev1.Affinity{
+				NodeAffinity: &corev1.NodeAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+						NodeSelectorTerms: []corev1.NodeSelectorTerm{
+							{
+								MatchExpressions: []corev1.NodeSelectorRequirement{
+									{
+										Key:      "key-1",
+										Values:   []string{"value-1"},
+										Operator: corev1.NodeSelectorOpIn,
+									},
+									{
+										Key:      "key-2",
+										Values:   []string{"value-2"},
+										Operator: corev1.NodeSelectorOpIn,
+									},
+									{
+										Key:      "key-3",
+										Values:   []string{"value-3-1", "value-3-2"},
+										Operator: corev1.NodeSelectorOpNotIn,
+									},
+									{
+										Key:      "key-4",
+										Values:   []string{"value-4-1", "value-4-2", "value-4-3"},
+										Operator: corev1.NodeSelectorOpDoesNotExist,
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			affinity := toSystemAffinity(test.loadAffinity)
+			assert.Equal(t, true, reflect.DeepEqual(affinity, test.expected))
+		})
+	}
+}
diff --git a/pkg/nodeagent/node_agent.go b/pkg/nodeagent/node_agent.go
index 7fad915ae..e3597e27a 100644
--- a/pkg/nodeagent/node_agent.go
+++ b/pkg/nodeagent/node_agent.go
@@ -68,7 +68,7 @@ type Configs struct {
 	LoadConcurrency *LoadConcurrency `json:"loadConcurrency,omitempty"`
 
 	// LoadAffinity is the config for data path load affinity.
-	LoadAffinity *LoadAffinity `json:"loadAffinity,omitempty"`
+	LoadAffinity []*LoadAffinity `json:"loadAffinity,omitempty"`
 }
 
 // IsRunning checks if the node agent daemonset is running properly. If not, return the error found