controller: fix the low_ru request missed (#8368)

client/go.mod

@@ -11,7 +11,7 @@ require (
  github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
  github.com/pingcap/kvproto v0.0.0-20230727073445-53e1f8730c30
  github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
- github.com/prometheus/client_golang v1.11.1
+ github.com/prometheus/client_golang v1.18.0
  github.com/stretchr/testify v1.8.2
  go.uber.org/atomic v1.10.0
  go.uber.org/goleak v1.1.11
@@ -26,18 +26,18 @@ require (
  github.com/cespare/xxhash/v2 v2.2.0 // indirect
  github.com/davecgh/go-spew v1.1.1 // indirect
  github.com/golang/protobuf v1.5.3 // indirect
- github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
  github.com/pkg/errors v0.9.1 // indirect
  github.com/pmezard/go-difflib v1.0.0 // indirect
- github.com/prometheus/client_model v0.2.0 // indirect
- github.com/prometheus/common v0.26.0 // indirect
- github.com/prometheus/procfs v0.6.0 // indirect
+ github.com/prometheus/client_model v0.5.0 // indirect
+ github.com/prometheus/common v0.46.0 // indirect
+ github.com/prometheus/procfs v0.12.0 // indirect
+ github.com/stretchr/objx v0.5.0 // indirect
  go.uber.org/multierr v1.11.0 // indirect
- golang.org/x/net v0.17.0 // indirect
- golang.org/x/sys v0.13.0 // indirect
- golang.org/x/text v0.13.0 // indirect
+ golang.org/x/net v0.20.0 // indirect
+ golang.org/x/sys v0.16.0 // indirect
+ golang.org/x/text v0.14.0 // indirect
  google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect
- google.golang.org/protobuf v1.30.0 // indirect
+ google.golang.org/protobuf v1.32.0 // indirect
  gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
  gopkg.in/yaml.v3 v3.0.1 // indirect
 )

client/resource_group/controller/controller.go

@@ -321,9 +321,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
  case notifyMsg := <-c.lowTokenNotifyChan:
  c.executeOnAllGroups((*groupCostController).updateRunState)
  c.executeOnAllGroups((*groupCostController).updateAvgRequestResourcePerSec)
- if len(c.run.currentRequests) == 0 {
- c.collectTokenBucketRequests(c.loopCtx, FromLowRU, lowToken /* select low tokens resource group */, notifyMsg)
- }
+ c.collectTokenBucketRequests(c.loopCtx, FromLowRU, lowToken /* select low tokens resource group */, notifyMsg)
  if c.run.inDegradedMode {
  c.executeOnAllGroups((*groupCostController).applyDegradedMode)
  }
@@ -1171,11 +1169,19 @@ func (gc *groupCostController) collectRequestAndConsumption(selectTyp selectType
  switch selectTyp {
  case periodicReport:
  selected = selected || gc.shouldReportConsumption()
+ failpoint.Inject("triggerPeriodicReport", func(val failpoint.Value) {
+ selected = gc.name == val.(string)
+ })
  fallthrough
  case lowToken:
  if counter.limiter.IsLowTokens() {
  selected = true
  }
+ failpoint.Inject("triggerLowRUReport", func(val failpoint.Value) {
+ if selectTyp == lowToken {
+ selected = gc.name == val.(string)
+ }
+ })
  }
  request := &rmpb.RequestUnitItem{
  Type: typ,

client/resource_group/controller/controller_test.go

@@ -24,8 +24,12 @@ import (
  "testing"
  "time"
 
+ "github.com/pingcap/failpoint"
+ "github.com/pingcap/kvproto/pkg/meta_storagepb"
  rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+ "github.com/stretchr/testify/mock"
  "github.com/stretchr/testify/require"
+ pd "github.com/tikv/pd/client"
  "github.com/tikv/pd/client/errs"
 )
 
@@ -132,3 +136,138 @@ func TestResourceGroupThrottledError(t *testing.T) {
  re.Error(err)
  re.True(errs.ErrClientResourceGroupThrottled.Equal(err))
 }
+
+// MockResourceGroupProvider is a mock implementation of the ResourceGroupProvider interface.
+type MockResourceGroupProvider struct {
+ mock.Mock
+}
+
+func (m *MockResourceGroupProvider) GetResourceGroup(ctx context.Context, resourceGroupName string) (*rmpb.ResourceGroup, error) {
+ args := m.Called(ctx, resourceGroupName)
+ return args.Get(0).(*rmpb.ResourceGroup), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) ListResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, error) {
+ args := m.Called(ctx)
+ return args.Get(0).([]*rmpb.ResourceGroup), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) AddResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
+ args := m.Called(ctx, metaGroup)
+ return args.String(0), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) ModifyResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
+ args := m.Called(ctx, metaGroup)
+ return args.String(0), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) DeleteResourceGroup(ctx context.Context, resourceGroupName string) (string, error) {
+ args := m.Called(ctx, resourceGroupName)
+ return args.String(0), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) AcquireTokenBuckets(ctx context.Context, request *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error) {
+ args := m.Called(ctx, request)
+ return args.Get(0).([]*rmpb.TokenBucketResponse), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) LoadResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, int64, error) {
+ args := m.Called(ctx)
+ return args.Get(0).([]*rmpb.ResourceGroup), args.Get(1).(int64), args.Error(2)
+}
+
+func (m *MockResourceGroupProvider) Watch(ctx context.Context, key []byte, opts ...pd.OpOption) (chan []*meta_storagepb.Event, error) {
+ args := m.Called(ctx, key, opts)
+ return args.Get(0).(chan []*meta_storagepb.Event), args.Error(1)
+}
+
+func (m *MockResourceGroupProvider) Get(ctx context.Context, key []byte, opts ...pd.OpOption) (*meta_storagepb.GetResponse, error) {
+ args := m.Called(ctx, key, opts)
+ return args.Get(0).(*meta_storagepb.GetResponse), args.Error(1)
+}
+
+func TestControllerWithTwoGroupRequestConcurrency(t *testing.T) {
+ re := require.New(t)
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ mockProvider := new(MockResourceGroupProvider)
+
+ mockProvider.On("Get", mock.Anything, mock.Anything, mock.Anything).Return(&meta_storagepb.GetResponse{}, nil)
+ // LoadResourceGroups
+ mockProvider.On("LoadResourceGroups", mock.Anything).Return([]*rmpb.ResourceGroup{}, int64(0), nil)
+ // Watch
+ mockProvider.On("Watch", mock.Anything, mock.Anything, mock.Anything).Return(make(chan []*meta_storagepb.Event), nil)
+
+ re.NoError(failpoint.Enable("github.com/tikv/pd/client/resource_group/controller/triggerPeriodicReport", fmt.Sprintf("return(\"%s\")", "default")))
+ defer failpoint.Disable("github.com/tikv/pd/client/resource_group/controller/triggerPeriodicReport")
+ re.NoError(failpoint.Enable("github.com/tikv/pd/client/resource_group/controller/triggerLowRUReport", fmt.Sprintf("return(\"%s\")", "test-group")))
+ defer failpoint.Disable("github.com/tikv/pd/client/resource_group/controller/triggerLowRUReport")
+
+ controller, _ := NewResourceGroupController(ctx, 1, mockProvider, nil)
+ controller.Start(ctx)
+
+ defaultResourceGroup := &rmpb.ResourceGroup{Name: "default", Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
+ testResourceGroup := &rmpb.ResourceGroup{Name: "test-group", Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
+ mockProvider.On("GetResourceGroup", mock.Anything, "default", mock.Anything).Return(defaultResourceGroup, nil)
+ mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil)
+
+ c1, err := controller.tryGetResourceGroup(ctx, "default")
+ re.NoError(err)
+ re.Equal(defaultResourceGroup, c1.meta)
+
+ c2, err := controller.tryGetResourceGroup(ctx, "test-group")
+ re.NoError(err)
+ re.Equal(testResourceGroup, c2.meta)
+
+ var expectResp []*rmpb.TokenBucketResponse
+ recTestGroupAcquireTokenRequest := make(chan bool)
+ mockProvider.On("AcquireTokenBuckets", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
+ request := args.Get(1).(*rmpb.TokenBucketsRequest)
+ var responses []*rmpb.TokenBucketResponse
+ for _, req := range request.Requests {
+ if req.ResourceGroupName == "default" {
+ // no response the default group request, that's mean `len(c.run.currentRequests) != 0` always.
+ time.Sleep(100 * time.Second)
+ responses = append(responses, &rmpb.TokenBucketResponse{
+ ResourceGroupName: "default",
+ GrantedRUTokens: []*rmpb.GrantedRUTokenBucket{
+ {
+ GrantedTokens: &rmpb.TokenBucket{
+ Tokens: 100000,
+ },
+ },
+ },
+ })
+ } else {
+ responses = append(responses, &rmpb.TokenBucketResponse{
+ ResourceGroupName: req.ResourceGroupName,
+ GrantedRUTokens: []*rmpb.GrantedRUTokenBucket{
+ {
+ GrantedTokens: &rmpb.TokenBucket{
+ Tokens: 100000,
+ },
+ },
+ },
+ })
+ }
+ }
+ // receive test-group request
+ if len(request.Requests) == 1 && request.Requests[0].ResourceGroupName == "test-group" {
+ recTestGroupAcquireTokenRequest <- true
+ }
+ expectResp = responses
+ }).Return(expectResp, nil)
+ // wait default group request token by PeriodicReport.
+ time.Sleep(2 * time.Second)
+ counter := c2.run.requestUnitTokens[0]
+ counter.limiter.mu.Lock()
+ counter.limiter.notify()
+ counter.limiter.mu.Unlock()
+ select {
+ case res := <-recTestGroupAcquireTokenRequest:
+ re.True(res)
+ case <-time.After(5 * time.Second):
+ re.Fail("timeout")
+ }
+}

