resource_manager: implement the token-consuming by client

pkg/mcs/resource_manager/server/token_buckets_test.go

@@ -81,14 +81,16 @@ func TestGroupTokenBucketRequest(t *testing.T) {
 	tb, trickle = gtb.request(time1, 60000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
 	re.LessOrEqual(math.Abs(tb.Tokens-22000), 1e-7)
 	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
-	tb, trickle = gtb.request(time1, 3000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
-	re.LessOrEqual(math.Abs(tb.Tokens-2000), 1e-7)
-	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	// Get reserved 10000 tokens = fillrate(2000) * 10 * defaultReserveRatio(0.5)
+	// Max loan tokens is 60000.
 	tb, trickle = gtb.request(time1, 3000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
 	re.LessOrEqual(math.Abs(tb.Tokens-1000), 1e-7)
 	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
-	time2 := time1.Add(10 * time.Second)
+	tb, trickle = gtb.request(time1, 12000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
+	re.LessOrEqual(math.Abs(tb.Tokens-10000), 1e-7)
+	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	time2 := time1.Add(30 * time.Second)
 	tb, trickle = gtb.request(time2, 20000, uint64(time.Second)*10/uint64(time.Millisecond), clientUniqueID)
-	re.LessOrEqual(math.Abs(tb.Tokens-19000), 1e-7)
+	re.LessOrEqual(math.Abs(tb.Tokens-20000), 1e-7)
 	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
 }

pkg/mcs/resource_manager/server/token_bukets.go

@@ -401,9 +401,10 @@ func (ts *TokenSlot) assignSlotTokens(neededTokens float64, targetPeriodMs uint6
 			}
 		}
 	}
-	if grantedTokens < defaultReserveRatio*float64(fillRate)*targetPeriodTimeSec {
-		ts.tokenCapacity -= defaultReserveRatio*float64(fillRate)*targetPeriodTimeSec - grantedTokens
-		grantedTokens = defaultReserveRatio * float64(fillRate) * targetPeriodTimeSec
+	if neededTokens > 0 && grantedTokens < defaultReserveRatio*float64(fillRate)*targetPeriodTimeSec {
+		reservedTokens := math.Min(neededTokens, defaultReserveRatio*float64(fillRate)*targetPeriodTimeSec)
+		ts.tokenCapacity -= reservedTokens - grantedTokens
+		grantedTokens = reservedTokens
 	}
 	res.Tokens = grantedTokens