Skip to content

Commit

Permalink
add tests to verify that circuit breaker is set for GetRegion calls
Browse files Browse the repository at this point in the history
Signed-off-by: artem_danilov <[email protected]>
  • Loading branch information
artem_danilov committed Jan 8, 2025
1 parent 01fadf5 commit 7ab69c9
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion internal/mockstore/mocktikv/pd.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,8 @@ import (
"github.com/tikv/pd/client/clients/tso"
"github.com/tikv/pd/client/opt"
"github.com/tikv/pd/client/pkg/caller"
sd "github.com/tikv/pd/client/servicediscovery"
"github.com/tikv/pd/client/pkg/circuitbreaker"
sd "github.com/tikv/pd/client/servicediscovery"

Check failure on line 59 in internal/mockstore/mocktikv/pd.go

View workflow job for this annotation

GitHub Actions / golangci

File is not `gofmt`-ed with `-s` (gofmt)
"go.uber.org/atomic"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
Expand Down Expand Up @@ -226,6 +227,7 @@ func (m *mockTSFuture) Wait() (int64, int64, error) {
}

func (c *pdClient) GetRegion(ctx context.Context, key []byte, opts ...opt.GetRegionOption) (*router.Region, error) {
enforceCircuitBreakerFor("GetRegion", ctx)
region, peer, buckets, downPeers := c.cluster.GetRegionByKey(key)
if len(opts) == 0 {
buckets = nil
Expand All @@ -244,6 +246,7 @@ func (c *pdClient) GetRegionFromMember(ctx context.Context, key []byte, memberUR
}

func (c *pdClient) GetPrevRegion(ctx context.Context, key []byte, opts ...opt.GetRegionOption) (*router.Region, error) {
enforceCircuitBreakerFor("GetPrevRegion", ctx)
region, peer, buckets, downPeers := c.cluster.GetPrevRegionByKey(key)
if len(opts) == 0 {
buckets = nil
Expand All @@ -252,6 +255,7 @@ func (c *pdClient) GetPrevRegion(ctx context.Context, key []byte, opts ...opt.Ge
}

func (c *pdClient) GetRegionByID(ctx context.Context, regionID uint64, opts ...opt.GetRegionOption) (*router.Region, error) {
enforceCircuitBreakerFor("GetRegionByID", ctx)
region, peer, buckets, downPeers := c.cluster.GetRegionByID(regionID)
return &router.Region{Meta: region, Leader: peer, Buckets: buckets, DownPeers: downPeers}, nil
}
Expand Down Expand Up @@ -465,3 +469,9 @@ func (m *pdClient) LoadResourceGroups(ctx context.Context) ([]*rmpb.ResourceGrou
func (m *pdClient) GetServiceDiscovery() sd.ServiceDiscovery { return nil }

func (m *pdClient) WithCallerComponent(caller.Component) pd.Client { return m }

func enforceCircuitBreakerFor(name string, ctx context.Context) {
if circuitbreaker.FromContext(ctx) == nil {
panic(fmt.Errorf("CircuitBreaker must be configured for %s", name))
}
}

0 comments on commit 7ab69c9

Please sign in to comment.