diff --git a/internal/locate/pd_codec.go b/internal/locate/pd_codec.go index 6cf75b595..90efe6c62 100644 --- a/internal/locate/pd_codec.go +++ b/internal/locate/pd_codec.go @@ -56,7 +56,7 @@ type CodecPDClient struct { // NewCodecPDClient creates a CodecPDClient in API v1. func NewCodecPDClient(mode apicodec.Mode, client pd.Client) *CodecPDClient { codec := apicodec.NewCodecV1(mode) - return &CodecPDClient{client, codec} + return &CodecPDClient{client.WithCallerComponent("CodecPDClient").(pd.Client), codec} } // NewCodecPDClientWithKeyspace creates a CodecPDClient in API v2 with keyspace name. diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go index 6fc55db5d..e8203b925 100644 --- a/internal/locate/region_cache.go +++ b/internal/locate/region_cache.go @@ -667,7 +667,7 @@ func NewRegionCache(pdClient pd.Client, opt ...RegionCacheOpt) *RegionCache { } c := &RegionCache{ - pdClient: pdClient, + pdClient: pdClient.WithCallerComponent("RegionCache").(pd.Client), requestHealthFeedbackCallback: options.requestHealthFeedbackCallback, } diff --git a/tikv/kv.go b/tikv/kv.go index 165e64883..a046338f4 100644 --- a/tikv/kv.go +++ b/tikv/kv.go @@ -905,7 +905,7 @@ func NewLockResolver(etcdAddrs []string, security config.Security, opts ...opt.C if err != nil { return nil, errors.WithStack(err) } - pdCli = util.InterceptedPDClient{Client: pdCli} + pdCli = util.NewInterceptedPDClient(pdCli) uuid := fmt.Sprintf("tikv-%v", pdCli.GetClusterID(context.TODO())) tlsConfig, err := security.ToTLSConfig() diff --git a/txnkv/client.go b/txnkv/client.go index aa3733359..bfcaad47a 100644 --- a/txnkv/client.go +++ b/txnkv/client.go @@ -76,7 +76,7 @@ func NewClient(pdAddrs []string, opts ...ClientOpt) (*Client, error) { return nil, errors.WithStack(err) } - pdClient = util.InterceptedPDClient{Client: pdClient} + pdClient = util.NewInterceptedPDClient(pdClient) // Construct codec from options. var codecCli *tikv.CodecPDClient diff --git a/util/pd_interceptor.go b/util/pd_interceptor.go index b8918d216..c66549ad8 100644 --- a/util/pd_interceptor.go +++ b/util/pd_interceptor.go @@ -62,6 +62,10 @@ type InterceptedPDClient struct { pd.Client } +func NewInterceptedPDClient(client pd.Client) *InterceptedPDClient { + return &InterceptedPDClient{client.WithCallerComponent("InterceptedPDClient").(pd.Client)} +} + // interceptedTsFuture is a PD's wrapper future to record stmt detail. type interceptedTsFuture struct { pd.TSFuture @@ -123,7 +127,7 @@ func (m InterceptedPDClient) GetRegionByID(ctx context.Context, regionID uint64, func (m InterceptedPDClient) ScanRegions(ctx context.Context, key, endKey []byte, limit int, opts ...opt.GetRegionOption) ([]*pd.Region, error) { start := time.Now() //nolint:staticcheck - r, err := m.Client.ScanRegions(ctx, key, endKey, limit, opts...) + r, err := m.Client.WithCallerComponent("InterceptedPDClient").ScanRegions(ctx, key, endKey, limit, opts...) recordPDWaitTime(ctx, start) return r, err }