From d900fbf6fc448f4eff36eb2d1438b73d64de51f3 Mon Sep 17 00:00:00 2001
From: Brendan Burns <bburns@google.com>
Date: Tue, 23 Sep 2014 16:14:54 -0700
Subject: [PATCH] Generalize the fit scheduler. Extract the "fit predicate"
 from random fit to make it extensible/plugable

---
 pkg/scheduler/randomfit.go      | 52 ++++++++++++++++++++++++---------
 pkg/scheduler/randomfit_test.go | 38 ++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 13 deletions(-)

diff --git a/pkg/scheduler/randomfit.go b/pkg/scheduler/randomfit.go
index 92e66dd5cc7..836cbd6c632 100644
--- a/pkg/scheduler/randomfit.go
+++ b/pkg/scheduler/randomfit.go
@@ -25,21 +25,49 @@ import (
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
 )
 
+// FitPredicate is a function that indicates if a pod fits into an existing node.
+type FitPredicate func(pod api.Pod, existingPods []api.Pod, node string) (bool, error)
+
 // RandomFitScheduler is a Scheduler which schedules a Pod on a random machine which matches its requirement.
 type RandomFitScheduler struct {
 	podLister  PodLister
+	predicates []FitPredicate
 	random     *rand.Rand
 	randomLock sync.Mutex
 }
 
+// NewRandomFitScheduler creates a random fit scheduler with the default set of fit predicates
 func NewRandomFitScheduler(podLister PodLister, random *rand.Rand) Scheduler {
+	return NewRandomFitSchedulerWithPredicates(podLister, random, []FitPredicate{podFitsPorts})
+}
+
+// NewRandomFitScheduler creates a random fit scheduler with the specified set of fit predicates.
+// All predicates must be true for the pod to be considered a fit.
+func NewRandomFitSchedulerWithPredicates(podLister PodLister, random *rand.Rand, predicates []FitPredicate) Scheduler {
 	return &RandomFitScheduler{
-		podLister: podLister,
-		random:    random,
+		podLister:  podLister,
+		random:     random,
+		predicates: predicates,
 	}
 }
 
-func (s *RandomFitScheduler) containsPort(pod api.Pod, port api.Port) bool {
+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
+				}
+			}
+		}
+	}
+	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 {
@@ -69,16 +97,14 @@ func (s *RandomFitScheduler) Schedule(pod api.Pod, minionLister MinionLister) (s
 	var machineOptions []string
 	for _, machine := range machines {
 		podFits := true
-		for _, scheduledPod := range machineToPods[machine] {
-			for _, container := range pod.DesiredState.Manifest.Containers {
-				for _, port := range container.Ports {
-					if port.HostPort == 0 {
-						continue
-					}
-					if s.containsPort(scheduledPod, port) {
-						podFits = false
-					}
-				}
+		for _, predicate := range s.predicates {
+			fits, err := predicate(pod, machineToPods[machine], machine)
+			if err != nil {
+				return "", err
+			}
+			if !fits {
+				podFits = false
+				break
 			}
 		}
 		if podFits {
diff --git a/pkg/scheduler/randomfit_test.go b/pkg/scheduler/randomfit_test.go
index 8b139cf1d3a..f4c3b908f33 100644
--- a/pkg/scheduler/randomfit_test.go
+++ b/pkg/scheduler/randomfit_test.go
@@ -76,3 +76,41 @@ func TestRandomFitSchedulerFirstScheduledImpossible(t *testing.T) {
 	}
 	st.expectFailure(newPod("", 8080, 8081))
 }
+
+func falsePredicate(pod api.Pod, existingPods []api.Pod, node string) (bool, error) {
+	return false, nil
+}
+
+func truePredicate(pod api.Pod, existingPods []api.Pod, node string) (bool, error) {
+	return true, nil
+}
+
+func TestRandomFitSchedulerFirstScheduledComplicatedWithMultiplePredicates(t *testing.T) {
+	fakeRegistry := FakePodLister{
+		newPod("m1", 80, 8080),
+		newPod("m2", 8081, 8082, 8083),
+		newPod("m3", 80, 443, 8085),
+	}
+	r := rand.New(rand.NewSource(0))
+	st := schedulerTester{
+		t:            t,
+		scheduler:    NewRandomFitSchedulerWithPredicates(fakeRegistry, r, []FitPredicate{podFitsPorts, truePredicate}),
+		minionLister: FakeMinionLister{"m1", "m2", "m3"},
+	}
+	st.expectSchedule(newPod("", 8080, 8081), "m3")
+}
+
+func TestRandomFitSchedulerFailureManyPredicates(t *testing.T) {
+	fakeRegistry := FakePodLister{
+		newPod("m1", 8080),
+		newPod("m2", 8081),
+		newPod("m3", 8080),
+	}
+	r := rand.New(rand.NewSource(0))
+	st := schedulerTester{
+		t:            t,
+		scheduler:    NewRandomFitSchedulerWithPredicates(fakeRegistry, r, []FitPredicate{truePredicate, falsePredicate}),
+		minionLister: FakeMinionLister{"m1", "m2", "m3"},
+	}
+	st.expectFailure(newPod("", 8080, 8081))
+}