client/resource_group/controller/limiter_test.go

@@ -44,6 +44,18 @@ var (
  t8 = t0.Add(time.Duration(8) * d)
 )
 
+func resetTime() {
+ t0 = time.Now()
+ t1 = t0.Add(time.Duration(1) * d)
+ t2 = t0.Add(time.Duration(2) * d)
+ t3 = t0.Add(time.Duration(3) * d)
+ t4 = t0.Add(time.Duration(4) * d)
+ t5 = t0.Add(time.Duration(5) * d)
+ t6 = t0.Add(time.Duration(6) * d)
+ t7 = t0.Add(time.Duration(7) * d)
+ t8 = t0.Add(time.Duration(8) * d)
+}
+
 type request struct {
  t time.Time
  n float64
@@ -144,6 +156,7 @@ func TestNotify(t *testing.T) {
 }
 
 func TestCancel(t *testing.T) {
+ resetTime()
  ctx := context.Background()
  ctx1, cancel1 := context.WithDeadline(ctx, t2)
  re := require.New(t)
@@ -161,8 +174,8 @@ func TestCancel(t *testing.T) {
  checkTokens(re, lim1, t2, 7)
  checkTokens(re, lim2, t2, 2)
  d, err := WaitReservations(ctx, t2, []*Reservation{r1, r2})
- re.Equal(d, 4*time.Second)
  re.Error(err)
+ re.Equal(4*time.Second, d)
  re.Contains(err.Error(), "estimated wait time 4s, ltb state is 1.00:-4.00")
  checkTokens(re, lim1, t3, 13)
  checkTokens(re, lim2, t3, 3)

tests/integrations/client/go.mod

@@ -108,7 +108,6 @@ require (
  github.com/mailru/easyjson v0.7.6 // indirect
  github.com/mattn/go-isatty v0.0.19 // indirect
  github.com/mattn/go-sqlite3 v1.14.15 // indirect
- github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
  github.com/minio/sio v0.3.0 // indirect
  github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
  github.com/modern-go/reflect2 v1.0.2 // indirect
@@ -126,10 +125,10 @@ require (
  github.com/pkg/errors v0.9.1 // indirect
  github.com/pmezard/go-difflib v1.0.0 // indirect
  github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect
- github.com/prometheus/client_golang v1.11.1 // indirect
- github.com/prometheus/client_model v0.2.0 // indirect
- github.com/prometheus/common v0.26.0 // indirect
- github.com/prometheus/procfs v0.6.0 // indirect
+ github.com/prometheus/client_golang v1.18.0 // indirect
+ github.com/prometheus/client_model v0.5.0 // indirect
+ github.com/prometheus/common v0.46.0 // indirect
+ github.com/prometheus/procfs v0.12.0 // indirect
  github.com/rs/cors v1.7.0 // indirect
  github.com/samber/lo v1.37.0 // indirect
  github.com/sasha-s/go-deadlock v0.2.0 // indirect
@@ -163,13 +162,13 @@ require (
  go.uber.org/fx v1.12.0 // indirect
  go.uber.org/multierr v1.11.0 // indirect
  golang.org/x/arch v0.3.0 // indirect
- golang.org/x/crypto v0.17.0 // indirect
+ golang.org/x/crypto v0.18.0 // indirect
  golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 // indirect
  golang.org/x/image v0.5.0 // indirect
- golang.org/x/net v0.17.0 // indirect
- golang.org/x/oauth2 v0.10.0 // indirect
+ golang.org/x/net v0.20.0 // indirect
+ golang.org/x/oauth2 v0.16.0 // indirect
  golang.org/x/sync v0.5.0 // indirect
- golang.org/x/sys v0.15.0 // indirect
+ golang.org/x/sys v0.16.0 // indirect
  golang.org/x/text v0.14.0 // indirect
  golang.org/x/time v0.1.0 // indirect
  golang.org/x/tools v0.6.0 // indirect