Add support to recompute partial predicate metadata upon adding/removing pods

plugin/pkg/scheduler/algorithm/predicates/metadata.go

@@ -17,17 +17,53 @@ limitations under the License.
 package predicates
 
 import (
+	"fmt"
 	"github.com/golang/glog"
 	"k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/kubernetes/plugin/pkg/scheduler/algorithm"
 	"k8s.io/kubernetes/plugin/pkg/scheduler/schedulercache"
 	schedutil "k8s.io/kubernetes/plugin/pkg/scheduler/util"
+	"sync"
 )
 
 type PredicateMetadataFactory struct {
 	podLister algorithm.PodLister
 }
 
+//  Note that predicateMetadata and matchingPodAntiAffinityTerm need to be declared in the same file
+//  due to the way declarations are processed in predicate declaration unit tests.
+type matchingPodAntiAffinityTerm struct {
+	term *v1.PodAffinityTerm
+	node *v1.Node
+}
+
+// NOTE: When new fields are added/removed or logic is changed, please make sure
+// that RemovePod and AddPod functions are updated to work with the new changes.
+type predicateMetadata struct {
+	pod           *v1.Pod
+	podBestEffort bool
+	podRequest    *schedulercache.Resource
+	podPorts      map[int]bool
+	//key is a pod full name with the anti-affinity rules.
+	matchingAntiAffinityTerms          map[string][]matchingPodAntiAffinityTerm
+	serviceAffinityInUse               bool
+	serviceAffinityMatchingPodList     []*v1.Pod
+	serviceAffinityMatchingPodServices []*v1.Service
+}
+
+// PredicateMetadataProducer: Helper types/variables...
+type PredicateMetadataProducer func(pm *predicateMetadata)
+
+var predicateMetaProducerRegisterLock sync.Mutex
+var predicateMetadataProducers map[string]PredicateMetadataProducer = make(map[string]PredicateMetadataProducer)
+
+func RegisterPredicateMetadataProducer(predicateName string, precomp PredicateMetadataProducer) {
+	predicateMetaProducerRegisterLock.Lock()
+	defer predicateMetaProducerRegisterLock.Unlock()
+	predicateMetadataProducers[predicateName] = precomp
+}
+
 func NewPredicateMetadataFactory(podLister algorithm.PodLister) algorithm.MetadataProducer {
 	factory := &PredicateMetadataFactory{
 		podLister,
@@ -52,9 +88,72 @@ func (pfactory *PredicateMetadataFactory) GetMetadata(pod *v1.Pod, nodeNameToInf
 		podPorts:                  schedutil.GetUsedPorts(pod),
 		matchingAntiAffinityTerms: matchingTerms,
 	}
-	for predicateName, precomputeFunc := range predicatePrecomputations {
+	for predicateName, precomputeFunc := range predicateMetadataProducers {
 		glog.V(10).Infof("Precompute: %v", predicateName)
 		precomputeFunc(predicateMetadata)
 	}
 	return predicateMetadata
 }
+
+// RemovePod changes predicateMetadata assuming that the given `deletedPod` is
+// deleted from the system.
+func (meta *predicateMetadata) RemovePod(deletedPod *v1.Pod) error {
+	deletedPodFullName := schedutil.GetPodFullName(deletedPod)
+	if deletedPodFullName == schedutil.GetPodFullName(meta.pod) {
+		return fmt.Errorf("deletedPod and meta.pod must not be the same.")
+	}
+	// Delete any anti-affinity rule from the deletedPod.
+	delete(meta.matchingAntiAffinityTerms, deletedPodFullName)
+	// All pods in the serviceAffinityMatchingPodList are in the same namespace.
+	// So, if the namespace of the first one is not the same as the namespace of the
+	// deletedPod, we don't need to check the list, as deletedPod isn't in the list.
+	if meta.serviceAffinityInUse &&
+		len(meta.serviceAffinityMatchingPodList) > 0 &&
+		deletedPod.Namespace == meta.serviceAffinityMatchingPodList[0].Namespace {
+		for i, pod := range meta.serviceAffinityMatchingPodList {
+			if schedutil.GetPodFullName(pod) == deletedPodFullName {
+				meta.serviceAffinityMatchingPodList = append(
+					meta.serviceAffinityMatchingPodList[:i],
+					meta.serviceAffinityMatchingPodList[i+1:]...)
+				break
+			}
+		}
+	}
+	return nil
+}
+
+// AddPod changes predicateMetadata assuming that `newPod` is added to the
+// system.
+func (meta *predicateMetadata) AddPod(addedPod *v1.Pod, nodeInfo *schedulercache.NodeInfo) error {
+	addedPodFullName := schedutil.GetPodFullName(addedPod)
+	if addedPodFullName == schedutil.GetPodFullName(meta.pod) {
+		return fmt.Errorf("addedPod and meta.pod must not be the same.")
+	}
+	if nodeInfo.Node() == nil {
+		return fmt.Errorf("Invalid node in nodeInfo.")
+	}
+	// Add matching anti-affinity terms of the addedPod to the map.
+	podMatchingTerms, err := getMatchingAntiAffinityTermsOfExistingPod(meta.pod, addedPod, nodeInfo.Node())
+	if err != nil {
+		return err
+	}
+	if len(podMatchingTerms) > 0 {
+		existingTerms, found := meta.matchingAntiAffinityTerms[addedPodFullName]
+		if found {
+			meta.matchingAntiAffinityTerms[addedPodFullName] = append(existingTerms,
+				podMatchingTerms...)
+		} else {
+			meta.matchingAntiAffinityTerms[addedPodFullName] = podMatchingTerms
+		}
+	}
+	// If addedPod is in the same namespace as the meta.pod, update the list
+	// of matching pods if applicable.
+	if meta.serviceAffinityInUse && addedPod.Namespace == meta.pod.Namespace {
+		selector := CreateSelectorFromLabels(meta.pod.Labels)
+		if selector.Matches(labels.Set(addedPod.Labels)) {
+			meta.serviceAffinityMatchingPodList = append(meta.serviceAffinityMatchingPodList,
+				addedPod)
+		}
+	}
+	return nil
+}

plugin/pkg/scheduler/algorithm/predicates/predicates_test.go

@@ -1590,7 +1590,7 @@ func TestServiceAffinity(t *testing.T) {
 			// Reimplementing the logic that the scheduler implements: Any time it makes a predicate, it registers any precomputations.
 			predicate, precompute := NewServiceAffinityPredicate(schedulertesting.FakePodLister(test.pods), schedulertesting.FakeServiceLister(test.services), FakeNodeListInfo(nodes), test.labels)
 			// Register a precomputation or Rewrite the precomputation to a no-op, depending on the state we want to test.
-			RegisterPredicatePrecomputation("checkServiceAffinity-unitTestPredicate", func(pm *predicateMetadata) {
+			RegisterPredicateMetadataProducer("ServiceAffinityMetaProducer", func(pm *predicateMetadata) {
 				if !skipPrecompute {
 					precompute(pm)
 				}

plugin/pkg/scheduler/algorithm/types.go

@@ -80,6 +80,9 @@ type PodLister interface {
 	// We explicitly return []*v1.Pod, instead of v1.PodList, to avoid
 	// performing expensive copies that are unneeded.
 	List(labels.Selector) ([]*v1.Pod, error)
+	// This is similar to "List()", but the returned slice does not
+	// contain pods that don't pass `podFilter`.
+	FilteredList(podFilter schedulercache.PodFilter, selector labels.Selector) ([]*v1.Pod, error)
 }
 
 // ServiceLister interface represents anything that can produce a list of services; the list is consumed by a scheduler.

plugin/pkg/scheduler/factory/plugins.go

@@ -129,7 +129,7 @@ func RegisterCustomFitPredicate(policy schedulerapi.PredicatePolicy) string {
 				)
 
 				// Once we generate the predicate we should also Register the Precomputation
-				predicates.RegisterPredicatePrecomputation(policy.Name, precomputationFunction)
+				predicates.RegisterPredicateMetadataProducer(policy.Name, precomputationFunction)
 				return predicate
 			}
 		} else if policy.Argument.LabelsPresence != nil {

plugin/pkg/scheduler/schedulercache/cache.go

@@ -93,12 +93,17 @@ func (cache *schedulerCache) UpdateNodeNameToInfoMap(nodeNameToInfo map[string]*
 }
 
 func (cache *schedulerCache) List(selector labels.Selector) ([]*v1.Pod, error) {
+	alwaysTrue := func(p *v1.Pod) bool { return true }
+	return cache.FilteredList(alwaysTrue, selector)
+}
+
+func (cache *schedulerCache) FilteredList(podFilter PodFilter, selector labels.Selector) ([]*v1.Pod, error) {
 	cache.mu.Lock()
 	defer cache.mu.Unlock()
 	var pods []*v1.Pod
 	for _, info := range cache.nodes {
 		for _, pod := range info.pods {
-			if selector.Matches(labels.Set(pod.Labels)) {
+			if podFilter(pod) && selector.Matches(labels.Set(pod.Labels)) {
 				pods = append(pods, pod)
 			}
 		}

plugin/pkg/scheduler/schedulercache/interface.go

@@ -21,6 +21,8 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 )
 
+type PodFilter func(*v1.Pod) bool
+
 // Cache collects pods' information and provides node-level aggregated information.
 // It's intended for generic scheduler to do efficient lookup.
 // Cache's operations are pod centric. It does incremental updates based on pod events.
@@ -93,4 +95,7 @@ type Cache interface {
 
 	// List lists all cached pods (including assumed ones).
 	List(labels.Selector) ([]*v1.Pod, error)
+
+	// FilteredList returns all cached pods that pass the filter.
+	FilteredList(filter PodFilter, selector labels.Selector) ([]*v1.Pod, error)
 }

plugin/pkg/scheduler/schedulercache/node_info.go

@@ -26,6 +26,7 @@ import (
 	clientcache "k8s.io/client-go/tools/cache"
 	v1helper "k8s.io/kubernetes/pkg/api/v1/helper"
 	priorityutil "k8s.io/kubernetes/plugin/pkg/scheduler/algorithm/priorities/util"
+	"k8s.io/kubernetes/plugin/pkg/scheduler/util"
 )
 
 var emptyResource = Resource{}
@@ -458,3 +459,19 @@ func (n *NodeInfo) RemoveNode(node *v1.Node) error {
 func getPodKey(pod *v1.Pod) (string, error) {
 	return clientcache.MetaNamespaceKeyFunc(pod)
 }
+
+// Filter implements PodFilter interface. It returns false only if the pod node name
+// matches NodeInfo.node and the pod is not found in the pods list. Otherwise,
+// returns true.
+func (n *NodeInfo) Filter(pod *v1.Pod) bool {
+	pFullName := util.GetPodFullName(pod)
+	if pod.Spec.NodeName != n.node.Name {
+		return true
+	}
+	for _, p := range n.pods {
+		if util.GetPodFullName(p) == pFullName {
+			return true
+		}
+	}
+	return false
+}

plugin/pkg/scheduler/testing/fake_cache.go

@@ -57,3 +57,7 @@ func (f *FakeCache) UpdateNodeNameToInfoMap(infoMap map[string]*schedulercache.N
 }
 
 func (f *FakeCache) List(s labels.Selector) ([]*v1.Pod, error) { return nil, nil }
+
+func (f *FakeCache) FilteredList(filter schedulercache.PodFilter, selector labels.Selector) ([]*v1.Pod, error) {
+	return nil, nil
+}

plugin/pkg/scheduler/testing/fake_lister.go

@@ -25,6 +25,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	. "k8s.io/kubernetes/plugin/pkg/scheduler/algorithm"
+	"k8s.io/kubernetes/plugin/pkg/scheduler/schedulercache"
 )
 
 var _ NodeLister = &FakeNodeLister{}
@@ -52,6 +53,15 @@ func (f FakePodLister) List(s labels.Selector) (selected []*v1.Pod, err error) {
 	return selected, nil
 }
 
+func (f FakePodLister) FilteredList(podFilter schedulercache.PodFilter, s labels.Selector) (selected []*v1.Pod, err error) {
+	for _, pod := range f {
+		if podFilter(pod) && s.Matches(labels.Set(pod.Labels)) {
+			selected = append(selected, pod)
+		}
+	}
+	return selected, nil
+}
+
 var _ ServiceLister = &FakeServiceLister{}
 
 // FakeServiceLister implements ServiceLister on []v1.Service for test purposes.

plugin/pkg/scheduler/util/utils.go

@@ -39,3 +39,10 @@ func GetUsedPorts(pods ...*v1.Pod) map[int]bool {
 	}
 	return ports
 }
+
+// GetPodFullName returns a name that uniquely identifies a pod.
+func GetPodFullName(pod *v1.Pod) string {
+	// Use underscore as the delimiter because it is not allowed in pod name
+	// (DNS subdomain format).
+	return pod.Name + "_" + pod.Namespace
+}