mirror of https://github.com/tikv/client-go.git
114 lines
3.8 KiB
Go
114 lines
3.8 KiB
Go
// 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 (
|
|
"context"
|
|
"sync/atomic"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
// Category identifies a trace event family emitted from client-go.
|
|
type Category uint32
|
|
|
|
const (
|
|
// CategoryTxn2PC traces two-phase commit prewrite and commit phases.
|
|
CategoryTxn2PC Category = iota
|
|
// CategoryTxnLockResolve traces lock resolution and conflict handling.
|
|
CategoryTxnLockResolve
|
|
// CategoryKVRequest traces individual KV request send and result events.
|
|
CategoryKVRequest
|
|
)
|
|
|
|
// TraceEventFunc is the function signature for recording trace events.
|
|
type TraceEventFunc func(ctx context.Context, category Category, name string, fields ...zap.Field)
|
|
|
|
// IsCategoryEnabledFunc is the function signature for checking if a category is enabled.
|
|
type IsCategoryEnabledFunc func(category Category) bool
|
|
|
|
// Default no-op implementations
|
|
func noopTraceEvent(context.Context, Category, string, ...zap.Field) {}
|
|
func noopIsCategoryEnabled(Category) bool { return false }
|
|
|
|
// Global function pointers stored independently
|
|
var (
|
|
globalTraceEventFunc atomic.Pointer[TraceEventFunc]
|
|
globalIsCategoryEnabledFunc atomic.Pointer[IsCategoryEnabledFunc]
|
|
)
|
|
|
|
func init() {
|
|
// Set default no-op implementations
|
|
defaultTraceEvent := TraceEventFunc(noopTraceEvent)
|
|
defaultIsCategoryEnabled := IsCategoryEnabledFunc(noopIsCategoryEnabled)
|
|
globalTraceEventFunc.Store(&defaultTraceEvent)
|
|
globalIsCategoryEnabledFunc.Store(&defaultIsCategoryEnabled)
|
|
}
|
|
|
|
// SetTraceEventFunc registers the trace event handler function.
|
|
// This is typically called once during application initialization (e.g., by TiDB).
|
|
// Passing nil will use a no-op implementation.
|
|
func SetTraceEventFunc(fn TraceEventFunc) {
|
|
if fn == nil {
|
|
fn = noopTraceEvent
|
|
}
|
|
globalTraceEventFunc.Store(&fn)
|
|
}
|
|
|
|
// SetIsCategoryEnabledFunc registers the category enablement check function.
|
|
// This can be updated independently of the trace event function.
|
|
// Passing nil will use a no-op implementation that returns false.
|
|
func SetIsCategoryEnabledFunc(fn IsCategoryEnabledFunc) {
|
|
if fn == nil {
|
|
fn = noopIsCategoryEnabled
|
|
}
|
|
globalIsCategoryEnabledFunc.Store(&fn)
|
|
}
|
|
|
|
// TraceEvent records a trace event using the registered function.
|
|
func TraceEvent(ctx context.Context, category Category, name string, fields ...zap.Field) {
|
|
fn := globalTraceEventFunc.Load()
|
|
(*fn)(ctx, category, name, fields...)
|
|
}
|
|
|
|
// IsCategoryEnabled checks if a category is enabled for tracing using the registered function.
|
|
func IsCategoryEnabled(category Category) bool {
|
|
fn := globalIsCategoryEnabledFunc.Load()
|
|
return (*fn)(category)
|
|
}
|
|
|
|
// Trace ID context management
|
|
|
|
// traceIDKey is the context key for storing trace IDs.
|
|
// This key is shared between TiDB and client-go for trace ID propagation.
|
|
type traceIDKey struct{}
|
|
|
|
// ContextWithTraceID returns a new context with the given trace ID attached.
|
|
func ContextWithTraceID(ctx context.Context, traceID []byte) context.Context {
|
|
return context.WithValue(ctx, traceIDKey{}, traceID)
|
|
}
|
|
|
|
// TraceIDFromContext extracts the trace ID from the context.
|
|
// Returns nil if no trace ID is present.
|
|
func TraceIDFromContext(ctx context.Context) []byte {
|
|
if ctx == nil {
|
|
return nil
|
|
}
|
|
if traceID, ok := ctx.Value(traceIDKey{}).([]byte); ok {
|
|
return traceID
|
|
}
|
|
return nil
|
|
}
|