fix(collector): propagate context to reportProfile for graceful shutdown

reporter/collector_reporter.go

@@ -57,7 +57,7 @@ func (r *CollectorReporter) Start(ctx context.Context) error {
 	ctx, cancelReporting := context.WithCancel(ctx)
 
 	r.runLoop.Start(ctx, r.cfg.ReportInterval, func() {
-		if err := r.reportProfile(context.Background()); err != nil {
+		if err := r.reportProfile(ctx); err != nil {
 			log.Errorf("Request failed: %v", err)
 		}
 	}, func() {

reporter/collector_reporter_test.go

@@ -3,14 +3,18 @@ package reporter
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/consumer/xconsumer"
 	"go.opentelemetry.io/collector/pdata/pprofile"
 
 	"go.opentelemetry.io/ebpf-profiler/libpf"
 	"go.opentelemetry.io/ebpf-profiler/reporter/samples"
+	"go.opentelemetry.io/ebpf-profiler/support"
 )
 
 func TestCollectorReporterReportTraceEvent(t *testing.T) {
@@ -60,3 +64,51 @@ func TestCollectorReporterReportTraceEvent(t *testing.T) {
 		})
 	}
 }
+
+func TestCollectorReporterShutdown(t *testing.T) {
+	var cancelled atomic.Bool
+	consumerStarted := make(chan struct{})
+	next, err := xconsumer.NewProfiles(func(ctx context.Context, _ pprofile.Profiles) error {
+		close(consumerStarted)
+		select {
+		case <-ctx.Done():
+			cancelled.Store(true)
+			return nil
+		}
+	})
+	require.NoError(t, err)
+
+	r, err := NewCollector(&Config{
+		ReportInterval: 10 * time.Millisecond,
+	}, next)
+	require.NoError(t, err)
+
+	traceEventsPtr := r.traceEvents.WLock()
+	tree := (*traceEventsPtr)
+	tree[libpf.NullString] = map[libpf.Origin]samples.KeyToEventMapping{
+		support.TraceOriginProbe: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+			{Pid: 1}: {
+				Frames: func() libpf.Frames {
+					frames := make(libpf.Frames, 0, 1)
+					frames.Append(&libpf.Frame{
+						Type:            libpf.KernelFrame,
+						AddressOrLineno: 0xef,
+						FunctionName:    libpf.Intern("func1"),
+					})
+					return frames
+				}(),
+				Timestamps: []uint64{1, 2, 3, 4},
+			},
+		},
+	}
+	r.traceEvents.WUnlock(&traceEventsPtr)
+
+	ctx, cancelFn := context.WithCancel(t.Context())
+	require.NoError(t, r.Start(ctx))
+	// BLOCK until the consumer is actually running
+	<-consumerStarted
+	cancelFn()
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
+		assert.True(collect, cancelled.Load())
+	}, 5*time.Second, 100*time.Millisecond, "consumer did not exit after context cancellation")
+}