fix(resourcemanager): add validation for token assignment anomaly (#9597)

ref #9455

Add safeguard checks and warnings during token assignment to handle potential invalid token values.

Signed-off-by: JmPotato <[email protected]>

Author: JmPotato

pkg/mcs/resourcemanager/server/resource_group.go

@@ -69,9 +69,9 @@ func (rus *RequestUnitSettings) Clone() *RequestUnitSettings {
 }
 
 // NewRequestUnitSettings creates a new RequestUnitSettings with the given token bucket.
-func NewRequestUnitSettings(tokenBucket *rmpb.TokenBucket) *RequestUnitSettings {
+func NewRequestUnitSettings(resourceGroupName string, tokenBucket *rmpb.TokenBucket) *RequestUnitSettings {
 	return &RequestUnitSettings{
-		RU: NewGroupTokenBucket(tokenBucket),
+		RU: NewGroupTokenBucket(resourceGroupName, tokenBucket),
 	}
 }
 
@@ -241,9 +241,9 @@ func FromProtoResourceGroup(group *rmpb.ResourceGroup) *ResourceGroup {
 	switch group.GetMode() {
 	case rmpb.GroupMode_RUMode:
 		if group.GetRUSettings() == nil {
-			rg.RUSettings = NewRequestUnitSettings(nil)
+			rg.RUSettings = NewRequestUnitSettings(rg.Name, nil)
 		} else {
-			rg.RUSettings = NewRequestUnitSettings(group.GetRUSettings().GetRU())
+			rg.RUSettings = NewRequestUnitSettings(rg.Name, group.GetRUSettings().GetRU())
 		}
 		if group.RUStats != nil {
 			rg.RUConsumption = group.RUStats
@@ -268,6 +268,9 @@ func (rg *ResourceGroup) RequestRU(
 	}
 	// First, try to get tokens from the resource group.
 	tb, trickleTimeMs := rg.RUSettings.RU.request(now, requiredToken, targetPeriodMs, clientUniqueID)
+	if tb == nil {
+		return nil
+	}
 	// Then, try to apply the service limit.
 	grantedTokens := tb.GetTokens()
 	limitedTokens := sl.applyServiceLimit(now, grantedTokens)

pkg/mcs/resourcemanager/server/resource_group_test.go

@@ -26,7 +26,7 @@ import (
 
 func TestPatchResourceGroup(t *testing.T) {
 	re := require.New(t)
-	rg := &ResourceGroup{Name: "test", Mode: rmpb.GroupMode_RUMode, RUSettings: NewRequestUnitSettings(nil)}
+	rg := &ResourceGroup{Name: testResourceGroupName, Mode: rmpb.GroupMode_RUMode, RUSettings: NewRequestUnitSettings(testResourceGroupName, nil)}
 	testCaseRU := []struct {
 		patchJSONString  string
 		expectJSONString string

pkg/mcs/resourcemanager/server/token_buckets.go

@@ -148,9 +148,24 @@ type tokenSlot struct {
 	lastReqTime       time.Time
 }
 
+func (ts *tokenSlot) logFields() []zap.Field {
+	return []zap.Field{
+		zap.Uint64("slot-fill-rate", ts.fillRate),
+		zap.Int64("slot-burst-limit", ts.burstLimit),
+		zap.Float64("slot-require-tokens-sum", ts.requireTokensSum),
+		zap.Float64("slot-token-capacity", ts.tokenCapacity),
+		zap.Float64("slot-last-token-capacity", ts.lastTokenCapacity),
+		zap.Time("slot-last-req-time", ts.lastReqTime),
+	}
+}
+
 // GroupTokenBucketState is the running state of TokenBucket.
 type GroupTokenBucketState struct {
-	Tokens float64 `json:"tokens,omitempty"`
+	Tokens      float64    `json:"tokens,omitempty"`
+	LastUpdate  *time.Time `json:"last_update,omitempty"`
+	Initialized bool       `json:"initialized"`
+
+	resourceGroupName string
 	// ClientUniqueID -> TokenSlot
 	tokenSlots                 map[uint64]*tokenSlot
 	clientConsumptionTokensSum float64
@@ -173,8 +188,6 @@ type GroupTokenBucketState struct {
 	// means the burst limit is overridden.
 	overrideBurstLimit int64
 
-	LastUpdate  *time.Time `json:"last_update,omitempty"`
-	Initialized bool       `json:"initialized"`
 	// settingChanged is used to avoid that the number of tokens returned is jitter because of changing fill rate.
 	settingChanged      bool
 	lastCheckExpireSlot time.Time
@@ -198,6 +211,7 @@ func (gts *GroupTokenBucketState) clone() *GroupTokenBucketState {
 		Tokens:                     gts.Tokens,
 		LastUpdate:                 lastUpdate,
 		Initialized:                gts.Initialized,
+		resourceGroupName:          gts.resourceGroupName,
 		tokenSlots:                 tokenSlots,
 		overrideFillRate:           gts.overrideFillRate,
 		overrideBurstLimit:         gts.overrideBurstLimit,
@@ -210,16 +224,11 @@ func (gts *GroupTokenBucketState) resetLoan() {
 	gts.settingChanged = false
 	gts.Tokens = 0
 	gts.clientConsumptionTokensSum = 0
-	evenRatio := 1.0
-	if l := len(gts.tokenSlots); l > 0 {
-		evenRatio = 1 / float64(l)
-	}
-
-	evenTokens := gts.Tokens * evenRatio
+	// Reset all slots.
 	for _, slot := range gts.tokenSlots {
 		slot.requireTokensSum = 0
-		slot.tokenCapacity = evenTokens
-		slot.lastTokenCapacity = evenTokens
+		slot.tokenCapacity = 0
+		slot.lastTokenCapacity = 0
 	}
 }
 
@@ -330,14 +339,15 @@ func (gtb *GroupTokenBucket) calcRateAndBurstLimit(ratio float64) (fillRate uint
 }
 
 // NewGroupTokenBucket returns a new GroupTokenBucket
-func NewGroupTokenBucket(tokenBucket *rmpb.TokenBucket) *GroupTokenBucket {
+func NewGroupTokenBucket(resourceGroupName string, tokenBucket *rmpb.TokenBucket) *GroupTokenBucket {
 	if tokenBucket == nil || tokenBucket.Settings == nil {
 		return &GroupTokenBucket{}
 	}
 	return &GroupTokenBucket{
 		Settings: tokenBucket.GetSettings(),
 		GroupTokenBucketState: GroupTokenBucketState{
 			Tokens:             tokenBucket.GetTokens(),
+			resourceGroupName:  resourceGroupName,
 			tokenSlots:         make(map[uint64]*tokenSlot),
 			overrideFillRate:   -1,
 			overrideBurstLimit: -1,
@@ -418,43 +428,93 @@ func (gtb *GroupTokenBucket) updateTokens(now time.Time, burstLimit int64, clien
 	gtb.balanceSlotTokens(clientUniqueID, requiredToken, elapseTokens)
 }
 
+func (gtb *GroupTokenBucket) inspectAnomalies(
+	tb *rmpb.TokenBucket,
+	slot *tokenSlot,
+	logFields []zap.Field,
+) bool {
+	var errMsg string
+	// Verify whether the allocated token is invalid, such as negative values, math.Inf, or math.NaN.
+	if tb.Tokens <= 0 || math.IsInf(tb.Tokens, 0) || math.IsNaN(tb.Tokens) {
+		errMsg = "assigned token is invalid"
+	}
+	// Verify whether the state of the slot is abnormal.
+	if math.IsInf(slot.tokenCapacity, 0) || math.IsNaN(slot.tokenCapacity) {
+		errMsg = "slot token capacity is invalid"
+	}
+	// If there is any error, reset the group token bucket to avoid the group token bucket is in a bad state.
+	isAnomaly := len(errMsg) > 0
+	if isAnomaly {
+		logFields = append(logFields,
+			append(
+				slot.logFields(),
+				zap.String("resource-group-name", gtb.resourceGroupName),
+				zap.String("settings", gtb.Settings.String()),
+				zap.Float64("tokens", gtb.Tokens),
+				zap.Float64("client-consumption-tokens-sum", gtb.clientConsumptionTokensSum),
+				zap.Int("slot-len", len(gtb.tokenSlots)),
+			)...,
+		)
+		log.Error(errMsg, logFields...)
+		// Reset after logging to keep the original context.
+		gtb.resetLoan()
+	}
+	return isAnomaly
+}
+
 // request requests tokens from the corresponding slot.
-func (gtb *GroupTokenBucket) request(now time.Time,
+func (gtb *GroupTokenBucket) request(
+	now time.Time,
 	requiredToken float64,
 	targetPeriodMs, clientUniqueID uint64,
 ) (*rmpb.TokenBucket, int64) {
 	burstLimit := gtb.getBurstLimit()
 	gtb.updateTokens(now, burstLimit, clientUniqueID, requiredToken)
 	slot, ok := gtb.tokenSlots[clientUniqueID]
 	if !ok {
-		return &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{BurstLimit: burstLimit}}, 0
+		return &rmpb.TokenBucket{
+			Settings: &rmpb.TokenLimitSettings{BurstLimit: burstLimit},
+			Tokens:   0.0,
+		}, 0
 	}
 	res, trickleDuration := slot.assignSlotTokens(requiredToken, targetPeriodMs)
+	// Inspect the group token bucket and the assigned token result to catch any anomalies.
+	if isAnomaly := gtb.inspectAnomalies(res, slot, []zap.Field{
+		zap.Time("now", now),
+		zap.Uint64("client-unique-id", clientUniqueID),
+		zap.Uint64("target-period-ms", targetPeriodMs),
+		zap.Float64("required-token", requiredToken),
+		zap.Float64("assigned-tokens", res.Tokens),
+	}); isAnomaly {
+		// Return nil here to prevent sending any unexpected result to the client.
+		// The client has to retry later to access the resource group whose state has been reset.
+		return nil, 0
+	}
 	// Update bucket to record all tokens.
 	gtb.Tokens -= slot.lastTokenCapacity - slot.tokenCapacity
 	slot.lastTokenCapacity = slot.tokenCapacity
-
 	return res, trickleDuration
 }
 
 func (ts *tokenSlot) assignSlotTokens(requiredToken float64, targetPeriodMs uint64) (*rmpb.TokenBucket, int64) {
-	var res rmpb.TokenBucket
-	burstLimit := ts.burstLimit
-	res.Settings = &rmpb.TokenLimitSettings{BurstLimit: burstLimit}
+	res := &rmpb.TokenBucket{
+		Settings: &rmpb.TokenLimitSettings{BurstLimit: ts.burstLimit},
+		Tokens:   0.0,
+	}
 	if getBurstableMode(res.Settings) == unlimited {
 		res.Tokens = requiredToken
-		return &res, 0
+		return res, 0
 	}
 	// FillRate is used for the token server unavailable in abnormal situation.
 	if requiredToken <= 0 {
-		return &res, 0
+		return res, 0
 	}
 	// If the current tokens can directly meet the requirement, returns the need token.
 	if ts.tokenCapacity >= requiredToken {
 		ts.tokenCapacity -= requiredToken
 		// granted the total request tokens
 		res.Tokens = requiredToken
-		return &res, 0
+		return res, 0
 	}
 
 	// Firstly allocate the remaining tokens
@@ -472,6 +532,7 @@ func (ts *tokenSlot) assignSlotTokens(requiredToken float64, targetPeriodMs uint
 		targetPeriodTimeSec = targetPeriodTime.Seconds()
 		trickleTime         = 0.
 		fillRate            = ts.fillRate
+		burstLimit          = ts.burstLimit
 	)
 
 	loanCoefficient := defaultLoanCoefficient
@@ -547,5 +608,5 @@ func (ts *tokenSlot) assignSlotTokens(requiredToken float64, targetPeriodMs uint
 	} else {
 		trickleDuration = targetPeriodTime
 	}
-	return &res, trickleDuration.Milliseconds()
+	return res, trickleDuration.Milliseconds()
 }

pkg/mcs/resourcemanager/server/token_buckets_test.go

@@ -24,6 +24,8 @@ import (
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 )
 
+const testResourceGroupName = "test"
+
 func TestGroupTokenBucketUpdateAndPatch(t *testing.T) {
 	re := require.New(t)
 	tbSetting := &rmpb.TokenBucket{
@@ -35,7 +37,7 @@ func TestGroupTokenBucketUpdateAndPatch(t *testing.T) {
 	}
 
 	clientUniqueID := uint64(0)
-	tb := NewGroupTokenBucket(tbSetting)
+	tb := NewGroupTokenBucket(testResourceGroupName, tbSetting)
 	time1 := time.Now()
 	tb.request(time1, 0, 0, clientUniqueID)
 	re.LessOrEqual(math.Abs(tbSetting.Tokens-tb.Tokens), 1e-7)
@@ -62,7 +64,7 @@ func TestGroupTokenBucketUpdateAndPatch(t *testing.T) {
 			BurstLimit: -1,
 		},
 	}
-	tb = NewGroupTokenBucket(tbSetting)
+	tb = NewGroupTokenBucket(testResourceGroupName, tbSetting)
 	tb.request(time2, 0, 0, clientUniqueID)
 	re.LessOrEqual(math.Abs(tbSetting.Tokens), 1e-7)
 	time3 := time.Now()
@@ -76,7 +78,7 @@ func TestGroupTokenBucketUpdateAndPatch(t *testing.T) {
 			BurstLimit: -1,
 		},
 	}
-	tb = NewGroupTokenBucket(tbSetting)
+	tb = NewGroupTokenBucket(testResourceGroupName, tbSetting)
 	tb.request(time3, 0, 0, clientUniqueID)
 	re.LessOrEqual(math.Abs(tbSetting.Tokens-200000), 1e-7)
 	time.Sleep(10 * time.Millisecond)
@@ -95,7 +97,7 @@ func TestGroupTokenBucketRequest(t *testing.T) {
 		},
 	}
 
-	gtb := NewGroupTokenBucket(tbSetting)
+	gtb := NewGroupTokenBucket(testResourceGroupName, tbSetting)
 	time1 := time.Now()
 	clientUniqueID := uint64(0)
 	tb, trickle := gtb.request(time1, 190000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
@@ -128,7 +130,7 @@ func TestGroupTokenBucketRequest(t *testing.T) {
 func TestGroupTokenBucketRequestBurstLimit(t *testing.T) {
 	re := require.New(t)
 	testGroupSetting := func(tbSetting *rmpb.TokenBucket, expectedFillRate, expectedBurstLimit int64) {
-		gtb := NewGroupTokenBucket(tbSetting)
+		gtb := NewGroupTokenBucket(testResourceGroupName, tbSetting)
 		time1 := time.Now()
 		clientUniqueID := uint64(0)
 		gtb.request(time1, 190000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
@@ -219,7 +221,7 @@ func TestGroupTokenBucketRequestLoop(t *testing.T) {
 		},
 	}
 
-	gtb := NewGroupTokenBucket(tbSetting)
+	gtb := NewGroupTokenBucket(testResourceGroupName, tbSetting)
 	clientUniqueID := uint64(0)
 	initialTime := time.Now()