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,50 @@ 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 {
+				return errors.New("resource group not found")
+			}
+			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(float64(re.Value), targetPeriodMs)
+						resp.GrantedRUTokens = append(resp.GrantedRUTokens, tokens)
+					case rmpb.RequestUnitType_WRU:
+						rg.UpdateWRU(now)
+						tokens := rg.RequestWRU(float64(re.Value), targetPeriodMs)
+						resp.GrantedRUTokens = append(resp.GrantedRUTokens, tokens)
+					}
+				}
+			case rmpb.GroupMode_NativeMode:
+				return errors.New("not supports the resource type")
+			}
+			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/token_bukets.go

@@ -25,12 +25,30 @@ const defaultRefillRate = 10000
 
 const defaultInitialTokens = 10 * 10000
 
+const defaultMaxTokens = 1e7
+
+const defaultLoanMaxPeriod = 10 * time.Second
+
+var loanReserveRatio 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"`
+	LoanExpireTime    *time.Time                `json:"loan_time,omitempty"`
+	LoanMaxPeriod     time.Duration             `json:"loan_max_perio,omitempty"`
+	MaxTokens         float64                   `json:"max_tokens,omitempty"`
+}
+
+func NewGroupTokenBucket(tokenBucket *rmpb.TokenBucket) GroupTokenBucket {
+	return GroupTokenBucket{
+		TokenBucket:   tokenBucket,
+		MaxTokens:     defaultMaxTokens,
+		LoanMaxPeriod: defaultLoanMaxPeriod,
+	}
 }
 
 // patch patches the token bucket settings.
@@ -51,8 +69,8 @@ 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 {
 		t.Settings.Fillrate = defaultRefillRate
 		t.Tokens = defaultInitialTokens
@@ -66,12 +84,65 @@ func (t *GroupTokenBucket) Update(now time.Time) {
 		t.Tokens += float64(t.Settings.Fillrate) * delta.Seconds()
 		t.LastUpdate = &now
 	}
+	if t.Tokens >= 0 {
+		t.LoanExpireTime = nil
+	}
+	if t.Tokens > defaultMaxTokens {
+		t.Tokens = defaultMaxTokens
+	}
+
 }
 
-// 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{}
+	// TODO: consider the shares for dispatch the fill rate
+	res.Settings.Fillrate = 0
+
+	if neededTokens <= 0 {
+		return &res, 0
+	}
+
+	if t.Tokens >= neededTokens {
+		t.Tokens -= neededTokens
+		// granted the total request tokens
+		res.Tokens = neededTokens
+		return &res, 0
+	}
+
+	var grantedTokens float64
+	if t.Tokens > 0 {
+		grantedTokens = t.Tokens
+		t.Tokens = 0
+		neededTokens -= grantedTokens
+	}
+
+	now := time.Now()
+	var periodFilled float64
+	var trickleTime int64 = int64(targetPeriodMs)
+	if t.LoanExpireTime != nil && t.LoanExpireTime.After(now) {
+		duration := t.LoanExpireTime.Sub(now)
+		periodFilled = float64(t.Settings.Fillrate) * (1 - loanReserveRatio) * duration.Seconds()
+		trickleTime = duration.Milliseconds()
+	} else {
+		now.Add(t.LoanMaxPeriod)
+		t.LoanExpireTime = &now
+		periodFilled = float64(t.Settings.Fillrate) * (1 - loanReserveRatio) * t.LoanMaxPeriod.Seconds()
+	}
+	periodFilled += t.Tokens
+	if periodFilled <= float64(t.Settings.Fillrate)*loanReserveRatio {
+		periodFilled = float64(t.Settings.Fillrate) * loanReserveRatio
+	}
+	if periodFilled >= neededTokens {
+		grantedTokens += neededTokens
+		t.Tokens -= neededTokens
+	} else {
+		grantedTokens += periodFilled
+		t.Tokens -= periodFilled
+	}
+	res.Tokens = grantedTokens
+	return &res, trickleTime
 }

pkg/mcs/resource_manager/server/types.go

@@ -18,6 +18,7 @@ package server
 import (
 	"encoding/json"
 	"sync"
+	"time"
 
 	"github.com/pingcap/errors"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
@@ -146,34 +147,64 @@ func FromProtoResourceGroup(group *rmpb.ResourceGroup) *ResourceGroup {
 	case rmpb.GroupMode_RUMode:
 		if settings := group.GetSettings().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_NativeMode:
 		if settings := group.GetSettings().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()