tracing: pass immediate logging down to kv request (#1785)

ekexium · web-flow · commit 1264c1278595 · 2025-11-12T11:31:23.000Z
ref pingcap/tidb#64008 Signed-off-by: ekexium <eke@fastmail.com>
diff --git a/examples/gcworker/go.mod b/examples/gcworker/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/rawkv/go.mod b/examples/rawkv/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/1pc_txn/go.mod b/examples/txnkv/1pc_txn/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/async_commit/go.mod b/examples/txnkv/async_commit/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/delete_range/go.mod b/examples/txnkv/delete_range/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/go.mod b/examples/txnkv/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/pessimistic_txn/go.mod b/examples/txnkv/pessimistic_txn/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/examples/txnkv/unsafedestoryrange/go.mod b/examples/txnkv/unsafedestoryrange/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122 // indirect
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 // indirect
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
diff --git a/go.mod b/go.mod
@@ -15,7 +15,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20241219054535-6b8c588c3122
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.20.5
diff --git a/go.sum b/go.sum
@@ -77,8 +77,8 @@ github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXH
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 h1:wAL6q8MTf0kL6EGFribwUhMnJW3YovdDUslFgcQ1iWU=
-github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 h1:ZgebBNwgma8INCfexAX8dfqZo7TWQrvMXHGABEmAY2Y=
+github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/integration_tests/go.mod b/integration_tests/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20250523034308-74f78ae071ee
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23
+	github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18
 	github.com/pingcap/tidb v1.1.0-beta.0.20250609033843-a165d9fd7c01
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.10.0
diff --git a/integration_tests/go.sum b/integration_tests/go.sum
@@ -1351,8 +1351,8 @@ github.com/pingcap/fn v1.0.0/go.mod h1:u9WZ1ZiOD1RpNhcI42RucFh/lBuzTu6rw88a+oF2Z
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20241113043844-e1fa7ea8c302/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
-github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23 h1:wAL6q8MTf0kL6EGFribwUhMnJW3YovdDUslFgcQ1iWU=
-github.com/pingcap/kvproto v0.0.0-20251023055424-e9d10f5dcd23/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18 h1:ZgebBNwgma8INCfexAX8dfqZo7TWQrvMXHGABEmAY2Y=
+github.com/pingcap/kvproto v0.0.0-20251109100001-1907922fbd18/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20250514022801-14f3b4ca066e h1:8AZZRv1Ox9FVGATVZBBgr6y1MrNBQFABEdovNJt1QIc=
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
@@ -999,6 +999,8 @@ func (s *sendReqState) next() (done bool) {
 	if traceID := trace.TraceIDFromContext(bo.GetCtx()); len(traceID) > 0 {
 		req.Context.TraceId = traceID
 	}
+	// Set trace control flags based on the extractor function
+	req.Context.TraceControlFlags = uint64(trace.GetTraceControlFlags(bo.GetCtx()))
 	req.Context.ClusterId = s.vars.rpcCtx.ClusterID
 	if req.InputRequestSource != "" && s.replicaSelector != nil {
 		patchRequestSource(req, s.replicaSelector.replicaType())
@@ -1320,6 +1322,8 @@ func (s *sendReqState) initForAsyncRequest() (ok bool) {
 	if traceID := trace.TraceIDFromContext(bo.GetCtx()); len(traceID) > 0 {
 		req.Context.TraceId = traceID
 	}
+	// Set trace control flags based on the extractor function
+	req.Context.TraceControlFlags = uint64(trace.GetTraceControlFlags(bo.GetCtx()))
 	req.Context.ClusterId = s.vars.rpcCtx.ClusterID
 	if req.InputRequestSource != "" && s.replicaSelector != nil {
 		patchRequestSource(req, s.replicaSelector.replicaType())
diff --git a/trace/flags.go b/trace/flags.go
@@ -0,0 +1,48 @@
+// Copyright 2025 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 trace
+
+// TraceControlFlags represents trace logging control flags.
+// Each bit controls a different aspect of trace logging behavior.
+type TraceControlFlags uint64
+
+const (
+	// FlagImmediateLog forces immediate logging without buffering.
+	// When set, TiKV will log trace events immediately instead of buffering them.
+	FlagImmediateLog TraceControlFlags = 1 << 0
+
+	// FlagTiKVCategoryRequest enables TiKV request category tracing.
+	// When set, TiKV will trace request processing.
+	FlagTiKVCategoryRequest TraceControlFlags = 1 << 1
+
+	// FlagTiKVCategoryWriteDetails enables TiKV detailed write operation tracing.
+	// When set, TiKV will trace detailed write operation steps.
+	FlagTiKVCategoryWriteDetails TraceControlFlags = 1 << 2
+
+	// FlagTiKVCategoryReadDetails enables TiKV detailed read operation tracing.
+	// When set, TiKV will trace detailed read operation steps.
+	FlagTiKVCategoryReadDetails TraceControlFlags = 1 << 3
+)
+
+// Has checks if the given flag is set.
+func (f TraceControlFlags) Has(flag TraceControlFlags) bool {
+	return f&flag != 0
+}
+
+// With returns a new flags value with the given flag set.
+// This method does not modify the original value.
+func (f TraceControlFlags) With(flag TraceControlFlags) TraceControlFlags {
+	return f | flag
+}
diff --git a/trace/flags_test.go b/trace/flags_test.go
@@ -0,0 +1,82 @@
+// Copyright 2025 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 trace
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTraceControlFlags_Has(t *testing.T) {
+	flags := FlagImmediateLog | FlagTiKVCategoryRequest
+
+	require.True(t, flags.Has(FlagImmediateLog))
+	require.True(t, flags.Has(FlagTiKVCategoryRequest))
+	require.False(t, flags.Has(FlagTiKVCategoryWriteDetails))
+	require.False(t, flags.Has(FlagTiKVCategoryReadDetails))
+
+	// Empty flags
+	var emptyFlags TraceControlFlags
+	require.False(t, emptyFlags.Has(FlagImmediateLog))
+}
+
+func TestTraceControlFlags_With(t *testing.T) {
+	var flags TraceControlFlags
+
+	// Set individual flags
+	flags = flags.With(FlagImmediateLog)
+	require.True(t, flags.Has(FlagImmediateLog))
+	require.False(t, flags.Has(FlagTiKVCategoryRequest))
+
+	flags = flags.With(FlagTiKVCategoryRequest)
+	require.True(t, flags.Has(FlagImmediateLog))
+	require.True(t, flags.Has(FlagTiKVCategoryRequest))
+
+	// Setting an already-set flag is idempotent
+	flags = flags.With(FlagImmediateLog)
+	require.True(t, flags.Has(FlagImmediateLog))
+}
+
+func TestTraceControlFlags_CombinedOperations(t *testing.T) {
+	// Build up flags step by step (fluent style)
+	flags := TraceControlFlags(0).
+		With(FlagImmediateLog).
+		With(FlagTiKVCategoryRequest).
+		With(FlagTiKVCategoryWriteDetails).
+		With(FlagTiKVCategoryReadDetails)
+
+	// All four flags should be set
+	require.True(t, flags.Has(FlagImmediateLog))
+	require.True(t, flags.Has(FlagTiKVCategoryRequest))
+	require.True(t, flags.Has(FlagTiKVCategoryWriteDetails))
+	require.True(t, flags.Has(FlagTiKVCategoryReadDetails))
+}
+
+func TestTraceControlFlags_BitValues(t *testing.T) {
+	// Verify the bit positions are correct
+	require.Equal(t, TraceControlFlags(1<<0), FlagImmediateLog)
+	require.Equal(t, TraceControlFlags(1<<1), FlagTiKVCategoryRequest)
+	require.Equal(t, TraceControlFlags(1<<2), FlagTiKVCategoryWriteDetails)
+	require.Equal(t, TraceControlFlags(1<<3), FlagTiKVCategoryReadDetails)
+
+	// Verify flags don't overlap
+	require.NotEqual(t, FlagImmediateLog, FlagTiKVCategoryRequest)
+	require.NotEqual(t, FlagTiKVCategoryWriteDetails, FlagTiKVCategoryReadDetails)
+
+	// Verify we can combine all flags
+	all := FlagImmediateLog | FlagTiKVCategoryRequest | FlagTiKVCategoryWriteDetails | FlagTiKVCategoryReadDetails
+	require.Equal(t, TraceControlFlags(0b1111), all)
+}
diff --git a/trace/trace.go b/trace/trace.go
@@ -41,22 +41,35 @@ type TraceEventFunc func(ctx context.Context, category Category, name string, fi
 // IsCategoryEnabledFunc is the function signature for checking if a category is enabled.
 type IsCategoryEnabledFunc func(category Category) bool
 
+// TraceControlExtractorFunc is the function signature for extracting trace control flags from context.
+// This is typically provided by TiDB to determine trace logging behavior.
+type TraceControlExtractorFunc func(ctx context.Context) TraceControlFlags
+
 // Default no-op implementations
 func noopTraceEvent(context.Context, Category, string, ...zap.Field) {}
 func noopIsCategoryEnabled(Category) bool                            { return false }
 
+// noopTraceControlExtractor returns the default trace control flags.
+// By default, we enable FlagTiKVCategoryRequest to allow TiKV request tracing.
+func noopTraceControlExtractor(context.Context) TraceControlFlags {
+	return FlagTiKVCategoryRequest
+}
+
 // Global function pointers stored independently
 var (
-	globalTraceEventFunc        atomic.Pointer[TraceEventFunc]
-	globalIsCategoryEnabledFunc atomic.Pointer[IsCategoryEnabledFunc]
+	globalTraceEventFunc            atomic.Pointer[TraceEventFunc]
+	globalIsCategoryEnabledFunc     atomic.Pointer[IsCategoryEnabledFunc]
+	globalTraceControlExtractorFunc atomic.Pointer[TraceControlExtractorFunc]
 )
 
 func init() {
 	// Set default no-op implementations
 	defaultTraceEvent := TraceEventFunc(noopTraceEvent)
 	defaultIsCategoryEnabled := IsCategoryEnabledFunc(noopIsCategoryEnabled)
+	defaultTraceControlExtractor := TraceControlExtractorFunc(noopTraceControlExtractor)
 	globalTraceEventFunc.Store(&defaultTraceEvent)
 	globalIsCategoryEnabledFunc.Store(&defaultIsCategoryEnabled)
+	globalTraceControlExtractorFunc.Store(&defaultTraceControlExtractor)
 }
 
 // SetTraceEventFunc registers the trace event handler function.
@@ -91,6 +104,28 @@ func IsCategoryEnabled(category Category) bool {
 	return (*fn)(category)
 }
 
+// SetTraceControlExtractor registers the trace control extractor function.
+// This is typically called once during application initialization (e.g., by TiDB).
+// Passing nil will use a no-op implementation that returns 0 (no flags set).
+func SetTraceControlExtractor(fn TraceControlExtractorFunc) {
+	if fn == nil {
+		fn = noopTraceControlExtractor
+	}
+	globalTraceControlExtractorFunc.Store(&fn)
+}
+
+// GetTraceControlFlags extracts trace control flags from the given context.
+func GetTraceControlFlags(ctx context.Context) TraceControlFlags {
+	fn := globalTraceControlExtractorFunc.Load()
+	return (*fn)(ctx)
+}
+
+// ImmediateLoggingEnabled checks if immediate logging is enabled for the given context.
+// This is a convenience helper that checks the FlagImmediateLog bit.
+func ImmediateLoggingEnabled(ctx context.Context) bool {
+	return GetTraceControlFlags(ctx).Has(FlagImmediateLog)
+}
+
 // Trace ID context management
 
 // traceIDKey is the context key for storing trace IDs.
diff --git a/trace/trace_test.go b/trace/trace_test.go
