Skip to content

server, client: expose store scheduling state in GetStore #11041

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
LykxSassinator wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

server, client: expose store scheduling state in GetStore #11041

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a31cb2c
client: expose store scheduling state in GetStore
LykxSassinator Jul 23, 2026
d09cfb9
client: keep GetStore response API source-compatible
LykxSassinator Jul 23, 2026
5458f09
tools: use kvproto PR branch
LykxSassinator Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion client/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,14 @@ type RPCClient interface {
ResourceManagerClient
}

// RPCClientExt is an optional extension for callers that need fields carried
// by complete RPC responses in addition to the existing RPCClient methods. It
// is deliberately separate from RPCClient so new optional response APIs do not
// break existing pd.Client implementations and mocks.
type RPCClientExt interface {
GetStoreResponse(ctx context.Context, storeID uint64, opts ...opt.GetStoreOption) (*pdpb.GetStoreResponse, error)
}

// Client is a PD (Placement Driver) RPC client.
// It should not be used after calling Close().
type Client interface {
Expand Down Expand Up @@ -1033,6 +1041,15 @@ func handleRegionsResponse(resp *pdpb.ScanRegionsResponse) []*router.Region {

// GetStore implements the RPCClient interface.
func (c *client) GetStore(ctx context.Context, storeID uint64, opts ...opt.GetStoreOption) (*metapb.Store, error) {
resp, err := c.GetStoreResponse(ctx, storeID, opts...)
if err != nil {
return nil, err
}
return handleStoreResponse(resp)
}

// GetStoreResponse implements the RPCClientExt interface.
func (c *client) GetStoreResponse(ctx context.Context, storeID uint64, opts ...opt.GetStoreOption) (*pdpb.GetStoreResponse, error) {
if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
span = span.Tracer().StartSpan("pdclient.GetStore", opentracing.ChildOf(span.Context()))
defer span.Finish()
Expand Down Expand Up @@ -1077,7 +1094,7 @@ func (c *client) GetStore(ctx context.Context, storeID uint64, opts ...opt.GetSt
if err = c.respForErr(metrics.CmdFailedDurationGetStore, start, err, resp.GetHeader()); err != nil {
return nil, err
}
return handleStoreResponse(resp)
return resp, nil
}

func handleStoreResponse(resp *pdpb.GetStoreResponse) (*metapb.Store, error) {
Expand Down
3 changes: 3 additions & 0 deletions client/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@ module github.com/tikv/pd/client

go 1.25.10

// Temporary replacement for kvproto PR #1503. Remove it after the PR merges.
replace github.com/pingcap/kvproto => github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0

require (
github.com/BurntSushi/toml v0.3.1
github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
Expand Down
4 changes: 2 additions & 2 deletions client/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0 h1:W1he+VK7QW2ZrVt8GiDK4u1yUF7FsQbi95XRT/pJZrs=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
Expand Down Expand Up @@ -53,8 +55,6 @@ 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-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0 h1:MalBpLjhK/cS9ndCoSAg2C7aBzVojAqFVfzTKmuy0ho=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
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.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
Expand Down
2 changes: 1 addition & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ go 1.25.10
// When you modify PD cooperatively with kvproto, this will be useful to submit the PR to PD and the PR to
// kvproto at the same time. You can run `go mod tidy` to make it replaced with go-mod style specification.
// After the PR to kvproto is merged, remember to comment this out and run `go mod tidy`.
// replace github.com/pingcap/kvproto => github.com/$YourPrivateRepo $YourPrivateBranch
replace github.com/pingcap/kvproto => github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0

require (
github.com/AlekSi/gocov-xml v1.0.0
Expand Down
2 changes: 2 additions & 0 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -479,6 +479,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-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0 h1:W1he+VK7QW2ZrVt8GiDK4u1yUF7FsQbi95XRT/pJZrs=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0 h1:MalBpLjhK/cS9ndCoSAg2C7aBzVojAqFVfzTKmuy0ho=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
Expand Down
16 changes: 16 additions & 0 deletions pkg/utils/grpcutil/cluster_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (

"github.com/stretchr/testify/suite"

"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pingcap/kvproto/pkg/pdpb"

"github.com/tikv/pd/pkg/core"
Expand Down Expand Up @@ -216,3 +217,18 @@ func (suite *clusterSuite) TestGetRegionByPreKey() {
})
}
}

func (suite *clusterSuite) TestGetStoreDoesNotReportSchedulingState() {
re := suite.Require()
suite.rc.PutStore(core.NewStoreInfo(
&metapb.Store{Id: 1, Address: "tikv-1"},
core.SlowStoreEvicted(),
))

resp, err := GetStore(suite.rc, &pdpb.GetStoreRequest{StoreId: 1})
re.NoError(err)
re.NotNil(resp.GetStore())
// BasicCluster is used by the router and does not retain dynamic PD leader
// scheduling decisions. Returning nil preserves the unknown state.
re.Nil(resp.GetSchedulingState())
}
3 changes: 3 additions & 0 deletions server/grpc_service.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -786,6 +786,9 @@ func (s *GrpcServer) GetStore(ctx context.Context, request *pdpb.GetStoreRequest
Header: grpcutil.WrapHeader(),
Store: store.GetMeta(),
Stats: store.GetStoreStats(),
SchedulingState: &pdpb.StoreSchedulingState{
EvictedAsSlowStore: store.EvictedAsSlowStore(),
},
}, nil
}

Expand Down
12 changes: 12 additions & 0 deletions tests/integrations/client/client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1276,6 +1276,18 @@ func (suite *clientStatelessTestSuite) TestGetStore() {
store.LastHeartbeat = n.LastHeartbeat
re.Equal(store, n)

// GetStoreResponse retains PD's dynamic scheduling state, which is not part
// of the store metadata returned by GetStore.
err = cluster.SlowStoreEvicted(store.GetId())
re.NoError(err)
storeResp, err := suite.client.GetStoreResponse(context.Background(), store.GetId())
re.NoError(err)
re.True(storeResp.GetSchedulingState().GetEvictedAsSlowStore())
cluster.SlowStoreRecovered(store.GetId())
storeResp, err = suite.client.GetStoreResponse(context.Background(), store.GetId())
re.NoError(err)
re.False(storeResp.GetSchedulingState().GetEvictedAsSlowStore())

actualStores, err := suite.client.GetAllStores(context.Background())
re.NoError(err)
re.Len(actualStores, len(stores))
Expand Down
2 changes: 2 additions & 0 deletions tests/integrations/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ module github.com/tikv/pd/tests/integrations
go 1.25.10

replace (
// Temporary replacement for kvproto PR #1503. Remove it after the PR merges.
github.com/pingcap/kvproto => github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0
github.com/tikv/pd => ../../
github.com/tikv/pd/client => ../../client
github.com/tikv/pd/tests/integrations/mcs => ./mcs
Expand Down
2 changes: 2 additions & 0 deletions tests/integrations/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -472,6 +472,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-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0 h1:W1he+VK7QW2ZrVt8GiDK4u1yUF7FsQbi95XRT/pJZrs=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0 h1:MalBpLjhK/cS9ndCoSAg2C7aBzVojAqFVfzTKmuy0ho=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
Expand Down
2 changes: 2 additions & 0 deletions tools/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ module github.com/tikv/pd/tools
go 1.25.10

replace (
// Temporary replacement for kvproto PR #1503. Remove it after the PR merges.
github.com/pingcap/kvproto => github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0
github.com/tikv/pd => ../
github.com/tikv/pd/client => ../client
)
Expand Down
2 changes: 2 additions & 0 deletions tools/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg
github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc=
github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0 h1:W1he+VK7QW2ZrVt8GiDK4u1yUF7FsQbi95XRT/pJZrs=
github.com/LykxSassinator/kvproto v0.0.0-20260723025057-6a82baebdef0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/Masterminds/semver v1.5.0 h1:H65muMkzWKEuNDnfl9d70GUjFniHKHRbFPGBuZ3QEww=
github.com/Masterminds/semver v1.5.0/go.mod h1:MB6lktGJrhw8PrUyiEoblNEGEQ+RzHPF078ddwwvV3Y=
github.com/PuerkitoBio/purell v1.1.1 h1:WEQqlqaGbrPkxLJWfBwQmfEAE1Z7ONdDLqrN38tNFfI=
Expand Down
Loading