*: Add caller information to pd client (#59911)

okJiang · web-flow · commit f1b85a769a05 · 2026-05-18T17:54:35.000Z
close #59910
diff --git a/br/pkg/pdutil/pd.go b/br/pkg/pdutil/pd.go
@@ -159,7 +159,7 @@ func NewPdController(
 		grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(maxMsgSize)),
 	}
 	pdClient, err := pd.NewClientWithContext(
-		ctx, caller.Component("br-pd-controller"), pdAddrs, securityOption,
+		ctx, caller.GetComponent(1), pdAddrs, securityOption,
 		opt.WithGRPCDialOptions(maxCallMsgSize...),
 		// If the time too short, we may scatter a region many times, because
 		// the interface `ScatterRegions` may time out.
diff --git a/br/pkg/restore/split/BUILD.bazel b/br/pkg/restore/split/BUILD.bazel
@@ -44,6 +44,7 @@ go_library(
         "@com_github_tikv_pd_client//clients/router",
         "@com_github_tikv_pd_client//http",
         "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials",
diff --git a/br/pkg/restore/split/client.go b/br/pkg/restore/split/client.go
@@ -36,6 +36,7 @@ import (
 	pd "github.com/tikv/pd/client"
 	pdhttp "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/multierr"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
@@ -159,6 +160,14 @@ func WithOnSplit(onSplit func(key [][]byte)) ClientOptionalParameter {
 	}
 }
 
+func withCallerComponent(client pd.Client, component caller.Component) pd.Client {
+	if _, ok := client.(*tikvclient.CodecPDClient); ok {
+		// Keep codec-aware clients intact so callers can retrieve the same wrapper.
+		return client
+	}
+	return client.WithCallerComponent(component)
+}
+
 // NewClient creates a SplitClient.
 //
 // splitBatchKeyCnt controls how many keys are sent to TiKV in a batch in split
@@ -172,7 +181,7 @@ func NewClient(
 	opts ...ClientOptionalParameter,
 ) SplitClient {
 	cli := &pdClient{
-		client:           client,
+		client:           withCallerComponent(client, caller.GetComponent(1)),
 		httpCli:          httpCli,
 		tlsConf:          tlsConf,
 		storeCache:       make(map[uint64]*metapb.Store),
diff --git a/br/pkg/restore/split/mock_pd_client.go b/br/pkg/restore/split/mock_pd_client.go
@@ -23,6 +23,7 @@ import (
 	"github.com/tikv/pd/client/clients/router"
 	pdhttp "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/status"
@@ -220,6 +221,10 @@ func newRegionNotFullyReplicatedErr(regionID uint64) error {
 	return status.Errorf(codes.Unknown, "region %d is not fully replicated", regionID)
 }
 
+func (c *MockPDClientForSplit) WithCallerComponent(_ caller.Component) pd.Client {
+	return c
+}
+
 func (c *MockPDClientForSplit) SetRegions(boundaries [][]byte) []*metapb.Region {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -561,6 +566,10 @@ func (fpdc *FakePDClient) SetRegions(regions []*router.Region) {
 	fpdc.regions = regions
 }
 
+func (fpdc *FakePDClient) WithCallerComponent(_ caller.Component) pd.Client {
+	return fpdc
+}
+
 func (fpdc *FakePDClient) GetAllStores(context.Context, ...opt.GetStoreOption) ([]*metapb.Store, error) {
 	return slices.Clone(fpdc.stores), nil
 }
diff --git a/br/pkg/streamhelper/BUILD.bazel b/br/pkg/streamhelper/BUILD.bazel
@@ -46,6 +46,7 @@ go_library(
         "@com_github_tikv_client_go_v2//txnkv/txnlock",
         "@com_github_tikv_pd_client//:client",
         "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@io_etcd_go_etcd_client_v3//:client",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
diff --git a/br/pkg/streamhelper/advancer_env.go b/br/pkg/streamhelper/advancer_env.go
@@ -17,6 +17,7 @@ import (
 	"github.com/tikv/client-go/v2/txnkv/txnlock"
 	pd "github.com/tikv/pd/client"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
@@ -140,6 +141,7 @@ func (t clusterEnv) ClearCache(ctx context.Context, storeID uint64) error {
 
 // CliEnv creates the Env for CLI usage.
 func CliEnv(cli *utils.StoreManager, tikvStore tikv.Storage, etcdCli *clientv3.Client) Env {
+	cli.ResetPDClientCallerComponent(caller.Pitr)
 	return clusterEnv{
 		clis:                 cli,
 		AdvancerExt:          &AdvancerExt{MetaDataClient: *NewMetaDataClient(etcdCli)},
@@ -154,13 +156,14 @@ func TiDBEnv(tikvStore tikv.Storage, pdCli pd.Client, etcdCli *clientv3.Client,
 	if err != nil {
 		return nil, err
 	}
+	pitrPDClient := pdCli.WithCallerComponent(caller.Pitr)
 	env := clusterEnv{
-		clis: utils.NewStoreManager(pdCli, keepalive.ClientParameters{
+		clis: utils.NewStoreManager(pitrPDClient, keepalive.ClientParameters{
 			Time:    time.Duration(conf.TiKVClient.GrpcKeepAliveTime) * time.Second,
 			Timeout: time.Duration(conf.TiKVClient.GrpcKeepAliveTimeout) * time.Second,
 		}, tconf),
 		AdvancerExt:          &AdvancerExt{MetaDataClient: *NewMetaDataClient(etcdCli)},
-		PDRegionScanner:      PDRegionScanner{Client: pdCli},
+		PDRegionScanner:      PDRegionScanner{Client: pitrPDClient},
 		AdvancerLockResolver: newAdvancerLockResolver(tikvStore),
 	}
 
diff --git a/br/pkg/utils/BUILD.bazel b/br/pkg/utils/BUILD.bazel
@@ -62,6 +62,7 @@ go_library(
         "@com_github_prometheus_client_golang//prometheus/promhttp",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_pd_client//:client",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@io_etcd_go_etcd_api_v3//v3rpc/rpctypes",
         "@io_etcd_go_etcd_client_v3//:client",
         "@org_golang_google_grpc//:grpc",
diff --git a/br/pkg/utils/store_manager.go b/br/pkg/utils/store_manager.go
@@ -16,6 +16,7 @@ import (
 	berrors "github.com/pingcap/tidb/br/pkg/errors"
 	"github.com/pingcap/tidb/br/pkg/logutil"
 	pd "github.com/tikv/pd/client"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
@@ -127,6 +128,10 @@ func (mgr *StoreManager) PDClient() pd.Client {
 	return mgr.pdClient
 }
 
+func (mgr *StoreManager) ResetPDClientCallerComponent(component caller.Component) {
+	mgr.pdClient = mgr.pdClient.WithCallerComponent(component)
+}
+
 func (mgr *StoreManager) getGrpcConnLocked(ctx context.Context, storeID uint64) (*grpc.ClientConn, error) {
 	failpoint.Inject("hint-get-backup-client", func(v failpoint.Value) {
 		log.Info("failpoint hint-get-backup-client injected, "+
diff --git a/pkg/ddl/BUILD.bazel b/pkg/ddl/BUILD.bazel
@@ -221,6 +221,7 @@ go_library(
         "@com_github_tikv_pd_client//clients/router",
         "@com_github_tikv_pd_client//http",
         "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@io_etcd_go_etcd_client_v3//:client",
         "@org_golang_x_sync//errgroup",
         "@org_uber_go_atomic//:atomic",
diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go
@@ -3383,10 +3383,7 @@ func estimateRowSizeFromRegion(ctx context.Context, store kv.Storage, tbl table.
 	if !ok {
 		return 0, fmt.Errorf("not a helper.Storage")
 	}
-	h := &helper.Helper{
-		Store:       hStore,
-		RegionCache: hStore.GetRegionCache(),
-	}
+	h := helper.NewHelper(hStore)
 	pdCli, err := h.TryGetPDHTTPClient()
 	if err != nil {
 		return 0, err
diff --git a/pkg/ddl/ingest/BUILD.bazel b/pkg/ddl/ingest/BUILD.bazel
@@ -58,6 +58,7 @@ go_library(
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//util",
         "@com_github_tikv_pd_client//:client",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@io_etcd_go_etcd_client_v3//:client",
         "@org_golang_x_exp//maps",
         "@org_uber_go_atomic//:atomic",
diff --git a/pkg/ddl/ingest/backend_mgr.go b/pkg/ddl/ingest/backend_mgr.go
@@ -33,6 +33,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/tikv/client-go/v2/tikv"
+	"github.com/tikv/pd/client/pkg/caller"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
@@ -133,7 +134,7 @@ func (b *BackendCtxBuilder) Build(cfg *ingestctrl.BackendConfig, bd *ingestctrl.
 	})
 
 	//nolint: forcetypeassert
-	pdCli := store.(tikv.Storage).GetRegionCache().PDClient()
+	pdCli := store.(tikv.Storage).GetRegionCache().PDClient().WithCallerComponent(caller.Ddl)
 	var cpOp CheckpointOperator
 
 	// Create checkpoint manager based on the configuration
diff --git a/pkg/ddl/partition.go b/pkg/ddl/partition.go
@@ -67,6 +67,7 @@ import (
 	kvutil "github.com/tikv/client-go/v2/util"
 	"github.com/tikv/pd/client/clients/router"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/zap"
 )
 
@@ -429,7 +430,7 @@ func checkPartitionReplica(replicaCount uint64, addingDefinitions []model.Partit
 	})
 
 	ctx := context.Background()
-	pdCli := jobCtx.store.(tikv.Storage).GetRegionCache().PDClient()
+	pdCli := jobCtx.store.(tikv.Storage).GetRegionCache().PDClient().WithCallerComponent(caller.Ddl)
 	stores, err := pdCli.GetAllStores(ctx)
 	if err != nil {
 		return needWait, errors.Trace(err)
diff --git a/pkg/domain/BUILD.bazel b/pkg/domain/BUILD.bazel
@@ -125,6 +125,7 @@ go_library(
         "@com_github_tikv_pd_client//constants",
         "@com_github_tikv_pd_client//http",
         "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@com_github_tikv_pd_client//pkg/circuitbreaker",
         "@com_github_tikv_pd_client//resource_group/controller",
         "@io_etcd_go_etcd_client_v3//:client",
diff --git a/pkg/domain/domain.go b/pkg/domain/domain.go
@@ -115,6 +115,7 @@ import (
 	pd "github.com/tikv/pd/client"
 	pdhttp "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	rmclient "github.com/tikv/pd/client/resource_group/controller"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/client/v3/concurrency"
@@ -1267,7 +1268,7 @@ func (do *Domain) AutoIDClient() *autoid.ClientDiscover {
 // GetPDClient returns the PD client.
 func (do *Domain) GetPDClient() pd.Client {
 	if store, ok := do.store.(kv.StorageWithPD); ok {
-		return store.GetPDClient()
+		return store.GetPDClient().WithCallerComponent(caller.GetComponent(1))
 	}
 	return nil
 }
diff --git a/pkg/executor/analyze_utils.go b/pkg/executor/analyze_utils.go
@@ -43,10 +43,7 @@ func adaptiveAnlayzeDistSQLConcurrency(ctx context.Context, sctx sessionctx.Cont
 		logutil.BgLogger().Warn("Information about TiKV store status can be gotten only when the storage is TiKV")
 		return vardef.DefAnalyzeDistSQLScanConcurrency
 	}
-	tikvHelper := &helper.Helper{
-		Store:       tikvStore,
-		RegionCache: tikvStore.GetRegionCache(),
-	}
+	tikvHelper := helper.NewHelper(tikvStore)
 	pdCli, err := tikvHelper.TryGetPDHTTPClient()
 	if err != nil {
 		logutil.BgLogger().Warn("fail to TryGetPDHTTPClient", zap.Error(err))
diff --git a/pkg/server/handler/tikvhandler/BUILD.bazel b/pkg/server/handler/tikvhandler/BUILD.bazel
@@ -58,9 +58,11 @@ go_library(
         "@com_github_pingcap_log//:log",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//util",
+        "@com_github_tikv_pd_client//:client",
         "@com_github_tikv_pd_client//clients/router",
         "@com_github_tikv_pd_client//http",
         "@com_github_tikv_pd_client//opt",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@org_uber_go_zap//:zap",
     ],
 )
diff --git a/pkg/server/handler/tikvhandler/tikv_handler.go b/pkg/server/handler/tikvhandler/tikv_handler.go
@@ -72,9 +72,11 @@ import (
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"github.com/tikv/client-go/v2/tikv"
+	pd "github.com/tikv/pd/client"
 	"github.com/tikv/pd/client/clients/router"
-	pd "github.com/tikv/pd/client/http"
+	pdhttp "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/client/opt"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/zap"
 )
 
@@ -144,12 +146,17 @@ func NewRegionHandler(tool *handler.TikvHandlerTool) *RegionHandler {
 // TableHandler is the handler for list table's regions.
 type TableHandler struct {
 	*handler.TikvHandlerTool
-	op string
+	pdClient pd.Client
+	op       string
 }
 
 // NewTableHandler creates a new TableHandler.
 func NewTableHandler(tool *handler.TikvHandlerTool, op string) *TableHandler {
-	return &TableHandler{tool, op}
+	return &TableHandler{
+		TikvHandlerTool: tool,
+		pdClient:        tool.RegionCache.PDClient().WithCallerComponent(caller.TikvHandler),
+		op:              op,
+	}
 }
 
 // DDLHistoryJobHandler is the handler for list job history.
@@ -1293,7 +1300,7 @@ func (h *TableHandler) addScatterSchedule(startKey, endKey []byte, name string)
 	if err != nil {
 		return err
 	}
-	scheduleURL := fmt.Sprintf("%s://%s%s", util.InternalHTTPSchema(), pdAddrs[0], pd.Schedulers)
+	scheduleURL := fmt.Sprintf("%s://%s%s", util.InternalHTTPSchema(), pdAddrs[0], pdhttp.Schedulers)
 	resp, err := util.InternalHTTPClient().Post(scheduleURL, "application/json", bytes.NewBuffer(v))
 	if err != nil {
 		return err
@@ -1309,7 +1316,7 @@ func (h *TableHandler) deleteScatterSchedule(name string) error {
 	if err != nil {
 		return err
 	}
-	scheduleURL := fmt.Sprintf("%s://%s%s", util.InternalHTTPSchema(), pdAddrs[0], pd.ScatterRangeSchedulerWithName(name))
+	scheduleURL := fmt.Sprintf("%s://%s%s", util.InternalHTTPSchema(), pdAddrs[0], pdhttp.ScatterRangeSchedulerWithName(name))
 	req, err := http.NewRequest(http.MethodDelete, scheduleURL, nil)
 	if err != nil {
 		return err
@@ -1445,7 +1452,7 @@ func (h *TableHandler) getRegionsByID(tbl table.Table, id int64, name string) (*
 	// for record
 	startKey, endKey := tablecodec.GetTableHandleKeyRange(id)
 	ctx := context.Background()
-	pdCli := h.RegionCache.PDClient()
+	pdCli := h.pdClient
 	regions, err := pdCli.BatchScanRegions(ctx, []router.KeyRange{{StartKey: startKey, EndKey: endKey}}, -1, opt.WithAllowFollowerHandle())
 	if err != nil {
 		return nil, err
diff --git a/pkg/store/gcworker/BUILD.bazel b/pkg/store/gcworker/BUILD.bazel
@@ -41,6 +41,7 @@ go_library(
         "@com_github_tikv_pd_client//:client",
         "@com_github_tikv_pd_client//clients/gc",
         "@com_github_tikv_pd_client//constants",
+        "@com_github_tikv_pd_client//pkg/caller",
         "@org_uber_go_zap//:zap",
     ],
 )
diff --git a/pkg/store/gcworker/gc_worker.go b/pkg/store/gcworker/gc_worker.go
@@ -60,6 +60,7 @@ import (
 	pd "github.com/tikv/pd/client"
 	pdgc "github.com/tikv/pd/client/clients/gc"
 	"github.com/tikv/pd/client/constants"
+	"github.com/tikv/pd/client/pkg/caller"
 	"go.uber.org/zap"
 )
 
@@ -96,14 +97,15 @@ func NewGCWorker(store kv.Storage, pdClient pd.Client) (*GCWorker, error) {
 	uuid := strconv.FormatUint(ver.Ver, 16)
 	resolverIdentifier := fmt.Sprintf("gc-worker-%s", uuid)
 	keyspaceID := uint32(store.GetCodec().GetKeyspaceID())
+	gcPDClient := pdClient.WithCallerComponent(caller.GcWorker)
 	worker := &GCWorker{
 		uuid:                 uuid,
 		desc:                 fmt.Sprintf("host:%s, pid:%d, start at %s", hostName, os.Getpid(), time.Now()),
 		keyspaceID:           keyspaceID,
 		store:                store,
 		tikvStore:            tikvStore,
-		pdClient:             pdClient,
-		pdGCControllerClient: pdClient.GetGCInternalController(keyspaceID),
+		pdClient:             gcPDClient,
+		pdGCControllerClient: gcPDClient.GetGCInternalController(keyspaceID),
 		gcIsRunning:          false,
 		lastFinish:           time.Now(),
 		regionLockResolver:   tikv.NewRegionLockResolver(resolverIdentifier, tikvStore),
@@ -1800,12 +1802,13 @@ func RunGCJob(ctx context.Context, regionLockResolver tikv.RegionLockResolver, s
 // This function may not finish immediately because it may take some time to do resolveLocks.
 // Param concurrency specifies the concurrency of resolveLocks phase.
 func RunDistributedGCJob(ctx context.Context, regionLockResolver tikv.RegionLockResolver, s tikv.Storage, pd pd.Client, safePoint uint64, identifier string, concurrency int) error {
+	gcPDClient := pd.WithCallerComponent(caller.DistributedGcJob)
 	gcWorker := &GCWorker{
 		tikvStore:            s,
 		uuid:                 identifier,
 		keyspaceID:           constants.NullKeyspaceID,
-		pdClient:             pd,
-		pdGCControllerClient: pd.GetGCInternalController(constants.NullKeyspaceID),
+		pdClient:             gcPDClient,
+		pdGCControllerClient: gcPDClient.GetGCInternalController(constants.NullKeyspaceID),
 		regionLockResolver:   regionLockResolver,
 	}
 
diff --git a/tests/realtikvtest/addindextest3/functional_test.go b/tests/realtikvtest/addindextest3/functional_test.go
@@ -86,7 +86,7 @@ func TestDDLTestEstimateTableRowSize(t *testing.T) {
 }
 
 func TestBackendCtxConcurrentUnregister(t *testing.T) {
-	store := realtikvtest.CreateMockStoreAndSetup(t)
+	store, _ := realtikvtest.CreateMockStoreAndDomainAndSetup(t)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test;")
 	tk.MustExec("create table t (a int);")

Original file line number	Diff line number	Diff line change
`@@ -159,7 +159,7 @@ func NewPdController(`
`159`	`159`	`grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(maxMsgSize)),`
`160`	`160`	`}`
`161`	`161`	`pdClient, err := pd.NewClientWithContext(`
`162`		`- ctx, caller.Component("br-pd-controller"), pdAddrs, securityOption,`
	`162`	`+ ctx, caller.GetComponent(1), pdAddrs, securityOption,`
`163`	`163`	`opt.WithGRPCDialOptions(maxCallMsgSize...),`
`164`	`164`	`// If the time too short, we may scatter a region many times, because`
`165`	`165`	// the interface `ScatterRegions` may time out.