ClusterThrottles support PodGroup aka Gang

.github/workflows/ci.yaml

@@ -71,4 +71,4 @@ jobs:
     - name: Validate GoReleaser
       uses: goreleaser/goreleaser-action@v2
       with:
-        args: release --snapshot --skip-publish --rm-dist --debug
+        args: release --snapshot --skip=publish --clean

.goreleaser.yaml

@@ -1,3 +1,4 @@
+version: 2
 before:
   hooks:
   - go mod download
@@ -62,4 +63,4 @@ docker_manifests:
 
 # change log will be generated by tagpr
 changelog:
-  skip: true
+  disable: true

pkg/apis/schedule/v1alpha1/clusterthrottle_types.go

@@ -27,7 +27,7 @@ type ClusterThrottleSpec struct {
 	Selector         ClusterThrottleSelector `json:"selector,omitempty"`
 }
 
-func (thr ClusterThrottle) CheckThrottledFor(pod *corev1.Pod, reservedResourceAmount ResourceAmount, isThrottledOnEqual bool) CheckThrottleStatus {
+func (thr ClusterThrottle) CheckThrottledFor(pod *corev1.Pod, reservedResourceAmount ResourceAmount, usedByGroup ResourceAmount, isThrottledOnEqual bool) CheckThrottleStatus {
 	threshold := thr.Spec.Threshold
 	if !thr.Status.CalculatedThreshold.CalculatedAt.Time.IsZero() {
 		threshold = thr.Status.CalculatedThreshold.Threshold
@@ -46,11 +46,16 @@ func (thr ClusterThrottle) CheckThrottledFor(pod *corev1.Pod, reservedResourceAm
 		return CheckThrottleStatusActive
 	}
 
-	used := ResourceAmount{}.Add(thr.Status.Used).Add(ResourceAmountOfPod(pod)).Add(reservedResourceAmount)
-	if threshold.IsThrottled(used, isThrottledOnEqual).IsThrottledFor(pod) {
+	usedWithPod := alreadyUsed.Add(ResourceAmountOfPod(pod))
+	if threshold.IsThrottled(usedWithPod, isThrottledOnEqual).IsThrottledFor(pod) {
 		return CheckThrottleStatusInsufficient
 	}
 
+	usedWithGroup := usedWithPod.Add(usedByGroup)
+	if threshold.IsThrottled(usedWithGroup, isThrottledOnEqual).IsThrottledFor(pod) {
+		return CheckThrottleStatusInsufficientIncludingPodGroup
+	}
+
 	return CheckThrottleStatusNotThrottled
 }
 

pkg/apis/schedule/v1alpha1/throttle_types.go

@@ -119,10 +119,11 @@ type ThrottleStatus struct {
 type CheckThrottleStatus string
 
 var (
-	CheckThrottleStatusNotThrottled                CheckThrottleStatus = "not-throttled"
-	CheckThrottleStatusActive                      CheckThrottleStatus = "active"
-	CheckThrottleStatusInsufficient                CheckThrottleStatus = "insufficient"
-	CheckThrottleStatusPodRequestsExceedsThreshold CheckThrottleStatus = "pod-requests-exceeds-threshold"
+	CheckThrottleStatusNotThrottled                  CheckThrottleStatus = "not-throttled"
+	CheckThrottleStatusActive                        CheckThrottleStatus = "active"
+	CheckThrottleStatusInsufficient                  CheckThrottleStatus = "insufficient"
+	CheckThrottleStatusInsufficientIncludingPodGroup CheckThrottleStatus = "insufficient-including-pod-group"
+	CheckThrottleStatusPodRequestsExceedsThreshold   CheckThrottleStatus = "pod-requests-exceeds-threshold"
 )
 
 func (thr Throttle) CheckThrottledFor(pod *corev1.Pod, reservedResourceAmount ResourceAmount, isThrottledOnEqual bool) CheckThrottleStatus {

pkg/controllers/clusterthrottle_controller.go

@@ -378,27 +378,86 @@ func (c *ClusterThrottleController) UnReserveOnClusterThrottle(pod *corev1.Pod,
 func (c *ClusterThrottleController) CheckThrottled(
 	pod *corev1.Pod,
 	isThrottledOnEqual bool,
+	groupNameAnnotation string,
 ) (
 	[]schedulev1alpha1.ClusterThrottle,
 	[]schedulev1alpha1.ClusterThrottle,
 	[]schedulev1alpha1.ClusterThrottle,
 	[]schedulev1alpha1.ClusterThrottle,
+	[]schedulev1alpha1.ClusterThrottle,
 	error,
 ) {
 	throttles, err := c.affectedClusterThrottles(pod)
 	if err != nil {
-		return nil, nil, nil, nil, err
+		return nil, nil, nil, nil, nil, err
 	}
 	affected := []schedulev1alpha1.ClusterThrottle{}
 	alreadyThrottled := []schedulev1alpha1.ClusterThrottle{}
 	insufficient := []schedulev1alpha1.ClusterThrottle{}
+	insufficientIncludingGroup := []schedulev1alpha1.ClusterThrottle{}
 	podRequestsExceedsThreshold := []schedulev1alpha1.ClusterThrottle{}
+
+	// Fetch the pods which have group name's annotation
+	var podGroup []*corev1.Pod
+	if groupNameAnnotation != "" {
+		groupName, isGroup := pod.Annotations[groupNameAnnotation]
+		if isGroup {
+			candidatePods, err := c.podInformer.Lister().Pods(pod.Namespace).List(labels.Everything())
+			if err != nil {
+				return nil, nil, nil, nil, nil, err
+			}
+
+			for _, candidatePod := range candidatePods {
+				if isScheduled(candidatePod) {
+					continue
+				}
+
+				if gn, ok := candidatePod.Annotations[groupNameAnnotation]; !ok || gn != groupName {
+					continue
+				}
+
+				// Don't count the scheduling pod itself.
+				if candidatePod.UID == pod.UID {
+					continue
+				}
+
+				podGroup = append(podGroup, candidatePod)
+			}
+		}
+	}
+
 	for _, thr := range throttles {
 		affected = append(affected, *thr)
 		reservedAmt, reservedPodNNs := c.cache.reservedResourceAmount(types.NamespacedName{Namespace: thr.Namespace, Name: thr.Name})
+
+		usedInGroup := schedulev1alpha1.ResourceAmount{}
+		for _, groupPod := range podGroup {
+			ns, err := c.namespaceInformer.Lister().Get(groupPod.Namespace)
+			if err != nil {
+				return nil, nil, nil, nil, nil, err
+			}
+
+			// If a pod of a group is already counted, skip it because it'll be counted as a reserved resrouce amount.
+			thrnn := types.NamespacedName{Namespace: thr.Namespace, Name: thr.Name}
+			if c.cache.exist(thrnn, groupPod) {
+				continue
+			}
+
+			match, err := thr.Spec.Selector.MatchesToPod(groupPod, ns)
+			if err != nil {
+				return nil, nil, nil, nil, nil, err
+			}
+			if !match {
+				continue
+			}
+
+			usedInGroup = usedInGroup.Add(schedulev1alpha1.ResourceAmountOfPod(groupPod))
+		}
+
 		checkStatus := thr.CheckThrottledFor(
 			pod,
 			reservedAmt,
+			usedInGroup,
 			isThrottledOnEqual,
 		)
 		klog.V(3).InfoS("CheckThrottled result",
@@ -408,6 +467,7 @@ func (c *ClusterThrottleController) CheckThrottled(
 			"Threashold", thr.Status.CalculatedThreshold.Threshold,
 			"RequestedByPod", schedulev1alpha1.ResourceAmountOfPod(pod),
 			"UsedInClusterThrottle", thr.Status.Used,
+			"UsedInPodGroup", usedInGroup,
 			"ReservedAmountInScheduler", reservedAmt,
 			"ReservedPodsInScheduler", strings.Join(sets.List(reservedPodNNs), ","),
 			"AmountForCheck", schedulev1alpha1.ResourceAmount{}.Add(thr.Status.Used).Add(schedulev1alpha1.ResourceAmountOfPod(pod)).Add(reservedAmt),
@@ -417,11 +477,13 @@ func (c *ClusterThrottleController) CheckThrottled(
 			alreadyThrottled = append(alreadyThrottled, *thr)
 		case schedulev1alpha1.CheckThrottleStatusInsufficient:
 			insufficient = append(insufficient, *thr)
+		case schedulev1alpha1.CheckThrottleStatusInsufficientIncludingPodGroup:
+			insufficientIncludingGroup = append(insufficientIncludingGroup, *thr)
 		case schedulev1alpha1.CheckThrottleStatusPodRequestsExceedsThreshold:
 			podRequestsExceedsThreshold = append(podRequestsExceedsThreshold, *thr)
 		}
 	}
-	return alreadyThrottled, insufficient, podRequestsExceedsThreshold, affected, nil
+	return alreadyThrottled, insufficient, insufficientIncludingGroup, podRequestsExceedsThreshold, affected, nil
 }
 
 // mustSetupEventHandler sets up event handlers. If something wrong happens, it will panic.

pkg/controllers/reserved_resource_amounts.go

@@ -63,6 +63,16 @@ func (c *reservedResourceAmounts) getPodResourceAmountMap(nn types.NamespacedNam
 	return c.cache[nn]
 }
 
+func (c *reservedResourceAmounts) exist(nn types.NamespacedName, pod *corev1.Pod) bool {
+	c.keyMutex.LockKey(nn.String())
+	defer func() {
+		_ = c.keyMutex.UnlockKey(nn.String())
+	}()
+
+	m := c.getPodResourceAmountMap(nn)
+	return m.exist(types.NamespacedName{Namespace: pod.Namespace, Name: pod.Name})
+}
+
 func (c *reservedResourceAmounts) addPod(nn types.NamespacedName, pod *corev1.Pod) bool {
 	c.keyMutex.LockKey(nn.String())
 	defer func() {
@@ -145,6 +155,11 @@ func (c podResourceAmountMap) removeByNN(nn types.NamespacedName) bool {
 	return ok
 }
 
+func (c podResourceAmountMap) exist(nn types.NamespacedName) bool {
+	_, ok := c[nn]
+	return ok
+}
+
 func (c podResourceAmountMap) totalResoruceAmount() (schedulev1alpha1.ResourceAmount, sets.Set[string]) {
 	result := schedulev1alpha1.ResourceAmount{}
 	nns := sets.New[string]()

pkg/scheduler_plugin/plugin.go

@@ -46,9 +46,10 @@ const (
 )
 
 type KubeThrottler struct {
-	fh                 framework.Handle
-	throttleCtr        *controllers.ThrottleController
-	clusterThrottleCtr *controllers.ClusterThrottleController
+	fh                  framework.Handle
+	throttleCtr         *controllers.ThrottleController
+	clusterThrottleCtr  *controllers.ClusterThrottleController
+	groupNameAnnotation string
 }
 
 var _ framework.PreFilterPlugin = &KubeThrottler{}
@@ -135,9 +136,10 @@ func NewPlugin(ctx context.Context, configuration runtime.Object, fh framework.H
 	}
 
 	pl := KubeThrottler{
-		fh:                 fh,
-		throttleCtr:        throttleController,
-		clusterThrottleCtr: clusterthrottleController,
+		fh:                  fh,
+		throttleCtr:         throttleController,
+		clusterThrottleCtr:  clusterthrottleController,
+		groupNameAnnotation: kubeThrottlerArgs.GroupNameAnnotation,
 	}
 
 	return &pl, nil
@@ -160,20 +162,21 @@ func (pl *KubeThrottler) PreFilter(
 		"#AffectedThrottles", len(thrAffected),
 	)
 
-	clthrActive, clthrInsufficient, clthrPodRequestsExceeds, clThrAffected, err := pl.clusterThrottleCtr.CheckThrottled(pod, false)
+	clthrActive, clthrInsufficient, clthrInsufficientGroup, clthrPodRequestsExceeds, clThrAffected, err := pl.clusterThrottleCtr.CheckThrottled(pod, false, pl.groupNameAnnotation)
 	if err != nil {
 		return nil, framework.NewStatus(framework.Error, err.Error())
 	}
 	klog.V(2).InfoS("PreFilter: clusterthrottle check result",
 		"Pod", pod.Namespace+"/"+pod.Name,
 		"#ActiveClusterThrottles", len(clthrActive),
 		"#InsufficientClusterThrottles", len(clthrInsufficient),
+		"#InsufficientClusterThrottlesIncludingGroup", len(clthrInsufficientGroup),
 		"#PodRequestsExceedsThresholdClusterThrottles", len(clthrPodRequestsExceeds),
 		"#AffectedClusterThrottles", len(clThrAffected),
 	)
 
 	if len(thrActive)+len(thrInsufficient)+len(thrPodRequestsExceeds)+
-		len(clthrActive)+len(clthrInsufficient)+len(clthrPodRequestsExceeds) == 0 {
+		len(clthrActive)+len(clthrInsufficient)+len(clthrInsufficientGroup)+len(clthrPodRequestsExceeds) == 0 {
 		return nil, framework.NewStatus(framework.Success)
 	}
 
@@ -206,6 +209,9 @@ func (pl *KubeThrottler) PreFilter(
 	if len(clthrInsufficient) != 0 {
 		reasons = append(reasons, fmt.Sprintf("clusterthrottle[%s]=%s", schedulev1alpha1.CheckThrottleStatusInsufficient, strings.Join(clusterThrottleNames(clthrInsufficient), ",")))
 	}
+	if len(clthrInsufficientGroup) != 0 {
+		reasons = append(reasons, fmt.Sprintf("clusterthrottle[%s]=%s", schedulev1alpha1.CheckThrottleStatusInsufficientIncludingPodGroup, strings.Join(clusterThrottleNames(clthrInsufficientGroup), ",")))
+	}
 	if len(thrInsufficient) != 0 {
 		reasons = append(reasons, fmt.Sprintf("throttle[%s]=%s", schedulev1alpha1.CheckThrottleStatusInsufficient, strings.Join(throttleNames(thrInsufficient), ",")))
 	}

pkg/scheduler_plugin/plugin_args.go

@@ -37,6 +37,7 @@ type KubeThrottlerPluginArgs struct {
 	TargetSchedulerName                 string        `json:"targetSchedulerName"`
 	ControllerThrediness                int           `json:"controllerThrediness"`
 	NumKeyMutex                         int           `json:"numKeyMutex"`
+	GroupNameAnnotation                 string        `json:"groupNameAnnotation"`
 }
 
 func DecodePluginArgs(configuration runtime.Object) (*KubeThrottlerPluginArgs, error) {

test/integration/clusterthrottle_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/everpeace/kube-throttler/pkg/apis/schedule/v1alpha1"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
+	"github.com/onsi/gomega/types"
 	corev1 "k8s.io/api/core/v1"
 )
 
@@ -161,8 +162,49 @@ var _ = Describe("Clusterthrottle Test", func() {
 				Consistently(PodIsNotScheduled(ctx, DefaultNs, pod.Name)).Should(Succeed())
 			})
 		})
+		Context("Pod Group", func() {
+			var (
+				podPassedGroup    []*corev1.Pod
+				podThrottledGroup []*corev1.Pod
+			)
+			BeforeEach(func() {
+				for i := 0; i < 2; i++ {
+					podPassedGroup = append(podPassedGroup, MustCreatePod(ctx, MakePod(DefaultNs, fmt.Sprintf("passed-pod%d", i), "50m").Annotation(groupNameAnnotation, "passed").Label(throttleKey, throttleName).Obj()))
+				}
+				for i := 0; i < 3; i++ {
+					podThrottledGroup = append(podThrottledGroup, MustCreatePod(ctx, MakePod(DefaultNs, fmt.Sprintf("throttled-pod%d", i), "50m").Annotation(groupNameAnnotation, "throttled").Label(throttleKey, throttleName).Obj()))
+				}
+			})
+			It("should not schedule pod3", func() {
+				Eventually(AsyncAll(
+					WakeupBackoffPod(ctx),
+					ClusterThottleHasStatus(
+						ctx, thr.Name,
+						ClthrOpts.WithCalculatedThreshold(thr.Spec.Threshold),
+						ClthrOpts.WithPodThrottled(true),
+						ClthrOpts.WithUsedPod(len(podPassedGroup)),
+						ClthrOpts.WithUsedCpuReq(fmt.Sprintf("%dm", len(podPassedGroup)*50)),
+						ClthrOpts.WithCpuThrottled(false),
+					),
+					func(g types.Gomega) {
+						for _, pod := range podPassedGroup {
+							PodIsScheduled(ctx, DefaultNs, pod.Name)
+						}
+					},
+					func(g types.Gomega) {
+						for _, pod := range podThrottledGroup {
+							MustPodFailedScheduling(ctx, DefaultNs, pod.Name, v1alpha1.CheckThrottleStatusInsufficientIncludingPodGroup)
+						}
+					},
+				)).Should(Succeed())
+				Consistently(func(g types.Gomega) {
+					for _, pod := range podPassedGroup {
+						PodIsScheduled(ctx, DefaultNs, pod.Name)
+					}
+				}).Should(Succeed())
+			})
+		})
 	})
-
 	When("Many pods are created at once", func() {
 		var thr *v1alpha1.ClusterThrottle
 		var scheduled = make([]*corev1.Pod, 20)

test/integration/integration_suite_test.go

@@ -40,9 +40,10 @@ import (
 )
 
 const (
-	SchedulerName = "kube-throttler-integration"
-	ThrottlerName = "kube-throttler"
-	DefaultNs     = "default"
+	SchedulerName       = "kube-throttler-integration"
+	ThrottlerName       = "kube-throttler"
+	DefaultNs           = "default"
+	groupNameAnnotation = "scheduling.k8s.pfn.io/group-name"
 )
 
 var (
@@ -111,12 +112,13 @@ func mustStartKubeThrottler() {
 			      name: %s
 			      targetSchedulerName: %s
 			      kubeconfig: %s
+			      groupNameAnnotation: %s
 			      controllerThrediness: 64
 			      numKeyMutex: 128
 		`,
-			kubeConfigPath,                               // clientConnection.kubeconfig
-			SchedulerName,                                // prifiles[0].scedulerName
-			ThrottlerName, SchedulerName, kubeConfigPath, // profiles[0].pluginConfig[0].args
+			kubeConfigPath,                                                    // clientConnection.kubeconfig
+			SchedulerName,                                                     // prifiles[0].scedulerName
+			ThrottlerName, SchedulerName, kubeConfigPath, groupNameAnnotation, // profiles[0].pluginConfig[0].args
 		),
 	))
 	Expect(err).NotTo(HaveOccurred())