Add baggage to request context (#131)

x/k6/baggage.go

@@ -24,13 +24,20 @@ type labelHandler struct {
 }
 
 func (lh *labelHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	var found bool
+	r, found = setBaggageContextFromHeader(r)
+	if !found {
+		lh.innerHandler.ServeHTTP(w, r)
+		return
+	}
+
 	labels := getBaggageLabels(r)
 	if labels == nil {
 		lh.innerHandler.ServeHTTP(w, r)
 		return
 	}
 
-	// Inlined version of pryoscope.TagWrapper and pprof.Do to reduce noise in
+	// Inlined version of pyroscope.TagWrapper and pprof.Do to reduce noise in
 	// the stack trace.
 	ctx := r.Context()
 	defer pprof.SetGoroutineLabels(ctx)
@@ -40,11 +47,26 @@ func (lh *labelHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	lh.innerHandler.ServeHTTP(w, r.WithContext(ctx))
 }
 
+func setBaggageContextFromHeader(r *http.Request) (*http.Request, bool) {
+	baggageHeader := r.Header.Get("Baggage")
+	if baggageHeader == "" {
+		return r, false
+	}
+
+	b, err := baggage.Parse(baggageHeader)
+	if err != nil {
+		return r, false
+	}
+
+	ctx := baggage.ContextWithBaggage(r.Context(), b)
+	return r.WithContext(ctx), true
+}
+
 // getBaggageLabels applies filters and transformations to request baggage and
 // returns the resulting LabelSet.
 func getBaggageLabels(r *http.Request) *pyroscope.LabelSet {
-	b, err := baggage.Parse(r.Header.Get("Baggage"))
-	if err != nil {
+	b := baggage.FromContext(r.Context())
+	if b.Len() == 0 {
 		return nil
 	}
 

x/k6/baggage_test.go

@@ -5,43 +5,124 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"runtime/pprof"
+	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/baggage"
 )
 
-func Test_getBaggageLabels(t *testing.T) {
-	t.Run("empty values are skipped", func(t *testing.T) {
+func TestLabelsFromBaggageHandler(t *testing.T) {
+	t.Run("adds_baggage_to_context", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "http://example.com", nil)
-		req = testRequestWithBaggage(t, req, map[string]string{
-			"blank": "",
-		})
+		req = testAddBaggageToRequest(t, req,
+			"k6.test_run_id", "123",
+			"not_k6.some_other_key", "value",
+		)
+
+		handler := LabelsFromBaggageHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			b := baggage.FromContext(r.Context())
+			require.NotNil(t, b)
+
+			testAssertEqualMembers(t, b.Members(),
+				"k6.test_run_id", "123",
+				"not_k6.some_other_key", "value",
+			)
+		}))
+
+		handler.ServeHTTP(httptest.NewRecorder(), req)
+	})
 
-		labelSet := getBaggageLabels(req)
-		gotLabels := testPprofLabelsToMap(t, *labelSet)
+	t.Run("passthrough_requests_with_no_baggage", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
 
-		expectedLabels := map[string]string{}
-		require.Equal(t, expectedLabels, gotLabels)
+		handler := LabelsFromBaggageHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			b := baggage.FromContext(r.Context())
+			require.Equal(t, 0, b.Len())
+		}))
+
+		handler.ServeHTTP(httptest.NewRecorder(), req)
 	})
+}
 
-	t.Run("with K6Options", func(t *testing.T) {
+func Test_setBaggageContextFromHeader(t *testing.T) {
+	t.Run("sets_baggage_context", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
+		req = testAddBaggageToRequest(t, req,
+			"k6.test_run_id", "123",
+			"not_k6.some_other_key", "value",
+		)
+
+		req, found := setBaggageContextFromHeader(req)
+		require.True(t, found)
+
+		b := baggage.FromContext(req.Context())
+		testAssertEqualMembers(t, b.Members(),
+			"k6.test_run_id", "123",
+
+			// Also passthrough non-k6 baggage. We should avoid clobbering
+			// other baggage members that the system might also be using.
+			"not_k6.some_other_key", "value",
+		)
+	})
+
+	t.Run("does_not_set_baggage_context_with_no_baggage", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "http://example.com", nil)
-		req = testRequestWithBaggage(t, req, map[string]string{
-			"k6.test_run_id":        "123",
-			"not_k6.some_other_key": "value",
-		})
+
+		req, found := setBaggageContextFromHeader(req)
+		require.False(t, found)
+
+		b := baggage.FromContext(req.Context())
+		require.Equal(t, 0, b.Len())
+	})
+
+	t.Run("does_not_set_baggage_context_invalid_baggage", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
+		req.Header.Add("Baggage", "invalid")
+
+		req, found := setBaggageContextFromHeader(req)
+		require.False(t, found)
+
+		b := baggage.FromContext(req.Context())
+		require.Equal(t, 0, b.Len())
+	})
+}
+
+func Test_getBaggageLabels(t *testing.T) {
+	t.Run("with_k6_baggage", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
+		req = testAddBaggageToRequest(t, req,
+			"k6.test_run_id", "123",
+			"not_k6.some_other_key", "value",
+		)
 
 		labelSet := getBaggageLabels(req)
-		gotLabels := testPprofLabelsToMap(t, *labelSet)
+		require.NotNil(t, labelSet)
 
+		gotLabels := testPprofLabelsToMap(t, *labelSet)
 		expectedLabels := map[string]string{
 			"k6_test_run_id": "123",
 		}
+
 		require.Equal(t, expectedLabels, gotLabels)
 	})
 
