resource_manager: implement token assignment in server

pkg/mcs/resource_manager/server/grpc_service.go

@@ -16,12 +16,16 @@ package server
 
 import (
 	"context"
+	"io"
 	"net/http"
+	"time"
 
 	"github.com/pingcap/errors"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/mcs/registry"
 	"github.com/tikv/pd/server"
+	"go.uber.org/zap"
 	"google.golang.org/grpc"
 )
 
@@ -125,5 +129,52 @@ func (s *Service) ModifyResourceGroup(ctx context.Context, req *rmpb.PutResource
 
 // AcquireTokenBuckets implements ResourceManagerServer.AcquireTokenBuckets.
 func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBucketsServer) error {
-	return errors.New("Not implemented")
+	for {
+		select {
+		case <-s.ctx.Done():
+			return errors.New("server closed")
+		default:
+		}
+		request, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		targetPeriodMs := request.GetTargetRequestPeriodMs()
+		resps := &rmpb.TokenBucketsResponse{}
+		for _, req := range request.Requests {
+			rg := s.manager.GetResourceGroup(req.ResourceGroupName)
+			if rg == nil {
+				log.Warn("resource group not found", zap.String("resource-group", req.ResourceGroupName))
+				continue
+			}
+			now := time.Now()
+			resp := &rmpb.TokenBucketResponse{
+				ResourceGroupName: rg.Name,
+			}
+			switch rg.Mode {
+			case rmpb.GroupMode_RUMode:
+				for _, re := range req.GetRuItems().GetRequestRU() {
+					switch re.Type {
+					case rmpb.RequestUnitType_RRU:
+						rg.UpdateRRU(now)
+						tokens := rg.RequestRRU(re.Value, targetPeriodMs)
+						resp.GrantedRUTokens = append(resp.GrantedRUTokens, tokens)
+					case rmpb.RequestUnitType_WRU:
+						rg.UpdateWRU(now)
+						tokens := rg.RequestWRU(re.Value, targetPeriodMs)
+						resp.GrantedRUTokens = append(resp.GrantedRUTokens, tokens)
+					}
+				}
+			case rmpb.GroupMode_RawMode:
+				log.Warn("not supports the resource type", zap.String("resource-group", req.ResourceGroupName), zap.String("mode", rmpb.GroupMode_name[int32(rmpb.GroupMode_RawMode)]))
+				continue
+			}
+			log.Debug("finish token request from", zap.String("resource group", req.ResourceGroupName))
+			resps.Responses = append(resps.Responses, resp)
+		}
+		stream.Send(resps)
+	}
 }

pkg/mcs/resource_manager/server/resource_group.go

@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"path"
 	"sync"
+	"time"
 
 	"github.com/pingcap/errors"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
@@ -153,34 +154,64 @@ func FromProtoResourceGroup(group *rmpb.ResourceGroup) *ResourceGroup {
 	case rmpb.GroupMode_RUMode:
 		if settings := group.GetRUSettings(); settings != nil {
 			ruSettings = &RequestUnitSettings{
-				RRU: GroupTokenBucket{
-					TokenBucket: settings.GetRRU(),
-				},
-				WRU: GroupTokenBucket{
-					TokenBucket: settings.GetWRU(),
-				},
+				RRU: NewGroupTokenBucket(settings.GetRRU()),
+				WRU: NewGroupTokenBucket(settings.GetWRU()),
 			}
 			rg.RUSettings = ruSettings
 		}
 	case rmpb.GroupMode_RawMode:
 		if settings := group.GetResourceSettings(); settings != nil {
 			resourceSettings = &NativeResourceSettings{
-				CPU: GroupTokenBucket{
-					TokenBucket: settings.GetCpu(),
-				},
-				IOReadBandwidth: GroupTokenBucket{
-					TokenBucket: settings.GetIoRead(),
-				},
-				IOWriteBandwidth: GroupTokenBucket{
-					TokenBucket: settings.GetIoWrite(),
-				},
+				CPU:              NewGroupTokenBucket(settings.GetCpu()),
+				IOReadBandwidth:  NewGroupTokenBucket(settings.GetIoRead()),
+				IOWriteBandwidth: NewGroupTokenBucket(settings.GetIoWrite()),
 			}
 			rg.ResourceSettings = resourceSettings
 		}
 	}
 	return rg
 }
 
