From b213255bc806f0aa0bbde6131b3c431f22cbad15 Mon Sep 17 00:00:00 2001
From: Abdullah Gharaibeh <ahg@google.com>
Date: Wed, 23 Oct 2019 19:19:36 -0400
Subject: [PATCH] interpod affinity prioritize

---
 .../algorithm/priorities/interpod_affinity.go | 47 +++++++++++++------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/pkg/scheduler/algorithm/priorities/interpod_affinity.go b/pkg/scheduler/algorithm/priorities/interpod_affinity.go
index 3c7eadb4460..ae0d14ac745 100644
--- a/pkg/scheduler/algorithm/priorities/interpod_affinity.go
+++ b/pkg/scheduler/algorithm/priorities/interpod_affinity.go
@@ -18,7 +18,7 @@ package priorities
 
 import (
 	"context"
-	"sync/atomic"
+	"sync"
 
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -49,17 +49,20 @@ func NewInterPodAffinityPriority(
 	return interPodAffinity.CalculateInterPodAffinityPriority
 }
 
+type topologyPairToScore map[string]map[string]int64
+
 type podAffinityPriorityMap struct {
 	// nodes contain all nodes that should be considered.
 	nodes []*v1.Node
-	// counts store the so-far computed score for each node.
-	counts []int64
+	// tracks a topology pair score so far.
+	topologyScore topologyPairToScore
+	sync.Mutex
 }
 
 func newPodAffinityPriorityMap(nodes []*v1.Node) *podAffinityPriorityMap {
 	return &podAffinityPriorityMap{
-		nodes:  nodes,
-		counts: make([]int64, len(nodes)),
+		nodes:         nodes,
+		topologyScore: make(topologyPairToScore),
 	}
 }
 
@@ -69,13 +72,19 @@ func (p *podAffinityPriorityMap) processTerm(term *v1.PodAffinityTerm, podDefini
 	if err != nil {
 		return err
 	}
+	if len(fixedNode.Labels) == 0 {
+		return nil
+	}
+
 	match := priorityutil.PodMatchesTermsNamespaceAndSelector(podToCheck, namespaces, selector)
-	if match {
-		for i, node := range p.nodes {
-			if priorityutil.NodesHaveSameTopologyKey(node, fixedNode, term.TopologyKey) {
-				atomic.AddInt64(&p.counts[i], weight)
-			}
+	tpValue, tpValueExist := fixedNode.Labels[term.TopologyKey]
+	if match && tpValueExist {
+		p.Lock()
+		if p.topologyScore[term.TopologyKey] == nil {
+			p.topologyScore[term.TopologyKey] = make(map[string]int64)
 		}
+		p.topologyScore[term.TopologyKey][tpValue] += weight
+		p.Unlock()
 	}
 	return nil
 }
@@ -205,12 +214,20 @@ func (ipa *InterPodAffinity) CalculateInterPodAffinityPriority(pod *v1.Pod, node
 		return nil, err
 	}
 
+	counts := make([]int64, len(nodes))
 	for i := range nodes {
-		if pm.counts[i] > maxCount {
-			maxCount = pm.counts[i]
+		if nodes[i].Labels != nil {
+			for tpKey, tpValues := range pm.topologyScore {
+				if v, exist := nodes[i].Labels[tpKey]; exist {
+					counts[i] += tpValues[v]
+				}
+			}
 		}
-		if pm.counts[i] < minCount {
-			minCount = pm.counts[i]
+		if counts[i] > maxCount {
+			maxCount = counts[i]
+		}
+		if counts[i] < minCount {
+			minCount = counts[i]
 		}
 	}
 
@@ -220,7 +237,7 @@ func (ipa *InterPodAffinity) CalculateInterPodAffinityPriority(pod *v1.Pod, node
 	for i, node := range nodes {
 		fScore := float64(0)
 		if maxMinDiff > 0 {
-			fScore = float64(framework.MaxNodeScore) * (float64(pm.counts[i]-minCount) / float64(maxCount-minCount))
+			fScore = float64(framework.MaxNodeScore) * (float64(counts[i]-minCount) / float64(maxCount-minCount))
 		}
 		result = append(result, framework.NodeScore{Name: node.Name, Score: int64(fScore)})
 		if klog.V(10) {