feat: refactor to maintain policy state in-memory

.github/workflows/workflow.yml

@@ -50,7 +50,7 @@ jobs:
   e2e-test:
     name: "E2E Test"
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
     steps:
     - name: Set up Go 1.17
       uses: actions/setup-go@v2

pkg/plugins/placementpolicy/core/core.go

@@ -5,52 +5,36 @@ import (
 	"sort"
 
 	"github.com/Azure/placement-policy-scheduler-plugins/apis/v1alpha1"
-	ppclientset "github.com/Azure/placement-policy-scheduler-plugins/pkg/client/clientset/versioned"
 	ppinformers "github.com/Azure/placement-policy-scheduler-plugins/pkg/client/informers/externalversions/apis/v1alpha1"
 	pplisters "github.com/Azure/placement-policy-scheduler-plugins/pkg/client/listers/apis/v1alpha1"
 	"github.com/Azure/placement-policy-scheduler-plugins/pkg/utils"
 
 	corev1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
-	"k8s.io/client-go/kubernetes"
-	corelisters "k8s.io/client-go/listers/core/v1"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
 )
 
 // Manager defines the interfaces for PlacementPolicy management.
 type Manager interface {
 	GetPlacementPolicyForPod(context.Context, *corev1.Pod) (*v1alpha1.PlacementPolicy, error)
-	GetPodsWithLabels(context.Context, map[string]string) ([]*corev1.Pod, error)
-	AnnotatePod(context.Context, *corev1.Pod, *v1alpha1.PlacementPolicy, bool) (*corev1.Pod, error)
-	GetPlacementPolicy(context.Context, string, string) (*v1alpha1.PlacementPolicy, error)
+	GetPolicyInfo(*v1alpha1.PlacementPolicy) *PolicyInfo
+	MatchPodToPolicy(context.Context, *corev1.Pod, *PolicyInfo) (*PolicyInfo, error)
+	AddPodToPolicy(context.Context, *corev1.Pod, *PolicyInfo) (*PolicyInfo, error)
+	RemovePodFromPolicy(*corev1.Pod) error
 }
 
 type PlacementPolicyManager struct {
-	// client is a clientset for the kube API server.
-	client kubernetes.Interface
-	// client is a placementPolicy client
-	ppClient ppclientset.Interface
-	// podLister is pod lister
-	podLister corelisters.PodLister
-	// snapshotSharedLister is pod shared list
-	snapshotSharedLister framework.SharedLister
 	// ppLister is placementPolicy lister
 	ppLister pplisters.PlacementPolicyLister
+	// available policies by namespace
+	policies PolicyInfos
 }
 
 func NewPlacementPolicyManager(
-	client kubernetes.Interface,
-	ppClient ppclientset.Interface,
-	snapshotSharedLister framework.SharedLister,
-	ppInformer ppinformers.PlacementPolicyInformer,
-	podLister corelisters.PodLister) *PlacementPolicyManager {
+	ppInformer ppinformers.PlacementPolicyInformer) *PlacementPolicyManager {
 	return &PlacementPolicyManager{
-		client:               client,
-		ppClient:             ppClient,
-		snapshotSharedLister: snapshotSharedLister,
-		ppLister:             ppInformer.Lister(),
-		podLister:            podLister,
+		ppLister: ppInformer.Lister(),
+		policies: NewPolicyInfos(),
 	}
 }
 
@@ -73,38 +57,117 @@ func (m *PlacementPolicyManager) GetPlacementPolicyForPod(ctx context.Context, p
 	return ppList[0], nil
 }
 
