add caller component to codec/regioncache/interceptor

Signed-off-by: okJiang <[email protected]>
tikv · Dec 10, 2024 · ce615a4 · ce615a4
1 parent 7183eff
commit ce615a4
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
 }