From 178f6fa01aab1d0f674b2cc031d4769b1d67a680 Mon Sep 17 00:00:00 2001 From: ShuNing Date: Fri, 7 Jul 2023 15:02:42 +0800 Subject: [PATCH] add explicit request source type to label the external request like lightning/br (#868) Signed-off-by: nolouch --- tikvrpc/interceptor/interceptor.go | 3 ++ txnkv/transaction/txn.go | 5 +++ util/request_source.go | 47 +++++++++++++++++---- util/request_source_test.go | 68 ++++++++++++++++++++++++++++++ 4 files changed, 115 insertions(+), 8 deletions(-) create mode 100644 util/request_source_test.go diff --git a/tikvrpc/interceptor/interceptor.go b/tikvrpc/interceptor/interceptor.go index 3c232fdab..0daf20aee 100644 --- a/tikvrpc/interceptor/interceptor.go +++ b/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) } diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index f9873a33f..a12027aa6 100644 --- a/txnkv/transaction/txn.go +++ b/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) +} diff --git a/util/request_source.go b/util/request_source.go index eba974773..2723062d5 100644 --- a/util/request_source.go +++ b/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. diff --git a/util/request_source_test.go b/util/request_source_test.go new file mode 100644 index 000000000..4f991a9b0 --- /dev/null +++ b/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) +}