+// UpdateRRU updates the RRU of the resource group.
+func (rg *ResourceGroup) UpdateRRU(now time.Time) {
+	rg.Lock()
+	defer rg.Unlock()
+	if rg.RUSettings != nil {
+		rg.RUSettings.RRU.update(now)
+	}
+}
+
+// UpdateWRU updates the WRU of the resource group.
+func (rg *ResourceGroup) UpdateWRU(now time.Time) {
+	rg.Lock()
+	defer rg.Unlock()
+	if rg.RUSettings != nil {
+		rg.RUSettings.WRU.update(now)
+	}
+}
+
+// RequestRRU requests the RRU of the resource group.
+func (rg *ResourceGroup) RequestRRU(neededTokens float64, targetPeriodMs uint64) *rmpb.GrantedRUTokenBucket {
+	rg.Lock()
+	defer rg.Unlock()
+	if rg.RUSettings == nil {
+		return nil
+	}
+	tb, trickleTimeMs := rg.RUSettings.RRU.request(neededTokens, targetPeriodMs)
+	return &rmpb.GrantedRUTokenBucket{Type: rmpb.RequestUnitType_RRU, GrantedTokens: tb, TrickleTimeMs: trickleTimeMs}
+}
+
+// RequestWRU requests the WRU of the resource group.
+func (rg *ResourceGroup) RequestWRU(neededTokens float64, targetPeriodMs uint64) *rmpb.GrantedRUTokenBucket {
+	rg.Lock()
+	defer rg.Unlock()
+	if rg.RUSettings == nil {
+		return nil
+	}
+	tb, trickleTimeMs := rg.RUSettings.WRU.request(neededTokens, targetPeriodMs)
+	return &rmpb.GrantedRUTokenBucket{Type: rmpb.RequestUnitType_WRU, GrantedTokens: tb, TrickleTimeMs: trickleTimeMs}
+}
+
 // IntoProtoResourceGroup converts a ResourceGroup to a rmpb.ResourceGroup.
 func (rg *ResourceGroup) IntoProtoResourceGroup() *rmpb.ResourceGroup {
 	rg.RLock()

pkg/mcs/resource_manager/server/token_buckets_test.go

@@ -0,0 +1,90 @@
+// Copyright 2022 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,g
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package server
+
+import (
+	"math"
+	"testing"
+	"time"
+
+	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGroupTokenBucketUpdateAndPatch(t *testing.T) {
+	re := require.New(t)
+	tbSetting := &rmpb.TokenBucket{
+		Tokens: 200000,
+		Settings: &rmpb.TokenLimitSettings{
+			FillRate:   2000,
+			BurstLimit: 20000000,
+		},
+	}
+
+	tb := NewGroupTokenBucket(tbSetting)
+	time1 := time.Now()
+	tb.update(time1)
+	re.LessOrEqual(math.Abs(tbSetting.Tokens-tb.Tokens), 1e-7)
+	re.Equal(tbSetting.Settings.FillRate, tb.Settings.FillRate)
+
+	tbSetting = &rmpb.TokenBucket{
+		Tokens: -100000,
+		Settings: &rmpb.TokenLimitSettings{
+			FillRate:   1000,
+			BurstLimit: 10000000,
+		},
+	}
+	tb.patch(tbSetting)
+
+	time2 := time.Now()
+	tb.update(time2)
+	re.LessOrEqual(math.Abs(100000-tb.Tokens), time2.Sub(time1).Seconds()*float64(tbSetting.Settings.FillRate)+1e7)
+	re.Equal(tbSetting.Settings.FillRate, tb.Settings.FillRate)
+}
+
+func TestGroupTokenBucketRequest(t *testing.T) {
+	re := require.New(t)
+	tbSetting := &rmpb.TokenBucket{
+		Tokens: 200000,
+		Settings: &rmpb.TokenLimitSettings{
+			FillRate:   2000,
+			BurstLimit: 20000000,
+		},
+	}
+
+	gtb := NewGroupTokenBucket(tbSetting)
+	time1 := time.Now()
+	gtb.update(time1)
+	tb, trickle := gtb.request(100000, uint64(time.Second)*10/uint64(time.Millisecond))
+	re.LessOrEqual(math.Abs(tb.Tokens-100000), 1e-7)
+	re.Equal(trickle, int64(0))
+	tb, trickle = gtb.request(101000, uint64(time.Second)*10/uint64(time.Millisecond))
+	re.LessOrEqual(math.Abs(tb.Tokens-101000), 1e-7)
+	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	tb, trickle = gtb.request(17000, uint64(time.Second)*10/uint64(time.Millisecond))
+	re.LessOrEqual(math.Abs(tb.Tokens-17000), 1e-7)
+	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	tb, trickle = gtb.request(4000, uint64(time.Second)*10/uint64(time.Millisecond))
+	re.LessOrEqual(math.Abs(tb.Tokens-4000), 1e-7)
+	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	tb, trickle = gtb.request(19000, uint64(time.Second)*10/uint64(time.Millisecond))
+	re.LessOrEqual(math.Abs(tb.Tokens-18000), 1e-7)
+	re.Equal(trickle, int64(time.Second)*10/int64(time.Millisecond))
+	time2 := time1.Add(10 * time.Second)
+	gtb.update(time2)
+	tb, trickle = gtb.request(20000, uint64(time.Second)*10/uint64(time.Millisecond))
+	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

@@ -15,6 +15,7 @@
 package server
 
 import (
+	"math"
 	"time"
 
 	"github.com/gogo/protobuf/proto"
@@ -23,14 +24,31 @@ import (
 
 const defaultRefillRate = 10000
 
-const defaultInitialTokens = 10 * 10000
+const (
+	defaultInitialTokens = 10 * 10000
+	defaultMaxTokens     = 1e7
+)
+
+var reserveRatio float64 = 0.05
 
 // GroupTokenBucket is a token bucket for a resource group.
+// TODO: statistics Consumption
 type GroupTokenBucket struct {
 	*rmpb.TokenBucket `json:"token_bucket,omitempty"`
-	Consumption       *rmpb.TokenBucketsRequest `json:"consumption,omitempty"`
-	LastUpdate        *time.Time                `json:"last_update,omitempty"`
-	Initialized       bool                      `json:"initialized"`
+	// MaxTokens limits the number of tokens that can be accumulated
+	MaxTokens float64 `json:"max_tokens,omitempty"`
+
+	Consumption *rmpb.TokenBucketsRequest `json:"consumption,omitempty"`
+	LastUpdate  *time.Time                `json:"last_update,omitempty"`
+	Initialized bool                      `json:"initialized"`
+}
+
+// NewGroupTokenBucket returns a new GroupTokenBucket
+func NewGroupTokenBucket(tokenBucket *rmpb.TokenBucket) GroupTokenBucket {
+	return GroupTokenBucket{
+		TokenBucket: tokenBucket,
+		MaxTokens:   defaultMaxTokens,
+	}
 }
 
 // patch patches the token bucket settings.
@@ -51,15 +69,19 @@ func (t *GroupTokenBucket) patch(settings *rmpb.TokenBucket) {
 	t.TokenBucket = tb
 }
 
-// Update updates the token bucket.
-func (t *GroupTokenBucket) Update(now time.Time) {
+// update updates the token bucket.
+func (t *GroupTokenBucket) update(now time.Time) {
 	if !t.Initialized {
 		if t.Settings.FillRate == 0 {
 			t.Settings.FillRate = defaultRefillRate
 		}
 		if t.Tokens < defaultInitialTokens {
 			t.Tokens = defaultInitialTokens
 		}
+		// TODO: If we support init or modify MaxTokens in the future, we can move following code.
+		if t.Tokens > t.MaxTokens {
+			t.MaxTokens = t.Tokens
+		}
 		t.LastUpdate = &now
 		t.Initialized = true
 		return
@@ -70,12 +92,68 @@ func (t *GroupTokenBucket) Update(now time.Time) {
 		t.Tokens += float64(t.Settings.FillRate) * delta.Seconds()
 		t.LastUpdate = &now
 	}
+	if t.Tokens > t.MaxTokens {
+		t.Tokens = t.MaxTokens
+	}
 }
 
-// Request requests tokens from the token bucket.
-func (t *GroupTokenBucket) Request(
+// request requests tokens from the token bucket.
+func (t *GroupTokenBucket) request(
 	neededTokens float64, targetPeriodMs uint64,
-) *rmpb.TokenBucket {
-	// TODO: Implement the token bucket algorithm.
-	return nil
+) (*rmpb.TokenBucket, int64) {
+	var res rmpb.TokenBucket
+	res.Settings = &rmpb.TokenLimitSettings{}
+	// FillRate is used for the token server unavailable in abnormal situation.
+	if neededTokens <= 0 {
+		return &res, 0
+	}
+
+	// If the current tokens can directly meet the requirement, returns the need token
+	if t.Tokens >= neededTokens {
+		t.Tokens -= neededTokens
+		// granted the total request tokens
+		res.Tokens = neededTokens
+		return &res, 0
+	}
+
+	// Firstly allocate the remaining tokens
+	var grantedTokens float64
+	if t.Tokens > 0 {
+		grantedTokens = t.Tokens
+		neededTokens -= grantedTokens
+	}
+
+	var trickleTime = time.Duration(targetPeriodMs) * time.Millisecond
+	availableRate := float64(t.Settings.FillRate)
+	// When there are debt, the allotment will match the fill rate.
+	// We will have a threshold, beyond which the token allocation will be a minimum.
+	// the current threshold is fill rate * target period * 2.
+	// 				|
+	// fill rate	|· · · · · · · · ·
+	// 				|					·
+	//				|						·
+	// 				| 				  			·
+	// 				|								·
+	// reserve rate |									· · · ·
+	// 				|
+	// rate		0 	-----------------------------------------------
+	// 				debt 		period token		2*period token
+	if debt := -t.Tokens; debt > 0 {
+		debt -= float64(t.Settings.FillRate) * trickleTime.Seconds()
+		if debt > 0 {
+			debtRate := debt / float64(targetPeriodMs/1000)
+			availableRate -= debtRate
+			availableRate = math.Max(availableRate, reserveRatio*float64(t.Settings.FillRate))
+		}
+	}
+
+	consumptionDuration := time.Duration(float64(time.Second) * (neededTokens / availableRate))
+	if consumptionDuration <= trickleTime {
+		grantedTokens += neededTokens
+	} else {
+		grantedTokens += availableRate * trickleTime.Seconds()
+	}
+	t.Tokens -= grantedTokens
+	res.Tokens = grantedTokens
+	return &res, trickleTime.Milliseconds()
 }