Skip to content

Commit 78966d8

Browse files
committedJul 17, 2023
update test
Signed-off-by: husharp <jinhao.hu@pingcap.com>
1 parent 40eaa35 commit 78966d8

File tree

19 files changed

+268
-75
lines changed

19 files changed

+268
-75
lines changed
 

‎client/go.mod

+1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ require (
1414
github.com/stretchr/testify v1.8.2
1515
go.uber.org/goleak v1.1.11
1616
go.uber.org/zap v1.24.0
17+
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4
1718
google.golang.org/grpc v1.54.0
1819
)
1920

‎client/go.sum

+2
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,8 @@ golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnf
152152
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
153153
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
154154
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
155+
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4 h1:QLureRX3moex6NVu/Lr4MGakp9FdA7sBHGBmvRW7NaM=
156+
golang.org/x/exp v0.0.0-20230711005742-c3f37128e5a4/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc=
155157
golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
156158
golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
157159
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=

‎client/resource_group/controller/controller.go

+28-3
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"context"
1919
"encoding/json"
2020
"math"
21+
"strings"
2122
"sync"
2223
"sync/atomic"
2324
"time"
@@ -32,6 +33,7 @@ import (
3233
pd "github.com/tikv/pd/client"
3334
"github.com/tikv/pd/client/errs"
3435
"go.uber.org/zap"
36+
"golang.org/x/exp/slices"
3537
)
3638

3739
const (
@@ -56,8 +58,10 @@ const (
5658
type ResourceGroupKVInterceptor interface {
5759
// OnRequestWait is used to check whether resource group has enough tokens. It maybe needs to wait some time.
5860
OnRequestWait(ctx context.Context, resourceGroupName string, info RequestInfo) (*rmpb.Consumption, *rmpb.Consumption, error)
59-
// OnResponse is used to consume tokens after receiving response
61+
// OnResponse is used to consume tokens after receiving response.
6062
OnResponse(resourceGroupName string, req RequestInfo, resp ResponseInfo) (*rmpb.Consumption, error)
63+
// IsBackgroundRequest If the resource group has background jobs, we should not record consumption and wait for it.
64+
IsBackgroundRequest(ctx context.Context, resourceGroupName, requestResource string) bool
6165
}
6266

6367
// ResourceGroupProvider provides some api to interact with resource manager server.
@@ -454,7 +458,6 @@ func (c *ResourceGroupsController) OnRequestWait(
454458
) (*rmpb.Consumption, *rmpb.Consumption, error) {
455459
gc, err := c.tryGetResourceGroup(ctx, resourceGroupName)
456460
if err != nil {
457-
failedRequestCounter.WithLabelValues(resourceGroupName).Inc()
458461
return nil, nil, err
459462
}
460463
return gc.onRequestWait(ctx, info)
@@ -472,6 +475,28 @@ func (c *ResourceGroupsController) OnResponse(
472475
return tmp.(*groupCostController).onResponse(req, resp)
473476
}
474477

478+
// IsBackgroundRequest If the resource group has background jobs, we should not record consumption and wait for it.
479+
func (c *ResourceGroupsController) IsBackgroundRequest(ctx context.Context,
480+
resourceGroupName, requestResource string) bool {
481+
gc, err := c.tryGetResourceGroup(ctx, resourceGroupName)
482+
if err != nil {
483+
failedRequestCounter.WithLabelValues(resourceGroupName).Inc()
484+
return false
485+
}
486+
487+
gc.metaLock.RLock()
488+
defer gc.metaLock.RUnlock()
489+
if bg := gc.meta.BackgroundSettings; bg != nil {
490+
if len(requestResource) == 0 || len(bg.JobTypes) == 0 {
491+
return false
492+
}
493+
if idx := strings.LastIndex(requestResource, "_"); idx != -1 {
494+
return slices.Contains(bg.JobTypes, requestResource[idx+1:])
495+
}
496+
}
497+
return false
498+
}
499+
475500
// GetResourceGroup returns the meta setting of the given resource group name.
476501
func (c *ResourceGroupsController) GetResourceGroup(resourceGroupName string) (*rmpb.ResourceGroup, error) {
477502
gc, err := c.tryGetResourceGroup(c.loopCtx, resourceGroupName)
@@ -518,7 +543,7 @@ type groupCostController struct {
518543
lastRequestTime time.Time
519544

520545
// requestInProgress is set true when sending token bucket request.
521-
// And it is set false when reciving token bucket response.
546+
// And it is set false when receiving token bucket response.
522547
// This triggers a retry attempt on the next tick.
523548
requestInProgress bool
524549

‎client/resource_group/controller/controller_test.go

+3
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ func createTestGroupCostController(re *require.Assertions) *groupCostController
3939
},
4040
},
4141
},
42+
BackgroundSettings: &rmpb.BackgroundSettings{
43+
JobTypes: []string{"lightning", "br"},
44+
},
4245
}
4346
ch1 := make(chan struct{})
4447
ch2 := make(chan *groupCostController)

‎client/resource_group/controller/model.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import (
1818
"os"
1919
"time"
2020

21-
"github.com/cloudfoundry/gosigar"
21+
sigar "github.com/cloudfoundry/gosigar"
2222
"go.uber.org/zap"
2323

2424
rmpb "github.com/pingcap/kvproto/pkg/resource_manager"

‎go.mod

+1
Original file line numberDiff line numberDiff line change
@@ -205,3 +205,4 @@ replace google.golang.org/grpc v1.54.0 => google.golang.org/grpc v1.26.0
205205
// kvproto at the same time. You can run `go mod tidy` to make it replaced with go-mod style specification.
206206
// After the PR to kvproto is merged, remember to comment this out and run `go mod tidy`.
207207
// replace github.com/pingcap/kvproto => github.com/$YourPrivateRepo $YourPrivateBranch
208+
replace github.com/pingcap/kvproto => github.com/HuSharp/kvproto v0.0.0-20230717074128-3202750f4c06

0 commit comments

Comments
 (0)