From 55d3c82782122dec7e8671fee31be65b57976c46 Mon Sep 17 00:00:00 2001
From: ravisantoshgudimetla <ravisantoshgudimetla@gmail.com>
Date: Mon, 6 Mar 2017 09:24:51 -0500
Subject: [PATCH] Selector spreading improving code readability

---
 .../priorities/selector_spreading.go          | 33 ++++++++++---------
 .../priorities/selector_spreading_test.go     | 23 +++++++++++++
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/plugin/pkg/scheduler/algorithm/priorities/selector_spreading.go b/plugin/pkg/scheduler/algorithm/priorities/selector_spreading.go
index 02ee8aa2573..e7e16a5fe4c 100644
--- a/plugin/pkg/scheduler/algorithm/priorities/selector_spreading.go
+++ b/plugin/pkg/scheduler/algorithm/priorities/selector_spreading.go
@@ -61,7 +61,7 @@ func NewSelectorSpreadPriority(
 
 // Returns selectors of services, RCs and RSs matching the given pod.
 func getSelectors(pod *v1.Pod, sl algorithm.ServiceLister, cl algorithm.ControllerLister, rsl algorithm.ReplicaSetLister, ssl algorithm.StatefulSetLister) []labels.Selector {
-	selectors := make([]labels.Selector, 0, 3)
+	var selectors []labels.Selector
 	if services, err := sl.GetPodServices(pod); err == nil {
 		for _, service := range services {
 			selectors = append(selectors, labels.SelectorFromSet(service.Spec.Selector))
@@ -206,6 +206,21 @@ func NewServiceAntiAffinityPriority(podLister algorithm.PodLister, serviceLister
 	return antiAffinity.CalculateAntiAffinityPriority
 }
 
+// Classifies nodes into ones with labels and without labels.
+func (s *ServiceAntiAffinity) getNodeClassificationByLabels(nodes []*v1.Node) (map[string]string, []string) {
+	labeledNodes := map[string]string{}
+	nonLabeledNodes := []string{}
+	for _, node := range nodes {
+		if labels.Set(node.Labels).Has(s.label) {
+			label := labels.Set(node.Labels).Get(s.label)
+			labeledNodes[node.Name] = label
+		} else {
+			nonLabeledNodes = append(nonLabeledNodes, node.Name)
+		}
+	}
+	return labeledNodes, nonLabeledNodes
+}
+
 // CalculateAntiAffinityPriority spreads pods by minimizing the number of pods belonging to the same service
 // on machines with the same value for a particular label.
 // The label to be considered is provided to the struct (ServiceAntiAffinity).
@@ -228,17 +243,7 @@ func (s *ServiceAntiAffinity) CalculateAntiAffinityPriority(pod *v1.Pod, nodeNam
 	}
 
 	// separate out the nodes that have the label from the ones that don't
-	otherNodes := []string{}
-	labeledNodes := map[string]string{}
-	for _, node := range nodes {
-		if labels.Set(node.Labels).Has(s.label) {
-			label := labels.Set(node.Labels).Get(s.label)
-			labeledNodes[node.Name] = label
-		} else {
-			otherNodes = append(otherNodes, node.Name)
-		}
-	}
-
+	labeledNodes, nonLabeledNodes := s.getNodeClassificationByLabels(nodes)
 	podCounts := map[string]int{}
 	for _, pod := range nsServicePods {
 		label, exists := labeledNodes[pod.Spec.NodeName]
@@ -247,7 +252,6 @@ func (s *ServiceAntiAffinity) CalculateAntiAffinityPriority(pod *v1.Pod, nodeNam
 		}
 		podCounts[label]++
 	}
-
 	numServicePods := len(nsServicePods)
 	result := []schedulerapi.HostPriority{}
 	//score int - scale of 0-maxPriority
@@ -261,9 +265,8 @@ func (s *ServiceAntiAffinity) CalculateAntiAffinityPriority(pod *v1.Pod, nodeNam
 		result = append(result, schedulerapi.HostPriority{Host: node, Score: int(fScore)})
 	}
 	// add the open nodes with a score of 0
-	for _, node := range otherNodes {
+	for _, node := range nonLabeledNodes {
 		result = append(result, schedulerapi.HostPriority{Host: node, Score: 0})
 	}
-
 	return result, nil
 }
diff --git a/plugin/pkg/scheduler/algorithm/priorities/selector_spreading_test.go b/plugin/pkg/scheduler/algorithm/priorities/selector_spreading_test.go
index cfccd767444..4a6d924c3a6 100644
--- a/plugin/pkg/scheduler/algorithm/priorities/selector_spreading_test.go
+++ b/plugin/pkg/scheduler/algorithm/priorities/selector_spreading_test.go
@@ -739,6 +739,29 @@ func TestZoneSpreadPriority(t *testing.T) {
 	}
 }
 
+func TestGetNodeClassificationByLabels(t *testing.T) {
+	const machine01 = "machine01"
+	const machine02 = "machine02"
+	const zoneA = "zoneA"
+	zone1 := map[string]string{
+		"zone": zoneA,
+	}
+	labeledNodes := map[string]map[string]string{
+		machine01: zone1,
+	}
+	expectedNonLabeledNodes := []string{machine02}
+	serviceAffinity := ServiceAntiAffinity{label: "zone"}
+	newLabeledNodes, noNonLabeledNodes := serviceAffinity.getNodeClassificationByLabels(makeLabeledNodeList(labeledNodes))
+	noLabeledNodes, newnonLabeledNodes := serviceAffinity.getNodeClassificationByLabels(makeNodeList(expectedNonLabeledNodes))
+	label, _ := newLabeledNodes[machine01]
+	if label != zoneA && len(noNonLabeledNodes) != 0 {
+		t.Errorf("Expected only labeled node with label zoneA and no noNonLabeledNodes")
+	}
+	if len(noLabeledNodes) != 0 && newnonLabeledNodes[0] != machine02 {
+		t.Errorf("Expected only non labled nodes")
+	}
+}
+
 func makeLabeledNodeList(nodeMap map[string]map[string]string) []*v1.Node {
 	nodes := make([]*v1.Node, 0, len(nodeMap))
 	for nodeName, labels := range nodeMap {