Merge 023d47b into f994574

pkg/api/v1/types.go

@@ -20,6 +20,7 @@ import (
 	"k8s.io/client-go/pkg/runtime/schema"
 	"k8s.io/client-go/pkg/types"
 	"kubevirt.io/kubevirt/pkg/api"
+	"kubevirt.io/kubevirt/pkg/logging"
 	"kubevirt.io/kubevirt/pkg/mapper"
 	"kubevirt.io/kubevirt/pkg/precond"
 	"reflect"
@@ -394,3 +395,47 @@ func (ml *MigrationList) UnmarshalJSON(data []byte) error {
 	*ml = tmp2
 	return nil
 }
+
+// Given a pod, create a NodeSelectorTerm with anti-affinity for that pod's node.
+// This is useful for the case when a migration away from a node must occur.
+// In this context, pod is where a migration is starting
+func AntiAffinityFromNode(pod *v1.Pod) *v1.NodeSelectorTerm {
+	selector := v1.NodeSelectorTerm{
+		MatchExpressions: []v1.NodeSelectorRequirement{
+			v1.NodeSelectorRequirement{
+				Key:      "kubernetes.io/hostname",
+				Operator: v1.NodeSelectorOpNotIn,
+				Values:   []string{pod.Spec.NodeName},
+			},
+		},
+	}
+	return &selector
+}
+
+// Given a pod and an affinity rule, add the affinity to any others
+// associated with the pod (if any). In this context, pod is the destination
+// of a migration.
+func ApplyAntiAffinityToPod(pod *v1.Pod, selector *v1.NodeSelectorTerm) (*v1.Pod, error) {
+	logger := logging.DefaultLogger()
+	obj, err := kubeapi.Scheme.Copy(pod)
+	if err != nil {
+		logger.Error().Reason(err).Msg("could not copy pod object")
+		return nil, err
+	}
+	podCopy := obj.(*v1.Pod)
+	if podCopy.Spec.Affinity == nil {
+		newAffinity := v1.Affinity{}
+		podCopy.Spec.Affinity = &newAffinity
+	}
+	if podCopy.Spec.Affinity.NodeAffinity == nil {
+		newNodeAffinity := v1.NodeAffinity{}
+		podCopy.Spec.Affinity.NodeAffinity = &newNodeAffinity
+	}
+	if podCopy.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
+		newNodeSelector := v1.NodeSelector{}
+		podCopy.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &newNodeSelector
+	}
+	nodeSelector := podCopy.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
+	nodeSelector.NodeSelectorTerms = append(nodeSelector.NodeSelectorTerms, *selector)
+	return podCopy, nil
+}

pkg/api/v1/types_test.go

@@ -0,0 +1,86 @@
+package v1
+
+import (
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+
+	"k8s.io/client-go/pkg/api/v1"
+	"testing"
+)
+
+var _ = Describe("PodSelectors", func() {
+	Context("Pod affinity rules", func() {
+		var (
+			pod             = &v1.Pod{}
+			selector        = &v1.NodeSelectorTerm{}
+			podWithSelector = &v1.Pod{}
+		)
+
+		BeforeEach(func() {
+			pod = &v1.Pod{
+				Spec: v1.PodSpec{
+					NodeName: "testnode",
+				},
+			}
+			selector = &v1.NodeSelectorTerm{
+				MatchExpressions: []v1.NodeSelectorRequirement{
+					v1.NodeSelectorRequirement{
+						Key:      "kubernetes.io/hostname",
+						Operator: v1.NodeSelectorOpNotIn,
+						Values:   []string{pod.Spec.NodeName},
+					},
+				},
+			}
+			podWithSelector = &v1.Pod{
+				Spec: v1.PodSpec{
+					NodeName: "testnode",
+					Affinity: &v1.Affinity{
+						NodeAffinity: &v1.NodeAffinity{
+							RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
+								NodeSelectorTerms: []v1.NodeSelectorTerm{
+									v1.NodeSelectorTerm{
+										MatchExpressions: []v1.NodeSelectorRequirement{
+											v1.NodeSelectorRequirement{
+												Key:      "kubernetes.io/hostname",
+												Operator: v1.NodeSelectorOpIn,
+												Values:   []string{"test2"},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			}
+		})
+
+		AfterEach(func() {
+		})
+
+		It("should work", func() {
+			newSelector := AntiAffinityFromNode(pod)
+			Expect(len(newSelector.MatchExpressions)).To(Equal(1))
+		})
+		It("Should create missing structs", func() {
+			newPod, err := ApplyAntiAffinityToPod(pod, selector)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(newPod.Spec.Affinity).ToNot(BeNil())
+			Expect(newPod.Spec.Affinity.NodeAffinity).ToNot(BeNil())
+			Expect(newPod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution).ToNot(BeNil())
+			terms := newPod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
+			Expect(len(terms)).To(Equal(1))
+		})
+		It("Should append to existing node selectors", func() {
+			newPod, err := ApplyAntiAffinityToPod(podWithSelector, selector)
+                        Expect(err).ToNot(HaveOccurred())
+			terms := newPod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
+			Expect(len(terms)).To(Equal(2))
+		})
+	})
+})
+
+func TestSelectors(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "PodSelectors")
+}