add explicit request source type to label the external request like l…

…ightning/br (#868)

Signed-off-by: nolouch <[email protected]>

tikvrpc/interceptor/interceptor.go

@@ -205,6 +205,9 @@ var interceptorCtxKey = interceptorCtxKeyType{}
 
 // WithRPCInterceptor is a helper function used to bind RPCInterceptor with ctx.
 func WithRPCInterceptor(ctx context.Context, interceptor RPCInterceptor) context.Context {
+	if v := ctx.Value(interceptorCtxKey); v != nil {
+		interceptor = ChainRPCInterceptors(v.(RPCInterceptor), interceptor)
+	}
 	return context.WithValue(ctx, interceptorCtxKey, interceptor)
 }
 

txnkv/transaction/txn.go

@@ -1467,3 +1467,8 @@ func (txn *KVTxn) SetRequestSourceInternal(internal bool) {
 func (txn *KVTxn) SetRequestSourceType(tp string) {
 	txn.RequestSource.SetRequestSourceType(tp)
 }
+
+// SetExplicitRequestSourceType sets the explicit type of the request source.
+func (txn *KVTxn) SetExplicitRequestSourceType(tp string) {
+	txn.RequestSource.SetExplicitRequestSourceType(tp)
+}

util/request_source.go

@@ -27,11 +27,24 @@ const (
 	InternalTxnMeta = InternalTxnOthers
 )
 
+// explicit source types.
+const (
+	ExplicitTypeEmpty      = ""
+	ExplicitTypeDefault    = "default"
+	ExplicitTypeLightning  = "lightning"
+	ExplicitTypeBR         = "br"
+	ExplicitTypeDumpling   = "dumpling"
+	ExplicitTypeBackground = "background"
+)
+
+// ExplicitTypeList is the list of all explicit source types.
+var ExplicitTypeList = []string{ExplicitTypeEmpty, ExplicitTypeDefault, ExplicitTypeLightning, ExplicitTypeBR, ExplicitTypeDumpling, ExplicitTypeBackground}
+
 const (
 	// InternalRequest is the scope of internal queries
-	InternalRequest = "internal_"
+	InternalRequest = "internal"
 	// ExternalRequest is the scope of external queries
-	ExternalRequest = "external_"
+	ExternalRequest = "external"
 	// SourceUnknown keeps same with the default value(empty string)
 	SourceUnknown = "unknown"
 )
@@ -40,6 +53,9 @@ const (
 type RequestSource struct {
 	RequestSourceInternal bool
 	RequestSourceType     string
+	// ExplicitRequestSourceType is a type that is set from the session variable and may be specified by the client or users.
+	// It is a complement to the RequestSourceType and provides additional information about how a request was initiated.
+	ExplicitRequestSourceType string
 }
 
 // SetRequestSourceInternal sets the scope of the request source.
@@ -52,6 +68,11 @@ func (r *RequestSource) SetRequestSourceType(tp string) {
 	r.RequestSourceType = tp
 }
 
+// SetExplicitRequestSourceType sets the type of the request source.
+func (r *RequestSource) SetExplicitRequestSourceType(tp string) {
+	r.ExplicitRequestSourceType = tp
+}
+
 // WithInternalSourceType create context with internal source.
 func WithInternalSourceType(ctx context.Context, source string) context.Context {
 	return context.WithValue(ctx, RequestSourceKey, RequestSource{
@@ -71,15 +92,25 @@ func IsRequestSourceInternal(reqSrc *RequestSource) bool {
 
 // GetRequestSource gets the request_source field of the request.
 func (r *RequestSource) GetRequestSource() string {
-	// if r.RequestSourceType is not set, it's mostly possible that r.RequestSourceInternal is not set
-	// to avoid internal requests be marked as external(default value), return unknown source here.
-	if r == nil || r.RequestSourceType == "" {
-		return SourceUnknown
+	source := SourceUnknown
+	explicitSourceType := ExplicitTypeDefault
+	origin := ExternalRequest
+	if r == nil || (len(r.RequestSourceType) == 0 && len(r.ExplicitRequestSourceType) == 0) {
+		// if r.RequestSourceType and r.ExplicitRequestSourceType are not set, it's mostly possible that r.RequestSourceInternal is not set
+		// to avoid internal requests be marked as external(default value), return unknown source here.
+		return strings.Join([]string{source, explicitSourceType}, "_")
+	}
+
+	if len(r.RequestSourceType) > 0 {
+		source = r.RequestSourceType
+	}
+	if len(r.ExplicitRequestSourceType) > 0 {
+		explicitSourceType = r.ExplicitRequestSourceType
 	}
 	if r.RequestSourceInternal {
-		return InternalRequest + r.RequestSourceType
+		origin = InternalRequest
 	}
-	return ExternalRequest + r.RequestSourceType
+	return strings.Join([]string{origin, source, explicitSourceType}, "_")
 }
 
 // RequestSourceFromCtx extract source from passed context.

util/request_source_test.go

@@ -0,0 +1,68 @@
+// Copyright 2023 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package util
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetRequestSource(t *testing.T) {
+	rsi := true
+	rst := "test"
+	ers := "lightning"
+	rs := &RequestSource{
+		RequestSourceInternal:     rsi,
+		RequestSourceType:         rst,
+		ExplicitRequestSourceType: ers,
+	}
+
+	// Test internal request
+	expected := "internal_test_lightning"
+	actual := rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+
+	// Test external request
+	rs.RequestSourceInternal = false
+	expected = "external_test_lightning"
+	actual = rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+
+	// Test nil pointer
+	rs = nil
+	expected = "unknown_default"
+	actual = rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+
+	// Test empty RequestSourceType and ExplicitRequestSourceType
+	rs = &RequestSource{}
+	expected = "unknown_default"
+	actual = rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+
+	// Test empty ExplicitRequestSourceType
+	rs.RequestSourceType = "test"
+	expected = "external_test_default"
+	actual = rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+
+	// Test empty RequestSourceType
+	rs.RequestSourceType = ""
+	rs.ExplicitRequestSourceType = "lightning"
+	expected = "external_unknown_lightning"
+	actual = rs.GetRequestSource()
+	assert.Equal(t, expected, actual)
+}