[elasticapmintakereceiver] Propagate request metdata when include_metadata is set (#718)

* elasticapmintakereceiver: propagate request metdata when
include_metadata is set

* elasticapmintakereceiver: add metadata propagation test cases

Author: 1pkg

receiver/elasticapmintakereceiver/receiver.go

@@ -101,17 +101,18 @@ func newElasticAPMIntakeReceiver(fetcher agentCfgFetcherFactory, cfg *Config, se
 // Start runs an HTTP server for receiving data from Elastic APM agents.
 func (r *elasticAPMIntakeReceiver) Start(ctx context.Context, host component.Host) error {
 	ctx, r.cancelFn = context.WithCancel(ctx)
-	ecsCtx := withECSMappingMode(ctx)
-	if err := r.startHTTPServer(ecsCtx, host); err != nil {
-		return errors.Join(err, r.Shutdown(ecsCtx))
+	if err := r.startHTTPServer(ctx, host); err != nil {
+		return errors.Join(err, r.Shutdown(ctx))
 	}
 	return nil
 }
 
 func (r *elasticAPMIntakeReceiver) startHTTPServer(ctx context.Context, host component.Host) error {
 	httpMux := http.NewServeMux()
 
-	httpMux.HandleFunc(intakeV2EventsPath, r.newElasticAPMEventsHandler(ctx))
+	httpMux.HandleFunc(intakeV2EventsPath, r.newElasticAPMEventsHandler(func(req *http.Request) context.Context {
+		return withECSMappingMode(req.Context(), r.cfg.IncludeMetadata)
+	}))
 	httpMux.HandleFunc(agentConfigPath, r.newElasticAPMConfigsHandler(ctx, host))
 	// TODO rum v2, v3
 
@@ -157,7 +158,7 @@ func errorHandler(w http.ResponseWriter, r *http.Request, errMsg string, statusC
 	// TODO
 }
 
-func (r *elasticAPMIntakeReceiver) newElasticAPMEventsHandler(ctx context.Context) http.HandlerFunc {
+func (r *elasticAPMIntakeReceiver) newElasticAPMEventsHandler(ctxFunc func(*http.Request) context.Context) http.HandlerFunc {
 
 	var (
 		// TODO make semaphore size configurable and/or find a different way
@@ -185,7 +186,7 @@ func (r *elasticAPMIntakeReceiver) newElasticAPMEventsHandler(ctx context.Contex
 		baseEvent := &modelpb.APMEvent{}
 		baseEvent.Event = &modelpb.Event{}
 		streamErr := elasticapmProcessor.HandleStream(
-			ctx,
+			ctxFunc(req),
 			baseEvent,
 			req.Body,
 			batchSize,
@@ -460,14 +461,21 @@ func (r *elasticAPMIntakeReceiver) elasticSpanToOTelSpan(s *ptrace.Span, event *
 	}
 }
 
-func withECSMappingMode(ctx context.Context) context.Context {
-	return client.NewContext(ctx, withMappingMode(client.FromContext(ctx), "ecs"))
+func withECSMappingMode(ctx context.Context, includeMetadata bool) context.Context {
+	return client.NewContext(ctx, withMappingMode(client.FromContext(ctx), "ecs", includeMetadata))
 }
 
-func withMappingMode(info client.Info, mode string) client.Info {
+func withMappingMode(info client.Info, mode string, includeMetadata bool) client.Info {
+	newMeta := make(map[string][]string)
+	if includeMetadata {
+		for k := range info.Metadata.Keys() {
+			newMeta[k] = info.Metadata.Get(k)
+		}
+	}
+	newMeta["x-elastic-mapping-mode"] = []string{mode}
 	return client.Info{
 		Addr:     info.Addr,
 		Auth:     info.Auth,
-		Metadata: client.NewMetadata(map[string][]string{"x-elastic-mapping-mode": {mode}}),
+		Metadata: client.NewMetadata(newMeta),
 	}
 }

receiver/elasticapmintakereceiver/receiver_test.go

@@ -38,6 +38,7 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/pdatatest/ptracetest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/client"
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/component/componenttest"
 	"go.opentelemetry.io/collector/config/confighttp"
@@ -519,7 +520,65 @@ func TestTransactionsAndSpans(t *testing.T) {
 	}
 }
 
-func sendInput(t *testing.T, inputJsonFileName string, expectedYamlFileName string, testEndpoint string) {
+func TestMetadataPropagation(t *testing.T) {
+	table := map[string]struct {
+		includeMetadata  bool
+		expectedMetadata client.Metadata
+	}{
+		"when include_metadata is disabled only mappinmapping-mode is propagated": {
+			expectedMetadata: client.NewMetadata(map[string][]string{
+				"x-elastic-mapping-mode": {"ecs"},
+			}),
+		},
+		"when include_metadata is enabled all request metadata is propagated": {
+			includeMetadata: true,
+			expectedMetadata: client.NewMetadata(map[string][]string{
+				"content-type":           {"application/x-ndjson"},
+				"x-elastic-mapping-mode": {"ecs"},
+			}),
+		},
+	}
+	for tname, tcase := range table {
+		t.Run(tname, func(t *testing.T) {
+			factory := NewFactory()
+			testEndpoint := testutil.GetAvailableLocalAddress(t)
+			cfg := &Config{
+				ServerConfig: confighttp.ServerConfig{
+					Endpoint:        testEndpoint,
+					IncludeMetadata: tcase.includeMetadata,
+				},
+			}
+
+			set := receivertest.NewNopSettings(metadata.Type)
+			nextTrace := new(consumertest.TracesSink)
+			receiver, _ := factory.CreateTraces(context.Background(), set, cfg, nextTrace)
+
+			if err := receiver.Start(context.Background(), componenttest.NewNopHost()); err != nil {
+				t.Errorf("Starting receiver failed: %v", err)
+			}
+			defer func() {
+				if err := receiver.Shutdown(context.Background()); err != nil {
+					t.Errorf("Shutdown failed: %v", err)
+				}
+			}()
+
+			sendInput(t, "transactions_spans.ndjson", testEndpoint)
+
+			ctxs := nextTrace.Contexts()
+			require.GreaterOrEqual(t, len(ctxs), 1)
+			md := client.FromContext(ctxs[0]).Metadata
+			if tcase.includeMetadata {
+				for k := range tcase.expectedMetadata.Keys() {
+					require.Equal(t, tcase.expectedMetadata.Get(k), md.Get(k))
+				}
+			} else {
+				require.Equal(t, tcase.expectedMetadata, md)
+			}
+		})
+	}
+}
+
+func sendInput(t *testing.T, inputJsonFileName string, testEndpoint string) {
 	data, err := os.ReadFile(filepath.Join(testData, inputJsonFileName))
 	if err != nil {
 		t.Fatalf("failed to read file: %v", err)
@@ -542,7 +601,7 @@ func runComparisonForTraces(t *testing.T, inputJsonFileName string, expectedYaml
 ) {
 	nextTrace.Reset()
 
-	sendInput(t, inputJsonFileName, expectedYamlFileName, testEndpoint)
+	sendInput(t, inputJsonFileName, testEndpoint)
 	actualTraces := nextTrace.AllTraces()[0]
 	expectedFile := filepath.Join(testData, expectedYamlFileName)
 	// Use this line to generate the expected yaml file:
@@ -558,7 +617,7 @@ func runComparisonForErrors(t *testing.T, inputJsonFileName string, expectedYaml
 ) {
 	nextLog.Reset()
 
-	sendInput(t, inputJsonFileName, expectedYamlFileName, testEndpoint)
+	sendInput(t, inputJsonFileName, testEndpoint)
 	actualLogs := nextLog.AllLogs()[0]
 	expectedFile := filepath.Join(testData, expectedYamlFileName)
 	// Use this line to generate the expected yaml file:
@@ -573,7 +632,7 @@ func runComparisonForMetrics(t *testing.T, inputJsonFileName string, expectedYam
 ) {
 
 	nextMetric.Reset()
-	sendInput(t, inputJsonFileName, expectedYamlFileName, testEndpoint)
+	sendInput(t, inputJsonFileName, testEndpoint)
 	actualMetrics := nextMetric.AllMetrics()[0]
 	expectedFile := filepath.Join(testData, expectedYamlFileName)
 	// Use this line to generate the expected yaml file: