Refactor PodFitsPorts

pkg/scheduler/predicates.go

@@ -163,30 +163,29 @@ func (n *NodeSelector) PodSelectorMatches(pod api.Pod, existingPods []api.Pod, n
 }
 
 func PodFitsPorts(pod api.Pod, existingPods []api.Pod, node string) (bool, error) {
-	for _, scheduledPod := range existingPods {
-		for _, container := range pod.DesiredState.Manifest.Containers {
-			for _, port := range container.Ports {
-				if port.HostPort == 0 {
-					continue
-				}
-				if containsPort(scheduledPod, port) {
-					return false, nil
-				}
-			}
+	existingPorts := getUsedPorts(existingPods...)
+	wantPorts := getUsedPorts(pod)
+	for wport := range wantPorts {
+		if wport == 0 {
+			continue
+		}
+		if existingPorts[wport] {
+			return false, nil
 		}
 	}
 	return true, nil
 }
 
-func containsPort(pod api.Pod, port api.Port) bool {
-	for _, container := range pod.DesiredState.Manifest.Containers {
-		for _, podPort := range container.Ports {
-			if podPort.HostPort == port.HostPort {
-				return true
+func getUsedPorts(pods ...api.Pod) map[int]bool {
+	ports := make(map[int]bool)
+	for _, pod := range pods {
+		for _, container := range pod.DesiredState.Manifest.Containers {
+			for _, podPort := range container.Ports {
+				ports[podPort.HostPort] = true
 			}
 		}
 	}
-	return false
+	return ports
 }
 
 // MapPodsToMachines obtains a list of pods and pivots that list into a map where the keys are host names

pkg/scheduler/predicates_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package scheduler
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
@@ -182,6 +183,42 @@ func TestPodFitsPorts(t *testing.T) {
 	}
 }
 
+func TestGetUsedPorts(t *testing.T) {
+	tests := []struct {
+		pods []api.Pod
+
+		ports map[int]bool
+	}{
+		{
+			[]api.Pod{
+				newPod("m1", 9090),
+			},
+			map[int]bool{9090: true},
+		},
+		{
+			[]api.Pod{
+				newPod("m1", 9090),
+				newPod("m1", 9091),
+			},
+			map[int]bool{9090: true, 9091: true},
+		},
+		{
+			[]api.Pod{
+				newPod("m1", 9090),
+				newPod("m2", 9091),
+			},
+			map[int]bool{9090: true, 9091: true},
+		},
+	}
+
+	for _, test := range tests {
+		ports := getUsedPorts(test.pods...)
+		if !reflect.DeepEqual(test.ports, ports) {
+			t.Errorf("expect %v, got %v", test.ports, ports)
+		}
+	}
+}
+
 func TestDiskConflicts(t *testing.T) {
 	volState := api.PodState{
 		Manifest: api.ContainerManifest{