Let GPU samples be cached and fix bogus frame 0 references

Author: gnurizen

interpreter/gpu/cuda.go

@@ -34,11 +34,12 @@ var (
 )
 
 // SymbolizedCudaTrace holds a symbolized trace awaiting GPU timing information.
-// The CPU frames are already symbolized; only the CUDA kernel frame (frame[0])
+// The CPU frames are already symbolized; only the CUDA kernel frame
 // needs the kernel name from the timing event.
 type SymbolizedCudaTrace struct {
 	Trace         *libpf.Trace
 	Meta          *samples.TraceEventMeta
+	CUDAFrameIdx  int // index of CUDAKernelFrame in Trace.Frames
 	CorrelationID uint32
 	CBID          int32
 }
@@ -60,9 +61,9 @@ type CudaTraceOutput struct {
 // that launched the kernel."
 type gpuTraceFixer struct {
 	mu                  sync.Mutex
-	timesAwaitingTraces map[uint32][]CuptiTimingEvent    // keyed by correlation ID
-	tracesAwaitingTimes map[uint32]*SymbolizedCudaTrace  // keyed by correlation ID
-	maxCorrelationId    uint32                           // track highest ID for threshold-based clearing
+	timesAwaitingTraces map[uint32][]CuptiTimingEvent   // keyed by correlation ID
+	tracesAwaitingTimes map[uint32]*SymbolizedCudaTrace // keyed by correlation ID
+	maxCorrelationId    uint32                          // track highest ID for threshold-based clearing
 }
 
 type data struct {
@@ -343,7 +344,7 @@ func (f *gpuTraceFixer) prepTrace(st *SymbolizedCudaTrace, ev *CuptiTimingEvent)
 		out.Trace.CustomLabels["cuda_id"] = strconv.FormatUint(uint64(ev.Id), 10)
 	}
 
-	// Extract kernel name from timing event, demangle, and update frame[0]
+	// Extract kernel name from timing event, demangle, and update the CUDA frame.
 	nameBytes := ev.KernelName[:]
 	if idx := bytes.IndexByte(nameBytes, 0); idx >= 0 {
 		nameBytes = nameBytes[:idx]
@@ -356,11 +357,10 @@ func (f *gpuTraceFixer) prepTrace(st *SymbolizedCudaTrace, ev *CuptiTimingEvent)
 			funcName = libpf.Intern(demStr)
 		}
 
-		currentFrame := out.Trace.Frames[0].Value()
-		out.Trace.Frames[0] = unique.Make(libpf.Frame{
-			Type:            currentFrame.Type,
-			AddressOrLineno: currentFrame.AddressOrLineno,
-			FunctionName:    funcName,
+		fi := st.CUDAFrameIdx
+		out.Trace.Frames[fi] = unique.Make(libpf.Frame{
+			Type:         out.Trace.Frames[fi].Value().Type,
+			FunctionName: funcName,
 		})
 	}
 
@@ -375,7 +375,7 @@ func AddTrace(st *SymbolizedCudaTrace) []CudaTraceOutput {
 	pid := st.Meta.PID
 	value, ok := gpuFixers.Load(pid)
 	if !ok {
-		log.Warnf("no GPU fixer found for PID %d", pid)
+		log.Warnf("no GPU fixer found for PID %d in AddTrace", pid)
 		return nil
 	}
 	fixer := value.(*gpuTraceFixer)
@@ -387,7 +387,7 @@ func addTimeSingle(ev *CuptiTimingEvent) (CudaTraceOutput, bool) {
 	pid := libpf.PID(ev.Pid)
 	value, ok := gpuFixers.Load(pid)
 	if !ok {
-		log.Warnf("no GPU fixer found for PID %d", pid)
+		log.Warnf("no GPU fixer found for PID %d in AddTime", pid)
 		return CudaTraceOutput{}, false
 	}
 	fixer := value.(*gpuTraceFixer)
@@ -409,6 +409,7 @@ func AddTimes(events []CuptiTimingEvent) []CudaTraceOutput {
 	pid := libpf.PID(events[0].Pid)
 	value, ok := gpuFixers.Load(pid)
 	if !ok {
+		log.Warnf("no GPU fixer found for PID %d in AddTimes", pid)
 		return nil
 	}
 	fixer := value.(*gpuTraceFixer)
@@ -448,6 +449,7 @@ func MaybeClearAll() []metrics.Metric {
 		totalTraces += stats.tracesLen
 		totalTimesCleared += stats.timesCleared
 		totalTracesCleared += stats.tracesCleared
+
 		return true
 	})
 

parcagpu/parcagpu.go

@@ -73,6 +73,7 @@ func Start(ctx context.Context, tr *tracer.Tracer,
 				// avoiding duplicate-metric warnings from the metrics system.
 				metrics.AddSlice(gpu.MaybeClearAll())
 			case <-ctx.Done():
+				eventReader.Close()
 				return
 			default:
 				if err := eventReader.ReadInto(&data); err != nil {
@@ -105,7 +106,19 @@ func Start(ctx context.Context, tr *tracer.Tracer,
 			return false
 		}
 
-		// Find the CUDA kernel frame in the symbolized trace
+		// Extract correlation ID and CBID from the raw BPF trace (not the
+		// symbolized trace, which may be a cached template with stale values).
+		var correlationID uint32
+		var cbid int32
+		for i := range rawTrace.Frames {
+			if rawTrace.Frames[i].Type == libpf.CUDAKernelFrame {
+				correlationID = uint32(rawTrace.Frames[i].Lineno)
+				cbid = int32(rawTrace.Frames[i].Lineno >> 32)
+				break
+			}
+		}
+
+		// Find the CUDA kernel frame index in the symbolized trace.
 		cudaFrameIdx := -1
 		for i, uniqueFrame := range trace.Frames {
 			if uniqueFrame.Value().Type == libpf.CUDAKernelFrame {
@@ -118,13 +131,10 @@ func Start(ctx context.Context, tr *tracer.Tracer,
 			return false
 		}
 
-		frame := trace.Frames[cudaFrameIdx].Value()
-		correlationID := uint32(frame.AddressOrLineno)
-		cbid := int32(frame.AddressOrLineno >> 32)
-
 		st := &gpu.SymbolizedCudaTrace{
 			Trace:         trace,
 			Meta:          meta,
+			CUDAFrameIdx:  cudaFrameIdx,
 			CorrelationID: correlationID,
 			CBID:          cbid,
 		}

tracehandler/tracehandler.go

@@ -146,7 +146,10 @@ func (m *traceHandler) HandleTrace(bpfTrace *host.Trace) {
 	if trace, exists := m.traceCache.GetAndRefresh(bpfTrace.Hash,
 		traceCacheLifetime); exists {
 		m.traceCacheHit++
-		// Fast path
+		// Fast path: interceptor gets cached template + fresh rawTrace.
+		if m.interceptor != nil && m.interceptor(&trace, meta, bpfTrace) {
+			return
+		}
 		meta.APMServiceName = m.traceProcessor.MaybeNotifyAPMAgent(bpfTrace, trace.Hash, 1)
 		if err := m.reporter.ReportTraceEvent(&trace, meta); err != nil {
 			log.Errorf("Failed to report trace event: %v", err)
@@ -163,14 +166,13 @@ func (m *traceHandler) HandleTrace(bpfTrace *host.Trace) {
 	}
 	log.Debugf("Trace hash remap 0x%x -> 0x%x", bpfTrace.Hash, umTrace.Hash)
 
-	// If an interceptor is set and consumes the trace, skip caching and reporting.
-	// CUDA traces are always intercepted here and never cached.
+	m.traceCache.Add(bpfTrace.Hash, *umTrace)
+
+	// If an interceptor consumes the trace, skip reporting.
 	if m.interceptor != nil && m.interceptor(umTrace, meta, bpfTrace) {
 		return
 	}
 
-	m.traceCache.Add(bpfTrace.Hash, *umTrace)
-
 	meta.APMServiceName = m.traceProcessor.MaybeNotifyAPMAgent(bpfTrace, umTrace.Hash, 1)
 	if err := m.reporter.ReportTraceEvent(umTrace, meta); err != nil {
 		log.Errorf("Failed to report trace event: %v", err)

tracehandler/tracehandler_test.go

@@ -110,15 +110,15 @@ func TestTraceInterceptor(t *testing.T) {
 				libpf.NewTraceHash(2, 2): 1,
 			},
 		},
-		"consumed trace not cached on repeat": {
+		"consumed on both cache miss and hit": {
 			interceptReturn: map[host.TraceHash]bool{
 				host.TraceHash(0xcc): true,
 			},
 			input: []arguments{
 				{trace: &host.Trace{Hash: host.TraceHash(0xcc)}},
 				{trace: &host.Trace{Hash: host.TraceHash(0xcc)}},
 			},
-			expectedEvents: nil, // both consumed, nothing reported
+			expectedEvents: nil, // first miss + second hit, both consumed
 		},
 	}
 

tracer/tracer.go

@@ -934,19 +934,6 @@ func (t *Tracer) loadBpfTrace(raw []byte, cpu int) *host.Trace {
 		}
 	}
 
-	// Trace fields included in the hash:
-	//  - PID, kernel stack ID, length & frame array
-	// Intentionally excluded:
-	//  - ktime, COMM, APM trace, APM transaction ID, Origin and Off Time
-	ptr.Comm = [16]byte{}
-	ptr.Apm_trace_id = support.ApmTraceID{}
-	ptr.Apm_transaction_id = support.ApmSpanID{}
-	ptr.Ktime = 0
-	ptr.Origin = 0
-	ptr.Offtime = 0
-	ptr.Custom_labels = support.CustomLabelsArray{}
-	trace.Hash = host.TraceHash(xxh3.Hash128(raw).Lo)
-
 	if ptr.Kernel_stack_id >= 0 {
 		var err error
 		trace.KernelFrames, err = t.readKernelFrames(ptr.Kernel_stack_id)
@@ -955,6 +942,9 @@ func (t *Tracer) loadBpfTrace(raw []byte, cpu int) *host.Trace {
 		}
 	}
 
+	// Build host.Trace frames BEFORE zeroing per-sample fields, because
+	// ZeroPerSampleFields clears the CUDA frame's Addr_or_line which carries
+	// the correlation ID needed by the GPU interceptor.
 	trace.Frames = make([]host.Frame, ptr.Stack_len)
 	for i := 0; i < int(ptr.Stack_len); i++ {
 		rawFrame := &ptr.Frames[i]
@@ -967,6 +957,15 @@ func (t *Tracer) loadBpfTrace(raw []byte, cpu int) *host.Trace {
 			LJCallerPC:    uint32(rawFrame.Caller_pc_lo) + (uint32(rawFrame.Caller_pc_hi) << 16),
 		}
 	}
+
+	// Trace fields included in the hash:
+	//  - PID, kernel stack ID, length & frame array
+	// Intentionally excluded:
+	//  - ktime, COMM, APM trace, APM transaction ID, Origin and Off Time
+	//  - CUDA frame addr_or_line (correlation ID is per-sample, not trace identity)
+	ptr.ZeroPerSampleFields()
+	trace.Hash = host.TraceHash(xxh3.Hash128(raw).Lo)
+
 	return trace
 }