*: refactor traceevent to make some concept clear

pkg/ddl/executor.go

@@ -6801,16 +6801,22 @@ func (e *executor) doDDLJob2(ctx sessionctx.Context, job *model.Job, args model.
 // When fast create is enabled, we might merge multiple jobs into one, so do not
 // depend on job.ID, use JobID from jobSubmitResult.
 func (e *executor) DoDDLJobWrapper(ctx sessionctx.Context, jobW *JobWrapper) (resErr error) {
+	var traceID []byte
 	if traceCtx := ctx.GetTraceCtx(); traceCtx != nil {
 		r := tracing.StartRegion(traceCtx, "ddl.DoDDLJobWrapper")
 		defer r.End()
+		if v := tracing.GetTraceBuf(traceCtx); v != nil {
+			if traceBuf, ok := v.(*traceevent.TraceBuf); ok && traceBuf != nil {
+				traceID = traceBuf.GetTraceID()
+			}
+		}
 	}
 
 	job := jobW.Job
 	job.TraceInfo = &tracing.TraceInfo{
 		ConnectionID: ctx.GetSessionVars().ConnectionID,
 		SessionAlias: ctx.GetSessionVars().SessionAlias,
-		TraceID:      traceevent.TraceIDFromContext(ctx.GetTraceCtx()),
+		TraceID:      traceID,
 	}
 	if mci := ctx.GetSessionVars().StmtCtx.MultiSchemaInfo; mci != nil {
 		// In multiple schema change, we don't run the job.

pkg/ddl/job_scheduler.go

@@ -310,8 +310,8 @@ func (s *jobScheduler) schedule() error {
 	defer ticker.Stop()
 	s.mustReloadSchemas()
 
-	trace := traceevent.NewTrace()
-	ctx := tracing.WithFlightRecorder(s.schCtx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx := traceevent.WithTraceBuf(s.schCtx, trace)
 
 	for {
 		if err := s.schCtx.Err(); err != nil {
@@ -525,7 +525,7 @@ func (s *jobScheduler) deliveryJob(ctx context.Context, wk *worker, pool *worker
 			pool.put(wk)
 		}()
 
-		trace := traceevent.NewTrace()
+		trace := traceevent.NewTraceBuf()
 		jobCtx := s.getJobRunCtx(trace, jobW.ID, jobW.TraceInfo)
 		defer trace.DiscardOrFlush(jobCtx.ctx)
 
@@ -565,11 +565,11 @@ func (s *jobScheduler) deliveryJob(ctx context.Context, wk *worker, pool *worker
 	})
 }
 
-func (s *jobScheduler) getJobRunCtx(trace *traceevent.Trace, jobID int64, traceInfo *tracing.TraceInfo) *jobContext {
+func (s *jobScheduler) getJobRunCtx(trace *traceevent.TraceBuf, jobID int64, traceInfo *tracing.TraceInfo) *jobContext {
 	ch, _ := s.ddlJobDoneChMap.Load(jobID)
-	newCtx := tracing.WithFlightRecorder(s.schCtx, trace)
+	newCtx := traceevent.WithTraceBuf(s.schCtx, trace)
 	if len(traceInfo.TraceID) > 0 {
-		newCtx = traceevent.ContextWithTraceID(newCtx, traceInfo.TraceID)
+		trace.SetTraceID(traceInfo.TraceID)
 	}
 	return &jobContext{
 		ctx:                  newCtx,

pkg/ddl/notifier/subscribe.go

@@ -155,8 +155,8 @@ func (n *DDLNotifier) start() {
 
 	ctx := kv.WithInternalSourceType(n.ctx, kv.InternalDDLNotifier)
 	ctx = logutil.WithCategory(ctx, "ddl-notifier")
-	trace := traceevent.NewTrace()
-	ctx = tracing.WithFlightRecorder(ctx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx = traceevent.WithTraceBuf(ctx, trace)
 	ticker := time.NewTicker(n.pollInterval)
 	defer ticker.Stop()
 	for {

pkg/dxf/framework/scheduler/nodes.go

@@ -123,8 +123,8 @@ func (nm *NodeManager) maintainLiveNodes(ctx context.Context, taskMgr TaskManage
 func (nm *NodeManager) refreshNodesLoop(ctx context.Context, taskMgr TaskManager, slotMgr *SlotManager) {
 	ticker := time.NewTicker(nodesCheckInterval)
 	defer ticker.Stop()
-	trace := traceevent.NewTrace()
-	ctx = tracing.WithFlightRecorder(ctx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx = tracing.WithTraceBuf(ctx, trace)
 	for {
 		select {
 		case <-ctx.Done():

pkg/dxf/framework/scheduler/scheduler_manager.go

@@ -214,8 +214,8 @@ func (sm *Manager) scheduleTaskLoop() {
 	sm.logger.Info("schedule task loop start")
 	ticker := time.NewTicker(CheckTaskRunningInterval)
 	defer ticker.Stop()
-	trace := traceevent.NewTrace()
-	ctx := tracing.WithFlightRecorder(sm.ctx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx := traceevent.WithTraceBuf(sm.ctx, trace)
 	for {
 		select {
 		case <-sm.ctx.Done():
@@ -491,8 +491,8 @@ func (sm *Manager) collectLoop() {
 	defer func() {
 		metrics.Unregister(sm.metricCollector)
 	}()
-	trace := traceevent.NewTrace()
-	ctx := tracing.WithFlightRecorder(sm.ctx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx := traceevent.WithTraceBuf(sm.ctx, trace)
 	for {
 		select {
 		case <-sm.ctx.Done():

pkg/dxf/framework/taskexecutor/manager.go

@@ -72,7 +72,7 @@ type Manager struct {
 	logger       *zap.Logger
 	slotManager  *slotManager
 	nodeResource *proto.NodeResource
-	trace        *traceevent.Trace
+	trace        *traceevent.TraceBuf
 }
 
 // NewManager creates a new task executor Manager.
@@ -89,10 +89,10 @@ func NewManager(ctx context.Context, store kv.Storage, id string, taskTable Task
 		logger:       logger,
 		slotManager:  newSlotManager(resource.TotalCPU),
 		nodeResource: resource,
-		trace:        traceevent.NewTrace(),
+		trace:        traceevent.NewTraceBuf(),
 	}
 
-	ctx = tracing.WithFlightRecorder(ctx, m.trace)
+	ctx = traceevent.WithTraceBuf(ctx, m.trace)
 	m.ctx, m.cancel = context.WithCancel(ctx)
 	m.mu.taskExecutors = make(map[int64]TaskExecutor)
 

pkg/dxf/framework/taskexecutor/task_executor.go

@@ -43,7 +43,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/traceevent"
-	"github.com/pingcap/tidb/pkg/util/tracing"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 )
@@ -273,8 +272,8 @@ func (e *BaseTaskExecutor) Run() {
 		metering.UnregisterRecorder(e.GetTaskBase().ID)
 	}()
 
-	trace := traceevent.NewTrace()
-	ctx := tracing.WithFlightRecorder(e.ctx, trace)
+	trace := traceevent.NewTraceBuf()
+	ctx := traceevent.WithTraceBuf(e.ctx, trace)
 	// task executor occupies resources, if there's no subtask to run for 10s,
 	// we release the resources so that other tasks can use them.
 	// 300ms + 600ms + 1.2s + 2s * 4 = 10.1s

pkg/server/conn.go

@@ -1080,8 +1080,8 @@ func (cc *clientConn) Run(ctx context.Context) {
 	cc.addConnMetrics()
 
 	var traceInfo *tracing.TraceInfo
-	trace := traceevent.NewTrace()
-	ctx = tracing.WithFlightRecorder(ctx, trace)
+	traceBuf := traceevent.NewTraceBuf()
+	ctx = traceevent.WithTraceBuf(ctx, traceBuf)
 
 	// Usually, client connection status changes between [dispatching] <=> [reading].
 	// When some event happens, server may notify this client connection by setting
@@ -1176,7 +1176,7 @@ func (cc *clientConn) Run(ctx context.Context) {
 		err = cc.dispatch(ctx, data)
 		cc.ctx.GetSessionVars().ClearAlloc(&cc.chunkAlloc, err != nil)
 		cc.chunkAlloc.Reset()
-		trace.DiscardOrFlush(ctx)
+		traceBuf.DiscardOrFlush(ctx)
 
 		if err != nil {
 			cc.audit(context.Background(), plugin.Error) // tell the plugin API there was a dispatch error

pkg/session/session.go

@@ -1849,9 +1849,9 @@ func (s *session) getOomAlarmVariablesInfo() sessmgr.OOMAlarmVariablesInfo {
 }
 
 func (s *session) ExecuteInternal(ctx context.Context, sql string, args ...any) (rs sqlexec.RecordSet, err error) {
-	if sink := tracing.GetSink(ctx); sink == nil {
-		trace := traceevent.NewTrace()
-		ctx = tracing.WithFlightRecorder(ctx, trace)
+	if sink := tracing.GetTraceBuf(ctx); sink == nil {
+		trace := traceevent.NewTraceBuf()
+		ctx = traceevent.WithTraceBuf(ctx, trace)
 		defer trace.DiscardOrFlush(ctx)
 
 		// A developer debugging event so we can see what trace is missing!
@@ -2822,7 +2822,10 @@ func resetStmtTraceID(ctx context.Context, se *session) (context.Context, []byte
 	// The trace ID is generated from transaction start_ts and statement count
 	startTS := se.sessionVars.TxnCtx.StartTS
 	stmtCount := uint64(se.sessionVars.TxnCtx.StatementCount)
-	traceID := traceevent.GenerateTraceID(ctx, startTS, stmtCount)
+	traceID := traceevent.GenerateTraceID(startTS, stmtCount)
+	if traceBuf := traceevent.GetTraceBuf(ctx); traceBuf != nil {
+		traceBuf.SetTraceID(traceID)
+	}
 	ctx = trace.ContextWithTraceID(ctx, traceID)
 	se.currentCtx = ctx
 	// Store trace ID for next statement

pkg/util/traceevent/adapter.go

@@ -57,7 +57,7 @@ func handleTraceControlExtractor(ctx context.Context) trace.TraceControlFlags {
 	flags := trace.TraceControlFlags(0)
 
 	// Map TiDB categories to TiKV categories regardless of whether a Trace sink is present.
-	enabledCategories := tracing.GetEnabledCategories()
+	enabledCategories := GetEnabledCategories()
 	if enabledCategories&tracing.TiKVRequest != 0 {
 		flags = flags.With(trace.FlagTiKVCategoryRequest)
 	}
@@ -68,13 +68,8 @@ func handleTraceControlExtractor(ctx context.Context) trace.TraceControlFlags {
 		flags = flags.With(trace.FlagTiKVCategoryReadDetails)
 	}
 
-	// Extract Trace object from context
-	sink := tracing.GetSink(ctx)
-	if sink == nil {
-		return flags
-	}
-	t, ok := sink.(*Trace)
-	if !ok {
+	t := GetTraceBuf(ctx)
+	if t == nil {
 		return flags
 	}
 

pkg/util/traceevent/adapter_test.go

@@ -33,9 +33,9 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test with nil context (no sink)
 	t.Run("NoSink", func(t *testing.T) {
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
-		tracing.SetCategories(tracing.TiKVRequest)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
+		fr.SetCategories(tracing.TiKVRequest)
 
 		ctx := context.Background()
 		flags := handleTraceControlExtractor(ctx)
@@ -45,15 +45,15 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test with keep=false
 	t.Run("KeepFalse", func(t *testing.T) {
-		tr := NewTrace()
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		tr := NewTraceBuf()
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
 		// Enable only TiKVRequest
-		tracing.SetCategories(tracing.TiKVRequest)
+		fr.SetCategories(tracing.TiKVRequest)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.False(t, flags.Has(trace.FlagImmediateLog), "immediate log should not be set when keep=false")
@@ -62,17 +62,17 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test with keep=true
 	t.Run("KeepTrue", func(t *testing.T) {
-		tr := NewTrace()
+		tr := NewTraceBuf()
 		// This sets keep=true
 		tr.bits = GetFlightRecorder().truthTable[0]
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
 		// Enable only TiKVRequest
-		tracing.SetCategories(tracing.TiKVRequest)
+		fr.SetCategories(tracing.TiKVRequest)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.True(t, flags.Has(trace.FlagImmediateLog), "immediate log should be set when keep=true")
@@ -81,14 +81,14 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test category mapping: TiKVRequest
 	t.Run("CategoryTiKVRequest", func(t *testing.T) {
-		tr := NewTrace()
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		tr := NewTraceBuf()
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
-		tracing.SetCategories(tracing.TiKVRequest)
+		fr.SetCategories(tracing.TiKVRequest)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.True(t, flags.Has(trace.FlagTiKVCategoryRequest))
@@ -98,14 +98,14 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test category mapping: TiKVWriteDetails
 	t.Run("CategoryTiKVWriteDetails", func(t *testing.T) {
-		tr := NewTrace()
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		tr := NewTraceBuf()
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
-		tracing.SetCategories(tracing.TiKVWriteDetails)
+		fr.SetCategories(tracing.TiKVWriteDetails)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.False(t, flags.Has(trace.FlagTiKVCategoryRequest))
@@ -115,14 +115,14 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test category mapping: TiKVReadDetails
 	t.Run("CategoryTiKVReadDetails", func(t *testing.T) {
-		tr := NewTrace()
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		tr := NewTraceBuf()
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
-		tracing.SetCategories(tracing.TiKVReadDetails)
+		fr.SetCategories(tracing.TiKVReadDetails)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.False(t, flags.Has(trace.FlagTiKVCategoryRequest))
@@ -132,17 +132,17 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test multiple categories
 	t.Run("MultipleCategoriesAndKeep", func(t *testing.T) {
-		tr := NewTrace()
+		tr := NewTraceBuf()
 		// Set keep=true
 		tr.bits = GetFlightRecorder().truthTable[0]
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
 		// Enable all three TiKV categories
-		tracing.SetCategories(tracing.TiKVRequest | tracing.TiKVWriteDetails | tracing.TiKVReadDetails)
+		fr.SetCategories(tracing.TiKVRequest | tracing.TiKVWriteDetails | tracing.TiKVReadDetails)
 
 		flags := handleTraceControlExtractor(ctx)
 		require.True(t, flags.Has(trace.FlagImmediateLog), "immediate log should be set")
@@ -153,14 +153,14 @@ func TestTraceControlExtractor(t *testing.T) {
 
 	// Test concurrent access (should not race)
 	t.Run("ConcurrentAccess", func(t *testing.T) {
-		tr := NewTrace()
-		ctx := tracing.WithFlightRecorder(context.Background(), tr)
+		tr := NewTraceBuf()
+		ctx := tracing.WithTraceBuf(context.Background(), tr)
 
 		// Save old categories and restore after test
-		oldCategories := tracing.GetEnabledCategories()
-		defer tracing.SetCategories(oldCategories)
+		oldCategories := GetEnabledCategories()
+		defer fr.SetCategories(oldCategories)
 
-		tracing.SetCategories(tracing.TiKVRequest)
+		fr.SetCategories(tracing.TiKVRequest)
 
 		var wg sync.WaitGroup
 		// Run multiple concurrent extractors