client, resource_manager: resource manager client

client/client.go

@@ -137,6 +137,8 @@ type Client interface {
 
 	// KeyspaceClient manages keyspace metadata.
 	KeyspaceClient
+	// ResourceManagerClient
+	ResourceManagerClient
 	// Close closes the client.
 	Close()
 }
@@ -390,6 +392,8 @@ type client struct {
 	// dc-location -> *lastTSO
 	lastTSMap sync.Map // Same as map[string]*lastTSO
 
+	tokenDispatcher *tokenDispatcher
+
 	// For internal usage.
 	checkTSDeadlineCh    chan struct{}
 	leaderNetworkFailure int32
@@ -417,6 +421,7 @@ func NewClientWithContext(ctx context.Context, pdAddrs []string, security Securi
 	}
 	// Start the daemons.
 	c.updateTSODispatcher()
+	c.createTokenispatcher()
 	c.wg.Add(3)
 	go c.tsLoop()
 	go c.tsCancelLoop()

client/go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
-	github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172
+	github.com/pingcap/kvproto v0.0.0-20230105060948-64890fa4f6c1
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/prometheus/client_golang v1.11.0
 	github.com/stretchr/testify v1.7.0

client/go.sum

@@ -104,8 +104,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 h1:C3N3itkduZXDZFh4N3vQ5HEtld3S+Y+StULhWVvumU0=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
-github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172 h1:FYgKV9znRQmzVrrJDZ0gUfMIvKLAMU1tu1UKJib8bEQ=
-github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20230105060948-64890fa4f6c1 h1:jw4NjEiCleRJPPpHM7K6l8OKzOjnZAj62eKteCAY6ro=
+github.com/pingcap/kvproto v0.0.0-20230105060948-64890fa4f6c1/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

client/resourcemanager_client.go

@@ -0,0 +1,282 @@
+// Copyright 2023 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,
+// 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 pd
+
+import (
+	"context"
+	"time"
+
+	"github.com/pingcap/errors"
+	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+	"google.golang.org/grpc"
+)
+
+type actionType int
+
+const (
+	add    actionType = 0
+	modify actionType = 1
+)
+
+// ResourceManagerClient manages resource group info and token request.
+type ResourceManagerClient interface {
+	ListResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, error)
+	GetResourceGroup(ctx context.Context, resourceGroupName string) (*rmpb.ResourceGroup, error)
+	AddResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error)
+	ModifyResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error)
+	DeleteResourceGroup(ctx context.Context, resourceGroupName string) (string, error)
+	AcquireTokenBuckets(ctx context.Context, request *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error)
+}
+
+// resouceManagerClient gets the ResourceManager client of current PD leader.
+func (c *client) resouceManagerClient() rmpb.ResourceManagerClient {
+	if cc, ok := c.clientConns.Load(c.GetLeaderAddr()); ok {
+		return rmpb.NewResourceManagerClient(cc.(*grpc.ClientConn))
+	}
+	return nil
+}
+
+// ListResourceGroups loads and returns target keyspace's metadata.
+func (c *client) ListResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, error) {
+	req := &rmpb.ListResourceGroupsRequest{}
+	resp, err := c.resouceManagerClient().ListResourceGroups(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	resErr := resp.GetError()
+	if resErr != nil {
+		return nil, errors.Errorf("[resource_manager]" + resErr.Message)
+	}
+	return resp.GetGroups(), nil
+}
+
+func (c *client) GetResourceGroup(ctx context.Context, resourceGroupName string) (*rmpb.ResourceGroup, error) {
+	req := &rmpb.GetResourceGroupRequest{
+		ResourceGroupName: resourceGroupName,
+	}
+	resp, err := c.resouceManagerClient().GetResourceGroup(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	resErr := resp.GetError()
+	if resErr != nil {
+		return nil, errors.Errorf("[resource_manager]" + resErr.Message)
+	}
+	return resp.GetGroup(), nil
+}
+
+func (c *client) AddResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
+	return c.putResourceGroup(ctx, metaGroup, add)
+}
+
+func (c *client) ModifyResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
+	return c.putResourceGroup(ctx, metaGroup, modify)
+}
+
+func (c *client) putResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup, typ actionType) (str string, err error) {
+	req := &rmpb.PutResourceGroupRequest{
+		Group: metaGroup,
+	}
+	var resp *rmpb.PutResourceGroupResponse
+	switch typ {
+	case add:
+		resp, err = c.resouceManagerClient().AddResourceGroup(ctx, req)
+	case modify:
+		resp, err = c.resouceManagerClient().ModifyResourceGroup(ctx, req)
+	}
+	if err != nil {
+		return str, err
+	}
+	resErr := resp.GetError()
+	if resErr != nil {
+		return str, errors.Errorf("[resource_manager]" + resErr.Message)
+	}
+	str = resp.GetBody()
+	return
+}
+
+func (c *client) DeleteResourceGroup(ctx context.Context, resourceGroupName string) (string, error) {
+	req := &rmpb.DeleteResourceGroupRequest{
+		ResourceGroupName: resourceGroupName,
+	}
+	resp, err := c.resouceManagerClient().DeleteResourceGroup(ctx, req)
+	if err != nil {
+		return "", err
+	}
+	resErr := resp.GetError()
+	if resErr != nil {
+		return "", errors.Errorf("[resource_manager]" + resErr.Message)
+	}
+	return resp.GetBody(), nil
+}
+
+func (c *client) AcquireTokenBuckets(ctx context.Context, request *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error) {
+	req := &tokenRequest{
+		done:       make(chan error, 1),
+		requestCtx: ctx,
+		clientCtx:  c.ctx,
+		Requeset:   request,
+	}
+	c.tokenDispatcher.tokenBatchController.tokenRequestCh <- req
+	grantedTokens, err := req.Wait()
+	if err != nil {
+		return nil, err
+	}
+	return grantedTokens, err
+}
+
+type tokenRequest struct {
+	clientCtx    context.Context
+	requestCtx   context.Context
+	done         chan error
+	Requeset     *rmpb.TokenBucketsRequest
+	TokenBuckets []*rmpb.TokenBucketResponse
+}
+
+func (req *tokenRequest) Wait() (tokenBuckets []*rmpb.TokenBucketResponse, err error) {
+	select {
+	case err = <-req.done:
+		err = errors.WithStack(err)
+		if err != nil {
+			return nil, err
+		}
+		tokenBuckets = req.TokenBuckets
+		return
+	case <-req.requestCtx.Done():
+		return nil, errors.WithStack(req.requestCtx.Err())
+	case <-req.clientCtx.Done():
+		return nil, errors.WithStack(req.clientCtx.Err())
+	}
+}
+
+type tokenBatchController struct {
+	tokenRequestCh chan *tokenRequest
+}
+
+func newTokenBatchController(tokenRequestCh chan *tokenRequest) *tokenBatchController {
+	return &tokenBatchController{
+		tokenRequestCh: tokenRequestCh,
+	}
+}
+
+type tokenDispatcher struct {
+	dispatcherCancel     context.CancelFunc
+	tokenBatchController *tokenBatchController
+}
+
+type resourceManagerConnectionContext struct {
+	stream rmpb.ResourceManager_AcquireTokenBucketsClient
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (c *client) createTokenispatcher() {
+	dispatcherCtx, dispatcherCancel := context.WithCancel(c.ctx)
+	dispatcher := &tokenDispatcher{
+		dispatcherCancel: dispatcherCancel,
+		tokenBatchController: newTokenBatchController(
+			make(chan *tokenRequest, 1)),
+	}
+	go c.handleResouceTokenDispatcher(dispatcherCtx, dispatcher.tokenBatchController)
+	c.tokenDispatcher = dispatcher
+}
+
+func (c *client) handleResouceTokenDispatcher(dispatcherCtx context.Context, tbc *tokenBatchController) {
+	var connection resourceManagerConnectionContext
+	if err := c.tryResourceManagerConnect(dispatcherCtx, &connection); err != nil {
+		log.Warn("get stream error", zap.Error(err))
+	}
+
+	for {
+		var firstRequest *tokenRequest
+		select {
+		case <-dispatcherCtx.Done():
+			return
+		case firstRequest = <-tbc.tokenRequestCh:
+		}
+		stream, streamCtx, cancel := connection.stream, connection.ctx, connection.cancel
+		if stream == nil {
+			c.tryResourceManagerConnect(dispatcherCtx, &connection)
+			c.finishTokenRequest(firstRequest, nil, errors.Errorf("no stream"))
+			continue
+		}
+		select {
+		case <-streamCtx.Done():
+			log.Info("[resource_manager] resource manager stream is canceled")
+			cancel()
+			stream = nil
+			continue
+		default:
+		}
+		err := c.processTokenRequests(stream, firstRequest)
+		if err != nil {
+			log.Info("processTokenRequests error", zap.Error(err))
+			cancel()
+			connection.stream = nil
+		}
+	}
+}
+
+func (c *client) processTokenRequests(stream rmpb.ResourceManager_AcquireTokenBucketsClient, t *tokenRequest) error {
+	req := t.Requeset
+	if err := stream.Send(req); err != nil {
+		err = errors.WithStack(err)
+		c.finishTokenRequest(t, nil, err)
+		return err
+	}
+	resp, err := stream.Recv()
+	if err != nil {
+		err = errors.WithStack(err)
+		c.finishTokenRequest(t, nil, err)
+		return err
+	}
+	if resp.GetError() != nil {
+		return errors.Errorf("[resource_manager]" + resp.GetError().Message)
+	}
+	tokenBuckets := resp.GetResponses()
+	c.finishTokenRequest(t, tokenBuckets, nil)
+	return nil
+}
+
+func (c *client) finishTokenRequest(t *tokenRequest, tokenBuckets []*rmpb.TokenBucketResponse, err error) {
+	t.TokenBuckets = tokenBuckets
+	t.done <- err
+}
+
+func (c *client) tryResourceManagerConnect(ctx context.Context, connection *resourceManagerConnectionContext) error {
+	var (
+		err    error
+		stream rmpb.ResourceManager_AcquireTokenBucketsClient
+	)
+	for i := 0; i < maxRetryTimes; i++ {
+		cctx, cancel := context.WithCancel(ctx)
+		stream, err = c.resouceManagerClient().AcquireTokenBuckets(cctx)
+		if err == nil && stream != nil {
+			connection.cancel = cancel
+			connection.ctx = cctx
+			connection.stream = stream
+			return nil
+		}
+		cancel()
+		select {
+		case <-ctx.Done():
+			return err
+		case <-time.After(retryInterval):
+		}
+	}
+	return err
+}

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)
+	}
 }