refactor: move the methods shared between scheduler and binder to pkg…

…/plugins

pkg/plugins/interpodaffinity/util.go

@@ -0,0 +1,170 @@
+/*
+Copyright 2024 The Godel Scheduler Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package interpodaffinity
+
+import (
+	"context"
+	"sync/atomic"
+
+	framework "github.com/kubewharf/godel-scheduler/pkg/framework/api"
+	schedutil "github.com/kubewharf/godel-scheduler/pkg/scheduler/util"
+	"github.com/kubewharf/godel-scheduler/pkg/util/parallelize"
+	podutil "github.com/kubewharf/godel-scheduler/pkg/util/pod"
+	v1 "k8s.io/api/core/v1"
+)
+
+// TODO(Huang-Wei): It might be possible to use "make(map[TopologyPair]*int64)" so that
+// we can do atomic additions instead of using a global mutext, however we need to consider
+// how to init each TopologyToMatchedTermCount.
+type TopologyPair struct {
+	Key   string
+	Value string
+}
+
+type TopologyToMatchedTermCount map[TopologyPair]int64
+
+func (m TopologyToMatchedTermCount) append(toAppend TopologyToMatchedTermCount) {
+	for pair := range toAppend {
+		m[pair] += toAppend[pair]
+	}
+}
+
+func (m TopologyToMatchedTermCount) Clone() TopologyToMatchedTermCount {
+	copy := make(TopologyToMatchedTermCount, len(m))
+	copy.append(m)
+	return copy
+}
+
+// UpdateWithAffinityTerms updates the topologyToMatchedTermCount map with the specified value
+// for each affinity term if "targetPod" matches ALL terms.
+func (m TopologyToMatchedTermCount) UpdateWithAffinityTerms(targetPod *v1.Pod, nodeLbaels map[string]string, affinityTerms []framework.AffinityTerm, value int64) {
+	if PodMatchesAllAffinityTerms(targetPod, affinityTerms) {
+		for _, t := range affinityTerms {
+			if topologyValue, ok := nodeLbaels[t.TopologyKey]; ok {
+				pair := TopologyPair{Key: t.TopologyKey, Value: topologyValue}
+				m[pair] += value
+				// value could be a negative value, hence we delete the entry if
+				// the entry is down to zero.
+				if m[pair] == 0 {
+					delete(m, pair)
+				}
+			}
+		}
+	}
+}
+
+// UpdateWithAntiAffinityTerms updates the topologyToMatchedTermCount map with the specified value
+// for each anti-affinity term matched the target pod.
+func (m TopologyToMatchedTermCount) UpdateWithAntiAffinityTerms(targetPod *v1.Pod, nodeLabels map[string]string, antiAffinityTerms []framework.AffinityTerm, value int64) {
+	// Check anti-affinity terms.
+	for _, a := range antiAffinityTerms {
+		if schedutil.PodMatchesTermsNamespaceAndSelector(targetPod, a.Namespaces, a.Selector) {
+			if topologyValue, ok := nodeLabels[a.TopologyKey]; ok {
+				pair := TopologyPair{Key: a.TopologyKey, Value: topologyValue}
+				m[pair] += value
+				// value could be a negative value, hence we delete the entry if
+				// the entry is down to zero.
+				if m[pair] == 0 {
+					delete(m, pair)
+				}
+			}
+		}
+	}
+}
+
+// PodMatchesAllAffinityTerms returns true IFF the given pod matches all the given terms.
+func PodMatchesAllAffinityTerms(pod *v1.Pod, terms []framework.AffinityTerm) bool {
+	if len(terms) == 0 {
+		return false
+	}
+	for _, term := range terms {
+		if !schedutil.PodMatchesTermsNamespaceAndSelector(pod, term.Namespaces, term.Selector) {
+			return false
+		}
+	}
+	return true
+}
+
+// GetTPMapMatchingExistingAntiAffinity calculates the following for each existing pod on each node:
+// (1) Whether it has PodAntiAffinity
+// (2) Whether any AffinityTerm matches the incoming pod
+func GetTPMapMatchingExistingAntiAffinity(pod *v1.Pod, nodes []framework.NodeInfo, podLauncher podutil.PodLauncher) TopologyToMatchedTermCount {
+	topoMaps := make([]TopologyToMatchedTermCount, len(nodes))
+	index := int32(-1)
+	processNode := func(i int) {
+		nodeInfo := nodes[i]
+		topoMap := make(TopologyToMatchedTermCount)
+		for _, existingPod := range nodeInfo.GetPodsWithRequiredAntiAffinity() {
+			topoMap.UpdateWithAntiAffinityTerms(pod, nodeInfo.GetNodeLabels(podLauncher), existingPod.RequiredAntiAffinityTerms, 1)
+		}
+		if len(topoMap) != 0 {
+			topoMaps[atomic.AddInt32(&index, 1)] = topoMap
+		}
+	}
+	parallelize.Until(context.Background(), len(nodes), processNode)
+
+	result := make(TopologyToMatchedTermCount)
+	for i := 0; i <= int(index); i++ {
+		result.append(topoMaps[i])
+	}
+
+	return result
+}
+
+// GetTPMapMatchingIncomingAffinityAntiAffinity finds existing Pods that match affinity terms of the given "pod".
+// It returns a topologyToMatchedTermCount that are checked later by the affinity
+// predicate. With this topologyToMatchedTermCount available, the affinity predicate does not
+// need to check all the pods in the cluster.
+func GetTPMapMatchingIncomingAffinityAntiAffinity(podInfo *framework.PodInfo, allNodes []framework.NodeInfo, podLauncher podutil.PodLauncher) (TopologyToMatchedTermCount, TopologyToMatchedTermCount) {
+	matchedNodes := schedutil.FilterNodeInfosByPodLauncher(allNodes, podLauncher)
+
+	affinityCounts := make(TopologyToMatchedTermCount)
+	antiAffinityCounts := make(TopologyToMatchedTermCount)
+	if len(podInfo.RequiredAffinityTerms) == 0 && len(podInfo.RequiredAntiAffinityTerms) == 0 {
+		return affinityCounts, antiAffinityCounts
+	}
+
+	affinityCountsList := make([]TopologyToMatchedTermCount, len(matchedNodes))
+	antiAffinityCountsList := make([]TopologyToMatchedTermCount, len(matchedNodes))
+	index := int32(-1)
+	processNode := func(i int) {
+		nodeInfo := matchedNodes[i]
+		affinity := make(TopologyToMatchedTermCount)
+		antiAffinity := make(TopologyToMatchedTermCount)
+		for _, existingPod := range nodeInfo.GetPods() {
+			// Check affinity terms.
+			affinity.UpdateWithAffinityTerms(existingPod.Pod, nodeInfo.GetNodeLabels(podLauncher), podInfo.RequiredAffinityTerms, 1)
+
+			// Check anti-affinity terms.
+			antiAffinity.UpdateWithAntiAffinityTerms(existingPod.Pod, nodeInfo.GetNodeLabels(podLauncher), podInfo.RequiredAntiAffinityTerms, 1)
+		}
+
+		if len(affinity) > 0 || len(antiAffinity) > 0 {
+			k := atomic.AddInt32(&index, 1)
+			affinityCountsList[k] = affinity
+			antiAffinityCountsList[k] = antiAffinity
+		}
+	}
+	parallelize.Until(context.Background(), len(matchedNodes), processNode)
+
+	for i := 0; i <= int(index); i++ {
+		affinityCounts.append(affinityCountsList[i])
+		antiAffinityCounts.append(antiAffinityCountsList[i])
+	}
+
+	return affinityCounts, antiAffinityCounts
+}