use HTTPHeadersCarrier to extract multi-value headers

Signed-off-by: SATO Taiki <[email protected]>
DataDog · Nov 12, 2024 · 270eb93 · 270eb93
1 parent 6eda679
commit 270eb93
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 57 deletions.
diff --git a/pkg/serverless/trace/propagation/carriers.go b/pkg/serverless/trace/propagation/carriers.go
@@ -253,16 +253,6 @@ func headersCarrier(hdrs map[string]string) (tracer.TextMapReader, error) {
 	return tracer.TextMapCarrier(hdrs), nil
 }
 
-// multiValueHeadersCarrier returns the tracer.TextMapReader used to extract trace
-// context from a MultiValueHeaders field of form map[string][]string.
-func multiValueHeadersCarrier(hdrs map[string][]string) (tracer.TextMapReader, error) {
-	headers := make(map[string]string)
-	for k, v := range hdrs {
-		headers[k] = v[0]
-	}
-	return headersCarrier(headers)
-}
-
 // extractTraceContextFromStepFunctionContext extracts the execution ARN, state name, and state entered time and uses them to generate Trace ID and Parent ID
 // The logic is based on the trace context conversion in Logs To Traces, dd-trace-py, dd-trace-js, etc.
 func extractTraceContextFromStepFunctionContext(event events.StepFunctionPayload) (*TraceContext, error) {

diff --git a/pkg/serverless/trace/propagation/carriers_test.go b/pkg/serverless/trace/propagation/carriers_test.go
@@ -826,51 +826,6 @@ func TestHeadersCarrier(t *testing.T) {
 	}
 }
 
-func TestMultiValueHeadersCarrier(t *testing.T) {
-	testcases := []struct {
-		name   string
-		event  map[string]string
-		expMap map[string]string
-		expErr error
-	}{
-		{
-			name:   "nil-map",
-			event:  headersMapNone,
-			expMap: headersMapEmpty,
-			expErr: nil,
-		},
-		{
-			name:   "empty-map",
-			event:  headersMapEmpty,
-			expMap: headersMapEmpty,
-			expErr: nil,
-		},
-		{
-			name:   "headers-all",
-			event:  headersMapAll,
-			expMap: headersMapAll,
-			expErr: nil,
-		},
-	}
-
-	for _, tc := range testcases {
-		t.Run(tc.name, func(t *testing.T) {
-			event := map[string][]string{}
-			for k, v := range tc.event {
-				event[k] = []string{v, "dummy"}
-			}
-
-			tm, err := multiValueHeadersCarrier(event)
-			t.Logf("rawPayloadCarrier returned TextMapReader=%#v error=%#v", tm, err)
-			assert.Equal(t, tc.expErr != nil, err != nil)
-			if tc.expErr != nil && err != nil {
-				assert.Equal(t, tc.expErr.Error(), err.Error())
-			}
-			assert.Equal(t, tc.expMap, getMapFromCarrier(tm))
-		})
-	}
-}
-
 func Test_stringToDdSpanId(t *testing.T) {
 	type args struct {
 		execArn          string

diff --git a/pkg/serverless/trace/propagation/extractor.go b/pkg/serverless/trace/propagation/extractor.go
@@ -115,7 +115,7 @@ func (e Extractor) extract(event interface{}) (*TraceContext, error) {
 		if len(ev.Headers) > 0 {
 			carrier, err = headersCarrier(ev.Headers)
 		} else {
-			carrier, err = multiValueHeadersCarrier(ev.MultiValueHeaders)
+			carrier = tracer.HTTPHeadersCarrier(ev.MultiValueHeaders)
 		}
 	case events.LambdaFunctionURLRequest:
 		carrier, err = headersCarrier(ev.Headers)

diff --git a/pkg/serverless/trace/propagation/extractor_test.go b/pkg/serverless/trace/propagation/extractor_test.go
@@ -179,7 +179,7 @@ var (
 func toMultiValueHeaders(headers map[string]string) map[string][]string {
 	mvh := make(map[string][]string)
 	for k, v := range headers {
-		mvh[k] = []string{v, "dummy"}
+		mvh[k] = []string{v}
 	}
 	return mvh
 }