diff --git a/pkg/mcs/resourcemanager/server/token_buckets.go b/pkg/mcs/resourcemanager/server/token_buckets.go index 7e77fc6dfe..8cd958f2af 100644 --- a/pkg/mcs/resourcemanager/server/token_buckets.go +++ b/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) } } diff --git a/pkg/mcs/resourcemanager/server/token_buckets_test.go b/pkg/mcs/resourcemanager/server/token_buckets_test.go index d22085e32b..e855cbb7dc 100644 --- a/pkg/mcs/resourcemanager/server/token_buckets_test.go +++ b/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) + } + }) + } +}