resourcemanager: fix client allocation under service limits

pkg/mcs/resourcemanager/server/token_buckets.go

@@ -309,7 +309,8 @@ func (gtb *GroupTokenBucket) balanceSlotTokens(
 
 	var (
 		totalFillRate, totalBurstLimit = gtb.getFillRateAndBurstLimit()
-		basicFillRate                  = totalFillRate * evenRatio
+		allocationBudget               = totalFillRate
+		basicFillRate                  = allocationBudget * evenRatio
 		allocatedFillRate              = 0.0
 		allocationMap                  = make(map[uint64]float64, len(gtb.tokenSlots))
 		extraDemandSlots               = make(map[uint64]float64, len(gtb.tokenSlots))
@@ -335,6 +336,19 @@ func (gtb *GroupTokenBucket) balanceSlotTokens(
 		}
 		return
 	}
+	// A negative configured burst limit allows bursting, but even the default group
+	// can be constrained by the keyspace Service Limit. A positive override burst
+	// limit is the capacity assigned to this group by Service Limit coordination.
+	// Use that result, capped by the effective fill rate, as the client allocation
+	// budget. Using a huge fill rate (e.g. the default group's UnlimitedRate) instead
+	// makes client demand negligible relative to the budget, so distributing the
+	// unused budget evenly produces nearly equal shares despite unequal demand.
+	// This only changes how shares are calculated; the group refill rate and loan
+	// algorithm remain unchanged.
+	if gtb.overrideBurstLimit > 0 && gtb.getBurstLimitSetting() < 0 {
+		allocationBudget = math.Min(allocationBudget, float64(gtb.overrideBurstLimit))
+		basicFillRate = allocationBudget * evenRatio
+	}
 	if gtb.grt == nil {
 		gtb.grt = newGroupRUTracker()
 	}
@@ -357,7 +371,7 @@ func (gtb *GroupTokenBucket) balanceSlotTokens(
 		allocationMap[clientUniqueID] = allocation
 		allocatedFillRate += allocation
 	}
-	remainingFillRate := totalFillRate - allocatedFillRate
+	remainingFillRate := allocationBudget - allocatedFillRate
 	// For the remaining fill rate, allocate it proportionally to the high demand slots.
 	if remainingFillRate > 0 && len(extraDemandSlots) > 0 {
 		for clientUniqueID, extraDemand := range extraDemandSlots {
@@ -375,7 +389,7 @@ func (gtb *GroupTokenBucket) balanceSlotTokens(
 		// Distribute the fill rate.
 		fillRate := allocationMap[clientUniqueID]
 		// Distribute the burst limit and assign tokens based on the allocation ratio.
-		ratio := fillRate / totalFillRate
+		ratio := fillRate / allocationBudget
 		burstLimit := float64(totalBurstLimit) * ratio
 		assignTokens := tokensForBalance * ratio
 		// Need to reserve burst limit to next balance.
@@ -390,6 +404,9 @@ func (gtb *GroupTokenBucket) balanceSlotTokens(
 		slot.lastTokenCapacity += assignTokens
 		// Update the slot fill rate and burst limit.
 		slot.fillRate = fillRate
+		if allocationBudget != totalFillRate {
+			slot.fillRate = totalFillRate * ratio
+		}
 		slot.burstLimit = int64(burstLimit)
 	}
 }

pkg/mcs/resourcemanager/server/token_buckets_test.go

@@ -504,3 +504,72 @@ func TestFractionalAllocationWithinSmallGroupBudget(t *testing.T) {
 	}
 	re.InDelta(1.0, sum, 1e-12)
 }
+
+func TestServiceLimitedClientAllocation(t *testing.T) {
+	for _, tc := range []struct {
+		name          string
+		fillRate      uint64
+		burstLimit    int64
+		overrideFill  float64
+		overrideBurst int64
+		wantHot       int64
+		wantCold      int64
+	}{
+		{"unlimited group", UnlimitedRate, -1, -1, 160000, 92500, 67500},
+		{"lower service budget", UnlimitedRate, -1, -1, 120000, 72500, 47500},
+		{"hot demand above equal share", UnlimitedRate, -1, -1, 80000, 55000, 25000},
+		{"both demands above equal share", UnlimitedRate, -1, -1, 40000, 20000, 20000},
+		{"moderated group", UnlimitedRate, -2, -1, 160000, 92500, 67500},
+		{"finite refill below budget", 100000, -1, -1, 160000, 100000, 60000},
+		{"overridden refill", UnlimitedRate, -1, 100000, 100000, 62500, 37500},
+		{"fractional overridden refill", UnlimitedRate, -1, 100000.5, 100000, 62500, 37500},
+		{"explicit burst", UnlimitedRate, 160000, -1, 120000, 60000, 59999},
+		{"rate controlled", UnlimitedRate, 0, -1, 120000, 60000, 59999},
+		{"service disabled", UnlimitedRate, -1, -1, -1, -1, -1},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			re := require.New(t)
+			gtb := NewGroupTokenBucket(testResourceGroupName, &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{FillRate: tc.fillRate, BurstLimit: tc.burstLimit},
+			})
+			gtb.overrideFillRate = tc.overrideFill
+			gtb.overrideBurstLimit = tc.overrideBurst
+			gtb.grt = newGroupRUTracker()
+			now := time.Now()
+			for i, demand := range []float64{50000, 25000} {
+				id := uint64(i + 1)
+				gtb.tokenSlots[id] = newTokenSlot(id, now)
+				rt := gtb.grt.getOrCreateRUTracker(id)
+				rt.initialized = true
+				rt.lastSampleTime = now
+				rt.lastEMA = demand
+			}
+			fillRate := gtb.getFillRate()
+			const tokensForBalance = 10000.0
+			gtb.balanceSlotTokens(now, 1, 1, tokensForBalance)
+			re.Equal(fillRate, gtb.getFillRate())
+			re.Equal(tc.fillRate, gtb.Settings.FillRate)
+			re.Equal(tc.burstLimit, gtb.Settings.BurstLimit)
+			re.Equal(tc.wantHot, gtb.tokenSlots[1].burstLimit)
+			re.Equal(tc.wantCold, gtb.tokenSlots[2].burstLimit)
+			re.InDelta(fillRate, gtb.tokenSlots[1].fillRate+gtb.tokenSlots[2].fillRate, 1e-6)
+			if tc.overrideBurst > 0 {
+				for _, slot := range gtb.tokenSlots {
+					// Refill and newly assigned tokens must use the same share as capacity.
+					ratio := float64(slot.burstLimit) / float64(tc.overrideBurst)
+					re.InDelta(ratio, slot.fillRate/fillRate, 1/float64(tc.overrideBurst))
+					re.InDelta(ratio, slot.curTokenCapacity/tokensForBalance, 1/float64(tc.overrideBurst))
+					re.Equal(slot.curTokenCapacity, slot.lastTokenCapacity)
+				}
+				re.InDelta(tokensForBalance, gtb.tokenSlots[1].curTokenCapacity+gtb.tokenSlots[2].curTokenCapacity, 1e-7)
+				// Removing the other client restores the whole group allocation.
+				gtb.Tokens = tokensForBalance
+				gtb.balanceSlotTokens(now, 2, 0, 0)
+				re.Len(gtb.tokenSlots, 1)
+				re.Equal(fillRate, gtb.tokenSlots[1].fillRate)
+				re.Equal(tc.overrideBurst, gtb.tokenSlots[1].burstLimit)
+				re.Equal(tokensForBalance, gtb.tokenSlots[1].curTokenCapacity)
+			}
+		})
+	}
+}