-	t.Run("does not allocate with failure to parse baggage", func(t *testing.T) {
+	t.Run("empty_values_are_skipped", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
+		req = testAddBaggageToRequest(t, req)
+
+		labelSet := getBaggageLabels(req)
+		require.Nil(t, labelSet)
+	})
+
+	t.Run("skips_missing_baggage_header", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://example.com", nil)
+
+		labelSet := getBaggageLabels(req)
+		require.Nil(t, labelSet)
+	})
+
+	t.Run("skips_invalid_baggage_header", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "http://example.com", nil)
 		req.Header.Add("Baggage", "invalid")
 
@@ -50,12 +131,17 @@ func Test_getBaggageLabels(t *testing.T) {
 	})
 }
 
-func testRequestWithBaggage(t *testing.T, req *http.Request, bag map[string]string) *http.Request {
+func testAddBaggageToRequest(t *testing.T, req *http.Request, kvPairs ...string) *http.Request {
 	t.Helper()
 
-	members := []baggage.Member{}
-	for k, v := range bag {
-		member, err := baggage.NewMember(k, v)
+	require.Equal(t, 0, len(kvPairs)%2, "kvPairs must be a multiple of 2")
+
+	members := make([]baggage.Member, 0, len(kvPairs)/2)
+	for i := 0; i < len(kvPairs); i += 2 {
+		key := kvPairs[i]
+		value := kvPairs[i+1]
+
+		member, err := baggage.NewMember(key, value)
 		require.NoError(t, err)
 
 		members = append(members, member)
@@ -64,10 +150,64 @@ func testRequestWithBaggage(t *testing.T, req *http.Request, bag map[string]stri
 	b, err := baggage.New(members...)
 	require.NoError(t, err)
 
+	ctx := baggage.ContextWithBaggage(req.Context(), b)
+	req = req.WithContext(ctx)
 	req.Header.Add("Baggage", b.String())
+
 	return req
 }
 
+func testMustNewMember(t *testing.T, key string, value string) baggage.Member {
+	t.Helper()
+
+	member, err := baggage.NewMember(key, value)
+	require.NoError(t, err)
+
+	return member
+}
+
+// testAssertEqualMembers verifies the two slices of Members are equal by
+// sorting them and comparing them as pairs of key-value strings.
+//
+// This is necessary because Baggage.Members() returns an unordered slice of
+// Members.
+func testAssertEqualMembers(t *testing.T, got []baggage.Member, wants ...string) {
+	t.Helper()
+
+	type Pair struct {
+		Key   string
+		Value string
+	}
+
+	require.Equal(t, 0, len(wants)%2, "wants must be a multiple of 2")
+	require.Equal(t, len(wants)/2, len(got))
+
+	wantPairs := make([]Pair, 0, len(wants)/2)
+	for i := 0; i < len(wants); i += 2 {
+		key := wants[i]
+		value := wants[i+1]
+
+		wantPairs = append(wantPairs, Pair{
+			Key:   key,
+			Value: value,
+		})
+	}
+
+	gotPairs := make([]Pair, 0, len(got))
+	for _, m := range got {
+		gotPairs = append(gotPairs, Pair{m.Key(), m.Value()})
+	}
+
+	sort.Slice(wantPairs, func(i, j int) bool {
+		return wantPairs[i].Key < wantPairs[j].Key
+	})
+	sort.Slice(gotPairs, func(i, j int) bool {
+		return gotPairs[i].Key < gotPairs[j].Key
+	})
+
+	require.Equal(t, wantPairs, gotPairs)
+}
+
 func testPprofLabelsToMap(t *testing.T, labelSet pprof.LabelSet) map[string]string {
 	t.Helper()