Add TraceInterceptor to tracehandler

Add a TraceInterceptor callback that is invoked after ConvertTrace on
cache-miss. When the interceptor returns true the trace is consumed
(skipped for caching and reporting), allowing callers like the GPU
subsystem to divert specific traces for further processing.

Includes tests covering consume, pass-through, mixed, and
non-caching behavior.

Author: gnurizen

internal/controller/controller.go

@@ -181,7 +181,7 @@ func startTraceHandling(ctx context.Context, rep reporter.TraceReporter,
 	}
 
 	_, err := tracehandler.Start(ctx, rep, trc.TraceProcessor(),
-		traceCh, intervals, cacheSize)
+		traceCh, intervals, cacheSize, nil)
 	return err
 }
 

tracehandler/tracehandler.go

@@ -34,6 +34,12 @@ type Times interface {
 // symbolization efforts.
 var traceCacheLifetime = 5 * time.Minute
 
+// TraceInterceptor is called after ConvertTrace with the symbolized trace,
+// metadata, and original BPF trace. Return true to consume the trace
+// (skip caching and reporting), false to proceed normally.
+type TraceInterceptor func(trace *libpf.Trace, meta *samples.TraceEventMeta,
+	rawTrace *host.Trace) bool
+
 // TraceProcessor is an interface used by traceHandler to convert traces
 // from a form received from eBPF to the form we wish to dispatch to the
 // collection agent.
@@ -72,12 +78,17 @@ type traceHandler struct {
 	// reporter instance to use to send out traces.
 	reporter reporter.TraceReporter
 
+	// interceptor, if set, is called after ConvertTrace on cache-miss.
+	// If it returns true the trace is consumed and not cached or reported.
+	interceptor TraceInterceptor
+
 	times Times
 }
 
 // newTraceHandler creates a new traceHandler
 func newTraceHandler(ctx context.Context, rep reporter.TraceReporter,
-	traceProcessor TraceProcessor, intervals Times, cacheSize uint32) (*traceHandler, error) {
+	traceProcessor TraceProcessor, intervals Times, cacheSize uint32,
+	interceptor TraceInterceptor) (*traceHandler, error) {
 	traceCache, err := lru.NewSynced[host.TraceHash, libpf.Trace](
 		cacheSize, func(k host.TraceHash) uint32 { return uint32(k) })
 	if err != nil {
@@ -111,6 +122,7 @@ func newTraceHandler(ctx context.Context, rep reporter.TraceReporter,
 		traceProcessor: traceProcessor,
 		traceCache:     traceCache,
 		reporter:       rep,
+		interceptor:    interceptor,
 		times:          intervals,
 	}, nil
 }
@@ -150,6 +162,13 @@ func (m *traceHandler) HandleTrace(bpfTrace *host.Trace) {
 		panic(err)
 	}
 	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.
+	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)
@@ -164,9 +183,10 @@ func (m *traceHandler) HandleTrace(bpfTrace *host.Trace) {
 // to exit after a cancellation through the context.
 func Start(ctx context.Context, rep reporter.TraceReporter, traceProcessor TraceProcessor,
 	traceInChan <-chan *host.Trace, intervals Times, cacheSize uint32,
+	interceptor TraceInterceptor,
 ) (workerExited <-chan libpf.Void, err error) {
 	handler, err :=
-		newTraceHandler(ctx, rep, traceProcessor, intervals, cacheSize)
+		newTraceHandler(ctx, rep, traceProcessor, intervals, cacheSize, interceptor)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create traceHandler: %v", err)
 	}

tracehandler/tracehandler_test.go

@@ -67,6 +67,100 @@ func (m *mockReporter) ReportTraceEvent(trace *libpf.Trace, _ *samples.TraceEven
 	return nil
 }
 
+func TestTraceInterceptor(t *testing.T) {
+	tests := map[string]struct {
+		// interceptReturn controls what the interceptor returns for each trace.
+		// true = consumed (should NOT appear in reporter), false = pass through.
+		interceptReturn map[host.TraceHash]bool
+		input           []arguments
+		expectedEvents  map[libpf.TraceHash]uint16
+	}{
+		"interceptor consumes trace": {
+			interceptReturn: map[host.TraceHash]bool{
+				host.TraceHash(0xaa): true,
+			},
+			input: []arguments{
+				{trace: &host.Trace{Hash: host.TraceHash(0xaa)}},
+			},
+			expectedEvents: nil, // consumed, nothing reported
+		},
+		"interceptor passes through": {
+			interceptReturn: map[host.TraceHash]bool{
+				host.TraceHash(0xbb): false,
+			},
+			input: []arguments{
+				{trace: &host.Trace{Hash: host.TraceHash(0xbb)}},
+			},
+			expectedEvents: map[libpf.TraceHash]uint16{
+				libpf.NewTraceHash(0xbb, 0xbb): 1,
+			},
+		},
+		"mix of consumed and passed": {
+			interceptReturn: map[host.TraceHash]bool{
+				host.TraceHash(1): true,  // consumed
+				host.TraceHash(2): false, // passed
+				host.TraceHash(3): true,  // consumed
+			},
+			input: []arguments{
+				{trace: &host.Trace{Hash: host.TraceHash(1)}},
+				{trace: &host.Trace{Hash: host.TraceHash(2)}},
+				{trace: &host.Trace{Hash: host.TraceHash(3)}},
+			},
+			expectedEvents: map[libpf.TraceHash]uint16{
+				libpf.NewTraceHash(2, 2): 1,
+			},
+		},
+		"consumed trace not cached on repeat": {
+			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
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			r := &mockReporter{
+				t:       t,
+				reports: make(map[libpf.TraceHash]uint16),
+			}
+
+			var interceptCalls int
+			interceptor := func(trace *libpf.Trace, meta *samples.TraceEventMeta,
+				rawTrace *host.Trace) bool {
+				interceptCalls++
+				return test.interceptReturn[rawTrace.Hash]
+			}
+
+			traceChan := make(chan *host.Trace)
+			ctx, cancel := context.WithCancel(t.Context())
+			defer cancel()
+			exitNotify, err := tracehandler.Start(ctx, r, &fakeTraceProcessor{},
+				traceChan, defaultTimes(), 128, interceptor)
+			require.NoError(t, err)
+
+			for _, input := range test.input {
+				traceChan <- input.trace
+			}
+
+			cancel()
+			<-exitNotify
+
+			expected := test.expectedEvents
+			if expected == nil {
+				expected = map[libpf.TraceHash]uint16{}
+			}
+			if !maps.Equal(r.reports, expected) {
+				t.Fatalf("Expected %#v but got %#v", expected, r.reports)
+			}
+		})
+	}
+}
+
 func TestTraceHandler(t *testing.T) {
 	tests := map[string]struct {
 		input          []arguments
@@ -108,7 +202,7 @@ func TestTraceHandler(t *testing.T) {
 			ctx, cancel := context.WithCancel(t.Context())
 			defer cancel()
 			exitNotify, err := tracehandler.Start(ctx, r, &fakeTraceProcessor{},
-				traceChan, defaultTimes(), 128)
+				traceChan, defaultTimes(), 128, nil)
 			require.NoError(t, err)
 
 			for _, input := range test.input {