-func (m *PlacementPolicyManager) GetPodsWithLabels(ctx context.Context, podLabels map[string]string) ([]*corev1.Pod, error) {
-	return m.podLister.List(labels.Set(podLabels).AsSelector())
+func (m *PlacementPolicyManager) filterPlacementPolicyList(ppList []*v1alpha1.PlacementPolicy, pod *corev1.Pod) []*v1alpha1.PlacementPolicy {
+	var filteredPPList []*v1alpha1.PlacementPolicy
+	for _, pp := range ppList {
+		labels := pp.Spec.PodSelector.MatchLabels
+		if utils.HasMatchingLabels(pod.Labels, labels) {
+			filteredPPList = append(filteredPPList, pp)
+		}
+	}
+	return filteredPPList
+}
+
+func (m *PlacementPolicyManager) GetPolicyInfo(pp *v1alpha1.PlacementPolicy) *PolicyInfo {
+	policyNamespace := pp.Namespace
+	policyName := pp.Name
+	corePolicy := pp.Spec.Policy
+
+	namespacePolicies, namespaceExists := m.policies[policyNamespace]
+
+	if namespaceExists {
+		policy, policyExists := namespacePolicies[policyName]
+
+		if policyExists {
+			return policy
+		}
+
+		created := newPolicyInfo(policyNamespace, policyName, corePolicy.Action, corePolicy.TargetSize)
+		m.policies[policyNamespace][policyName] = created
+		return created
+	}
+
+	m.policies[policyNamespace] = make(map[string]*PolicyInfo)
+	created := newPolicyInfo(policyNamespace, policyName, corePolicy.Action, corePolicy.TargetSize)
+	m.policies[policyNamespace][policyName] = created
+	return created
 }
 
-// AnnotatePod annotates the pod with the placement policy.
-func (m *PlacementPolicyManager) AnnotatePod(ctx context.Context, pod *corev1.Pod, pp *v1alpha1.PlacementPolicy, preferredNodeWithMatchingLabels bool) (*corev1.Pod, error) {
-	annotations := map[string]string{}
-	if pod.Annotations != nil {
-		annotations = pod.Annotations
+type PodAction int16
+
+const (
+	Match PodAction = iota
+	Add
+	Remove
+)
+
+func (m *PlacementPolicyManager) MatchPodToPolicy(ctx context.Context, pod *corev1.Pod, policy *PolicyInfo) (*PolicyInfo, error) {
+	matchError := policy.addMatch(pod)
+	if matchError != nil {
+		return nil, matchError
 	}
 
-	preference := "false"
-	if preferredNodeWithMatchingLabels {
-		preference = "true"
+	m.updatePolicies(policy, Match)
+	return policy, nil
+}
+
+func (m *PlacementPolicyManager) AddPodToPolicy(ctx context.Context, pod *corev1.Pod, policy *PolicyInfo) (*PolicyInfo, error) {
+	addError := policy.addPodIfNotPresent(pod)
+	if addError != nil {
+		return nil, addError
 	}
-	annotations[v1alpha1.PlacementPolicyAnnotationKey] = pp.Name
-	annotations[v1alpha1.PlacementPolicyPreferenceAnnotationKey] = preference
-	pod.Annotations = annotations
-	return m.client.CoreV1().Pods(pod.Namespace).Update(ctx, pod, metav1.UpdateOptions{})
+
+	m.updatePolicies(policy, Add)
+	return policy, nil
 }
 
-func (m *PlacementPolicyManager) GetPlacementPolicy(ctx context.Context, namespace, name string) (*v1alpha1.PlacementPolicy, error) {
-	return m.ppLister.PlacementPolicies(namespace).Get(name)
+func (m *PlacementPolicyManager) RemovePodFromPolicy(pod *corev1.Pod) error {
+	key, keyError := framework.GetPodKey(pod)
+	if keyError != nil {
+		return keyError
+	}
+
+	podNamespace := pod.Namespace
+	ppList := m.policies[podNamespace]
+
+	if ppList != nil {
+		var matchingPolicy *PolicyInfo
+
+		for _, pp := range ppList {
+			if pp.PodQualifiesForPolicy(key) {
+				matchingPolicy = pp
+				break
+			}
+		}
+
+		if matchingPolicy != nil {
+			removeError := matchingPolicy.removePodIfPresent(pod)
+			if removeError != nil {
+				return removeError
+			}
+
+			m.updatePolicies(matchingPolicy, Remove)
+		}
+	}
+	return nil
 }
 
-func (m *PlacementPolicyManager) filterPlacementPolicyList(ppList []*v1alpha1.PlacementPolicy, pod *corev1.Pod) []*v1alpha1.PlacementPolicy {
-	var filteredPPList []*v1alpha1.PlacementPolicy
-	for _, pp := range ppList {
-		labels := pp.Spec.PodSelector.MatchLabels
-		if utils.HasMatchingLabels(pod.Labels, labels) {
-			filteredPPList = append(filteredPPList, pp)
+func (m *PlacementPolicyManager) updatePolicies(policy *PolicyInfo, act PodAction) {
+	namespace := policy.Namespace
+	name := policy.Name
+
+	if act == Remove {
+		qualifyingCount := len(policy.qualifiedPods)
+
+		if qualifyingCount == 0 {
+			delete(m.policies[namespace], name)
+			return
 		}
+
+		m.policies[namespace][name] = policy
+		return
 	}
-	return filteredPPList
+
+	existing := m.policies[namespace][name]
+	m.policies[namespace][name] = policy.merge(existing)
 }

pkg/plugins/placementpolicy/core/policy.go

@@ -0,0 +1,167 @@
+package core
+
+import (
+	"github.com/Azure/placement-policy-scheduler-plugins/apis/v1alpha1"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/kubernetes/pkg/scheduler/framework"
+)
+
+type PolicyInfo struct {
+	Namespace     string
+	Name          string
+	Action        v1alpha1.Action
+	TargetSize    *intstr.IntOrString
+	matchedPods   sets.String
+	qualifiedPods sets.String
+	managedPods   sets.String
+	targetMet     bool
+}
+
+func newPolicyInfo(namespace string, name string, action v1alpha1.Action, targetSize *intstr.IntOrString) *PolicyInfo {
+	policy := &PolicyInfo{
+		Namespace:     namespace,
+		Name:          name,
+		Action:        action,
+		TargetSize:    targetSize,
+		matchedPods:   sets.NewString(),
+		qualifiedPods: sets.NewString(),
+		managedPods:   sets.NewString(),
+		targetMet:     false,
+	}
+	return policy
+}
+
+type PolicyInfos map[string]map[string]*PolicyInfo
+
+func NewPolicyInfos() PolicyInfos {
+	return make(PolicyInfos)
+}
+
+func (p *PolicyInfo) merge(existing *PolicyInfo) *PolicyInfo {
+	existing.Action = p.Action
+	existing.TargetSize = p.TargetSize
+	existing.targetMet = p.targetMet
+
+	tempMatched := sets.NewString()
+	if len(p.matchedPods) > 0 {
+		pods := p.matchedPods.List()
+		for _, pod := range pods {
+			tempMatched = tempMatched.Insert(pod)
+		}
+	}
+	existing.matchedPods = tempMatched
+
+	tempQualified := sets.NewString()
+	if len(p.qualifiedPods) > 0 {
+		pods := p.qualifiedPods.List()
+		for _, pod := range pods {
+			tempQualified = tempQualified.Insert(pod)
+		}
+	}
+	existing.qualifiedPods = tempQualified
+
+	tempManaged := sets.NewString()
+	if len(p.managedPods) > 0 {
+		pods := p.managedPods.List()
+		for _, pod := range pods {
+			tempManaged = tempManaged.Insert(pod)
+		}
+	}
+	existing.managedPods = tempManaged
+
+	return existing
+}
+
+func (p *PolicyInfo) addMatch(pod *corev1.Pod) error {
+	key, keyError := framework.GetPodKey(pod)
+	if keyError != nil {
+		return keyError
+	}
+
+	p.matchedPods = p.managedPods.Insert(key)
+	return nil
+}
+
+func (p *PolicyInfo) removePodIfPresent(pod *corev1.Pod) error {
+	key, keyError := framework.GetPodKey(pod)
+	if keyError != nil {
+		return keyError
+	}
+
+	if !p.PodQualifiesForPolicy(key) {
+		return nil
+	}
+
+	p.qualifiedPods = p.qualifiedPods.Delete(key)
+
+	if p.PodIsManagedByPolicy(key) {
+		p.managedPods = p.managedPods.Delete(key)
+	}
+
+	err := p.setTargetMet()
+	return err
+}
+
+func (p *PolicyInfo) addPodIfNotPresent(pod *corev1.Pod) error {
+	key, keyError := framework.GetPodKey(pod)
+	if keyError != nil {
+		return keyError
+	}
+
+	//if pod is already in the list, do nothing
+	if p.PodQualifiesForPolicy(key) {
+		return nil
+	}
+
+	p.matchedPods = p.matchedPods.Delete(key) //once added, don't need to worry about matched anymore
+
+	p.qualifiedPods = p.qualifiedPods.Insert(key)
+
+	targetError := p.setTargetMet()
+	if targetError != nil {
+		return targetError
+	}
+
+	//if target met, pod doesn't need to be managed
+	if p.targetMet {
+		return nil
+	}
+
+	p.managedPods = p.managedPods.Insert(key)
+	err := p.setTargetMet()
+	return err
+}
+
+func (p *PolicyInfo) setTargetMet() error {
+	specTarget := p.TargetSize
+	lenAllPods := len(p.qualifiedPods)
+
+	target, calcError := intstr.GetScaledValueFromIntOrPercent(specTarget, lenAllPods, false)
+
+	if calcError != nil {
+		return calcError
+	}
+
+	if p.Action == v1alpha1.ActionMustNot {
+		target = lenAllPods - target
+	}
+
+	managedCount := len(p.managedPods)
+	p.targetMet = managedCount >= target //since the TargetSize is rounded down, the expectation that it will only meet/equal and never exceed
+	return nil
+}
+
+func (p *PolicyInfo) PodMatchesPolicy(podKey string) bool {
+	return p.matchedPods.Has(podKey)
+}
+
+func (p *PolicyInfo) PodQualifiesForPolicy(podKey string) bool {
+	return p.qualifiedPods.Has(podKey)
+}
+
+func (p *PolicyInfo) PodIsManagedByPolicy(podKey string) bool {
+	return p.managedPods.Has(podKey)
+}