refactor the internal logic of podGroupManager

Author: googs1025

pkg/coscheduling/core/core.go

@@ -66,6 +66,7 @@ func (s *PermitState) Clone() framework.StateData {
 type Manager interface {
 	PreFilter(context.Context, *corev1.Pod) error
 	Permit(context.Context, *framework.CycleState, *corev1.Pod) Status
+	Reserve(context.Context, *corev1.Pod)
 	Unreserve(context.Context, *corev1.Pod)
 	GetPodGroup(context.Context, *corev1.Pod) (string, *v1alpha1.PodGroup)
 	GetAssignedPodCount(string) int
@@ -95,7 +96,7 @@ type PodGroupManager struct {
 	sync.RWMutex
 }
 
-func AddPodFactory(pgMgr *PodGroupManager) func(obj interface{}) {
+func AddPodFactory(ctx context.Context, pgMgr *PodGroupManager) func(obj interface{}) {
 	return func(obj interface{}) {
 		p, ok := obj.(*corev1.Pod)
 		if !ok {
@@ -104,22 +105,52 @@ func AddPodFactory(pgMgr *PodGroupManager) func(obj interface{}) {
 		if p.Spec.NodeName == "" {
 			return
 		}
-		pgFullName, _ := pgMgr.GetPodGroup(context.Background(), p)
-		if pgFullName == "" {
+		pgMgr.Reserve(ctx, p)
+	}
+}
+
+func UpdatePodFactory(ctx context.Context, pgMgr *PodGroupManager) func(oldObj interface{}, newObj interface{}) {
+	return func(oldObj interface{}, newObj interface{}) {
+		oldPod, ok := oldObj.(*corev1.Pod)
+		if !ok {
 			return
 		}
-		pgMgr.RWMutex.Lock()
-		defer pgMgr.RWMutex.Unlock()
-		if assigned, exist := pgMgr.assignedPodsByPG[pgFullName]; exist {
-			assigned.Insert(p.Name)
-		} else {
-			pgMgr.assignedPodsByPG[pgFullName] = sets.New(p.Name)
+		newPod, ok := newObj.(*corev1.Pod)
+		if !ok {
+			return
+		}
+		// If the pod is assumed or bound, it should be reserved.
+		if oldPod.Spec.NodeName == "" && newPod.Spec.NodeName != "" {
+			pgMgr.Reserve(ctx, newPod)
+		}
+	}
+}
+
+func DelPodFactory(ctx context.Context, pgMgr *PodGroupManager) func(obj interface{}) {
+	return func(obj interface{}) {
+		var pod *corev1.Pod
+		switch t := obj.(type) {
+		case *corev1.Pod:
+			pod = t
+		case cache.DeletedFinalStateUnknown:
+			var ok bool
+			if pod, ok = t.Obj.(*corev1.Pod); !ok {
+				return
+			}
+		default:
+			return
+		}
+		if pod.Spec.NodeName == "" {
+			return
 		}
+		pgMgr.Unreserve(ctx, pod)
 	}
 }
 
 // NewPodGroupManager creates a new operation object.
-func NewPodGroupManager(client client.Client, snapshotSharedLister framework.SharedLister, scheduleTimeout *time.Duration, podInformer informerv1.PodInformer) *PodGroupManager {
+func NewPodGroupManager(
+	ctx context.Context, client client.Client, snapshotSharedLister framework.SharedLister,
+	scheduleTimeout *time.Duration, podInformer informerv1.PodInformer) *PodGroupManager {
 	pgMgr := &PodGroupManager{
 		client:               client,
 		snapshotSharedLister: snapshotSharedLister,
@@ -130,30 +161,9 @@ func NewPodGroupManager(client client.Client, snapshotSharedLister framework.Sha
 		assignedPodsByPG:     map[string]sets.Set[string]{},
 	}
 	podInformer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
-		AddFunc: AddPodFactory(pgMgr),
-		DeleteFunc: func(obj interface{}) {
-			switch t := obj.(type) {
-			case *corev1.Pod:
-				pod := t
-				if pod.Spec.NodeName == "" {
-					return
-				}
-				pgMgr.Unreserve(context.Background(), pod)
-				return
-			case cache.DeletedFinalStateUnknown:
-				pod, ok := t.Obj.(*corev1.Pod)
-				if !ok {
-					return
-				}
-				if pod.Spec.NodeName == "" {
-					return
-				}
-				pgMgr.Unreserve(context.Background(), pod)
-				return
-			default:
-				return
-			}
-		},
+		AddFunc:    AddPodFactory(ctx, pgMgr),
+		UpdateFunc: UpdatePodFactory(ctx, pgMgr),
+		DeleteFunc: DelPodFactory(ctx, pgMgr),
 	})
 	return pgMgr
 }
@@ -312,6 +322,21 @@ func (pgMgr *PodGroupManager) Permit(ctx context.Context, state *framework.Cycle
 	return Wait
 }
 
+// Reserve adds Pod to the assignedPodsByPG map when it is scheduled
+func (pgMgr *PodGroupManager) Reserve(ctx context.Context, pod *corev1.Pod) {
+	pgFullName, _ := pgMgr.GetPodGroup(ctx, pod)
+	if pgFullName == "" {
+		return
+	}
+	pgMgr.RWMutex.Lock()
+	defer pgMgr.RWMutex.Unlock()
+	if assigned, exist := pgMgr.assignedPodsByPG[pgFullName]; exist {
+		assigned.Insert(pod.Name)
+	} else {
+		pgMgr.assignedPodsByPG[pgFullName] = sets.New(pod.Name)
+	}
+}
+
 // Unreserve invalidates assigned pod from assignedPodsByPG when schedule or bind failed.
 func (pgMgr *PodGroupManager) Unreserve(ctx context.Context, pod *corev1.Pod) {
 	pgFullName, _ := pgMgr.GetPodGroup(ctx, pod)

pkg/coscheduling/core/core_test.go

@@ -266,13 +266,13 @@ func TestPermit(t *testing.T) {
 			informerFactory := informers.NewSharedInformerFactory(cs, 0)
 			podInformer := informerFactory.Core().V1().Pods()
 
-			pgMgr := NewPodGroupManager(client, tu.NewFakeSharedLister(tt.existingPods, nodes), &scheduleTimeout, podInformer)
+			pgMgr := NewPodGroupManager(ctx, client, tu.NewFakeSharedLister(tt.existingPods, nodes), &scheduleTimeout, podInformer)
 
 			informerFactory.Start(ctx.Done())
 			if !clicache.WaitForCacheSync(ctx.Done(), podInformer.Informer().HasSynced) {
 				t.Fatal("WaitForCacheSync failed")
 			}
-			addFunc := AddPodFactory(pgMgr)
+			addFunc := AddPodFactory(ctx, pgMgr)
 			for _, p := range tt.existingPods {
 				podInformer.Informer().GetStore().Add(p)
 				// we call add func here because we can not ensure existing pods are added before premit are called

pkg/coscheduling/coscheduling.go

@@ -84,6 +84,7 @@ func New(ctx context.Context, obj runtime.Object, handle framework.Handle) (fram
 
 	scheduleTimeDuration := time.Duration(args.PermitWaitingTimeSeconds) * time.Second
 	pgMgr := core.NewPodGroupManager(
+		ctx,
 		client,
 		handle.SnapshotSharedLister(),
 		&scheduleTimeDuration,
@@ -246,6 +247,7 @@ func (cs *Coscheduling) Permit(ctx context.Context, state *framework.CycleState,
 
 // Reserve is the functions invoked by the framework at "reserve" extension point.
 func (cs *Coscheduling) Reserve(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) *framework.Status {
+	cs.pgMgr.Reserve(ctx, pod)
 	return nil
 }
 

pkg/coscheduling/coscheduling_test.go

@@ -114,6 +114,7 @@ func TestPodGroupBackoffTime(t *testing.T) {
 			}
 
 			pgMgr := core.NewPodGroupManager(
+				ctx,
 				client,
 				tu.NewFakeSharedLister(tt.pods, nodes),
 				// In this UT, 5 seconds should suffice to test the PreFilter's return code.
@@ -421,7 +422,7 @@ func TestLess(t *testing.T) {
 			informerFactory := informers.NewSharedInformerFactory(cs, 0)
 			podInformer := informerFactory.Core().V1().Pods()
 
-			pl := &Coscheduling{pgMgr: core.NewPodGroupManager(client, nil, nil, podInformer)}
+			pl := &Coscheduling{pgMgr: core.NewPodGroupManager(ctx, client, nil, nil, podInformer)}
 
 			informerFactory.Start(ctx.Done())
 			if !clicache.WaitForCacheSync(ctx.Done(), podInformer.Informer().HasSynced) {
@@ -516,7 +517,7 @@ func TestPermit(t *testing.T) {
 
 			pl := &Coscheduling{
 				frameworkHandler: f,
-				pgMgr:            core.NewPodGroupManager(client, tu.NewFakeSharedLister(nil, nodes), nil, podInformer),
+				pgMgr:            core.NewPodGroupManager(ctx, client, tu.NewFakeSharedLister(nil, nodes), nil, podInformer),
 				scheduleTimeout:  &scheduleTimeout,
 			}
 
@@ -623,6 +624,7 @@ func TestPostFilter(t *testing.T) {
 			informerFactory := informers.NewSharedInformerFactory(cs, 0)
 			podInformer := informerFactory.Core().V1().Pods()
 			pgMgr := core.NewPodGroupManager(
+				ctx,
 				client,
 				tu.NewFakeSharedLister(tt.existingPods, nodes),
 				&scheduleTimeout,
@@ -638,7 +640,7 @@ func TestPostFilter(t *testing.T) {
 			if !clicache.WaitForCacheSync(ctx.Done(), podInformer.Informer().HasSynced) {
 				t.Fatal("WaitForCacheSync failed")
 			}
-			addFunc := core.AddPodFactory(pgMgr)
+			addFunc := core.AddPodFactory(ctx, pgMgr)
 			for _, p := range tt.existingPods {
 				podInformer.Informer().GetStore().Add(p)
 				// we call add func here because we can not ensure existing pods are added before premit are called