Pass Context through span processors (#6534)

## Which problem is this PR solving?
- Context is lost in the process, see
#6491 (comment)

## Description of the changes
- Add Context argument to handler and span processor methods

## How was this change tested?
- CI

---------

Signed-off-by: Yuri Shkuro <github@ysh.us>

Author: yurishkuro

cmd/collector/app/collector_test.go

@@ -263,7 +263,7 @@ func TestAggregator(t *testing.T) {
 			},
 		},
 	}
-	_, err := c.spanProcessor.ProcessSpans(processor.SpansV1{
+	_, err := c.spanProcessor.ProcessSpans(context.Background(), processor.SpansV1{
 		Spans: spans,
 		Details: processor.Details{
 			SpanFormat: processor.JaegerSpanFormat,

cmd/collector/app/handler/grpc_handler.go

@@ -75,7 +75,7 @@ func (c *batchConsumer) consume(ctx context.Context, batch *model.Batch) error {
 			span.Process = batch.Process
 		}
 	}
-	_, err = c.spanProcessor.ProcessSpans(processor.SpansV1{
+	_, err = c.spanProcessor.ProcessSpans(ctx, processor.SpansV1{
 		Spans: batch.Spans,
 		Details: processor.Details{
 			InboundTransport: c.spanOptions.InboundTransport,

cmd/collector/app/handler/grpc_handler_test.go

@@ -25,6 +25,8 @@ import (
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2"
 )
 
+var _ processor.SpanProcessor = (*mockSpanProcessor)(nil)
+
 type mockSpanProcessor struct {
 	expectedError error
 	mux           sync.Mutex
@@ -34,7 +36,7 @@ type mockSpanProcessor struct {
 	spanFormat    processor.SpanFormat
 }
 
-func (p *mockSpanProcessor) ProcessSpans(batch processor.Batch) ([]bool, error) {
+func (p *mockSpanProcessor) ProcessSpans(_ context.Context, batch processor.Batch) ([]bool, error) {
 	p.mux.Lock()
 	defer p.mux.Unlock()
 	batch.GetSpans(func(spans []*model.Span) {

cmd/collector/app/handler/http_thrift_handler.go

@@ -75,7 +75,7 @@ func (aH *APIHandler) SaveSpan(w http.ResponseWriter, r *http.Request) {
 	}
 	batches := []*tJaeger.Batch{batch}
 	opts := SubmitBatchOptions{InboundTransport: processor.HTTPTransport}
-	if _, err = aH.jaegerBatchesHandler.SubmitBatches(batches, opts); err != nil {
+	if _, err = aH.jaegerBatchesHandler.SubmitBatches(r.Context(), batches, opts); err != nil {
 		http.Error(w, fmt.Sprintf("Cannot submit Jaeger batch: %v", err), http.StatusInternalServerError)
 		return
 	}

cmd/collector/app/handler/http_thrift_handler_test.go

@@ -25,15 +25,18 @@ import (
 	"github.com/jaegertracing/jaeger/thrift-gen/jaeger"
 )
 
-var httpClient = &http.Client{Timeout: 2 * time.Second}
+var (
+	httpClient                      = &http.Client{Timeout: 2 * time.Second}
+	_          JaegerBatchesHandler = (*mockJaegerHandler)(nil)
+)
 
 type mockJaegerHandler struct {
 	err     error
 	mux     sync.Mutex
 	batches []*jaeger.Batch
 }
 
-func (p *mockJaegerHandler) SubmitBatches(batches []*jaeger.Batch, _ SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error) {
+func (p *mockJaegerHandler) SubmitBatches(_ context.Context, batches []*jaeger.Batch, _ SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error) {
 	p.mux.Lock()
 	defer p.mux.Unlock()
 	p.batches = append(p.batches, batches...)

cmd/collector/app/handler/thrift_span_handler.go

@@ -5,6 +5,8 @@
 package handler
 
 import (
+	"context"
+
 	"go.uber.org/zap"
 
 	"github.com/jaegertracing/jaeger/cmd/collector/app/processor"
@@ -24,13 +26,13 @@ type SubmitBatchOptions struct {
 // ZipkinSpansHandler consumes and handles zipkin spans
 type ZipkinSpansHandler interface {
 	// SubmitZipkinBatch records a batch of spans in Zipkin Thrift format
-	SubmitZipkinBatch(spans []*zipkincore.Span, options SubmitBatchOptions) ([]*zipkincore.Response, error)
+	SubmitZipkinBatch(ctx context.Context, spans []*zipkincore.Span, options SubmitBatchOptions) ([]*zipkincore.Response, error)
 }
 
 // JaegerBatchesHandler consumes and handles Jaeger batches
 type JaegerBatchesHandler interface {
 	// SubmitBatches records a batch of spans in Jaeger Thrift format
-	SubmitBatches(batches []*jaeger.Batch, options SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error)
+	SubmitBatches(ctx context.Context, batches []*jaeger.Batch, options SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error)
 }
 
 type jaegerBatchesHandler struct {
@@ -46,15 +48,15 @@ func NewJaegerSpanHandler(logger *zap.Logger, modelProcessor processor.SpanProce
 	}
 }
 
-func (jbh *jaegerBatchesHandler) SubmitBatches(batches []*jaeger.Batch, options SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error) {
+func (jbh *jaegerBatchesHandler) SubmitBatches(ctx context.Context, batches []*jaeger.Batch, options SubmitBatchOptions) ([]*jaeger.BatchSubmitResponse, error) {
 	responses := make([]*jaeger.BatchSubmitResponse, 0, len(batches))
 	for _, batch := range batches {
 		mSpans := make([]*model.Span, 0, len(batch.Spans))
 		for _, span := range batch.Spans {
 			mSpan := jConv.ToDomainSpan(span, batch.Process)
 			mSpans = append(mSpans, mSpan)
 		}
-		oks, err := jbh.modelProcessor.ProcessSpans(processor.SpansV1{
+		oks, err := jbh.modelProcessor.ProcessSpans(ctx, processor.SpansV1{
 			Spans: mSpans,
 			Details: processor.Details{
 				InboundTransport: options.InboundTransport,
@@ -98,7 +100,7 @@ func NewZipkinSpanHandler(logger *zap.Logger, modelHandler processor.SpanProcess
 }
 
 // SubmitZipkinBatch records a batch of spans already in Zipkin Thrift format.
-func (h *zipkinSpanHandler) SubmitZipkinBatch(spans []*zipkincore.Span, options SubmitBatchOptions) ([]*zipkincore.Response, error) {
+func (h *zipkinSpanHandler) SubmitZipkinBatch(ctx context.Context, spans []*zipkincore.Span, options SubmitBatchOptions) ([]*zipkincore.Response, error) {
 	mSpans := make([]*model.Span, 0, len(spans))
 	convCount := make([]int, len(spans))
 	for i, span := range spans {
@@ -108,7 +110,7 @@ func (h *zipkinSpanHandler) SubmitZipkinBatch(spans []*zipkincore.Span, options
 		convCount[i] = len(converted)
 		mSpans = append(mSpans, converted...)
 	}
-	bools, err := h.modelProcessor.ProcessSpans(processor.SpansV1{
+	bools, err := h.modelProcessor.ProcessSpans(ctx, processor.SpansV1{
 		Spans: mSpans,
 		Details: processor.Details{
 			InboundTransport: options.InboundTransport,

cmd/collector/app/handler/thrift_span_handler_test.go

@@ -5,6 +5,7 @@
 package handler
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"os"
@@ -36,7 +37,7 @@ func TestJaegerSpanHandler(t *testing.T) {
 	for _, tc := range testChunks {
 		logger := zap.NewNop()
 		h := NewJaegerSpanHandler(logger, &shouldIErrorProcessor{tc.expectedErr != nil})
-		res, err := h.SubmitBatches([]*jaeger.Batch{
+		res, err := h.SubmitBatches(context.Background(), []*jaeger.Batch{
 			{
 				Process: &jaeger.Process{ServiceName: "someServiceName"},
 				Spans:   []*jaeger.Span{{SpanId: 21345}},
@@ -57,9 +58,12 @@ type shouldIErrorProcessor struct {
 	shouldError bool
 }
 
-var errTestError = errors.New("Whoops")
+var (
+	_            processor.SpanProcessor = (*shouldIErrorProcessor)(nil)
+	errTestError                         = errors.New("Whoops")
+)
 
-func (s *shouldIErrorProcessor) ProcessSpans(batch processor.Batch) ([]bool, error) {
+func (s *shouldIErrorProcessor) ProcessSpans(_ context.Context, batch processor.Batch) ([]bool, error) {
 	if s.shouldError {
 		return nil, errTestError
 	}
@@ -121,7 +125,7 @@ func TestZipkinSpanHandler(t *testing.T) {
 					},
 				}
 			}
-			res, err := h.SubmitZipkinBatch(spans, SubmitBatchOptions{})
+			res, err := h.SubmitZipkinBatch(context.Background(), spans, SubmitBatchOptions{})
 			if tc.expectedErr != nil {
 				assert.Nil(t, res)
 				assert.Equal(t, tc.expectedErr, err)

cmd/collector/app/processor/processor.go

@@ -4,6 +4,7 @@
 package processor
 
 import (
+	"context"
 	"io"
 
 	"go.opentelemetry.io/collector/pdata/ptrace"
@@ -29,7 +30,7 @@ type Batch interface {
 // SpanProcessor handles spans
 type SpanProcessor interface {
 	// ProcessSpans processes spans and return with either a list of true/false success or an error
-	ProcessSpans(spans Batch) ([]bool, error)
+	ProcessSpans(ctx context.Context, spans Batch) ([]bool, error)
 	io.Closer
 }
 

cmd/collector/app/server/test.go

@@ -26,6 +26,6 @@ func (*mockSpanProcessor) Close() error {
 	return nil
 }
 
-func (*mockSpanProcessor) ProcessSpans(_ processor.Batch) ([]bool, error) {
+func (*mockSpanProcessor) ProcessSpans(_ context.Context, _ processor.Batch) ([]bool, error) {
 	return []bool{}, nil
 }

cmd/collector/app/span_processor.go

@@ -36,6 +36,8 @@ const (
 	minRequiredChange = 1.2
 )
 
+var _ processor.SpanProcessor = (*spanProcessor)(nil)
+
 type spanProcessor struct {
 	queue              *queue.BoundedQueue[queueItem]
 	otelExporter       exporter.Traces
@@ -239,19 +241,18 @@ func (sp *spanProcessor) countSpansInQueue(span *model.Span, _ string /* tenant
 	sp.spansProcessed.Add(1)
 }
 
-// TODO pass Context
-func (sp *spanProcessor) ProcessSpans(batch processor.Batch) ([]bool, error) {
+func (sp *spanProcessor) ProcessSpans(ctx context.Context, batch processor.Batch) ([]bool, error) {
 	// We call preProcessSpans on a batch, it's responsibility of implementation
 	// to understand v1/v2 distinction. Jaeger itself does not use pre-processors.
 	sp.preProcessSpans(batch)
 
 	var batchOks []bool
 	var batchErr error
 	batch.GetSpans(func(spans []*model.Span) {
-		batchOks, batchErr = sp.processSpans(batch, spans)
+		batchOks, batchErr = sp.processSpans(ctx, batch, spans)
 	}, func(traces ptrace.Traces) {
 		// TODO verify if the context will survive all the way to the consumer threads.
-		ctx := tenancy.WithTenant(context.Background(), batch.GetTenant())
+		ctx := tenancy.WithTenant(ctx, batch.GetTenant())
 
 		// the exporter will eventually call pushTraces from consumer threads.
 		if err := sp.otelExporter.ConsumeTraces(ctx, traces); err != nil {
@@ -266,7 +267,7 @@ func (sp *spanProcessor) ProcessSpans(batch processor.Batch) ([]bool, error) {
 	return batchOks, batchErr
 }
 
-func (sp *spanProcessor) processSpans(batch processor.Batch, spans []*model.Span) ([]bool, error) {
+func (sp *spanProcessor) processSpans(_ context.Context, batch processor.Batch, spans []*model.Span) ([]bool, error) {
 	sp.metrics.BatchSize.Update(int64(len(spans)))
 	retMe := make([]bool, len(spans))