WIP: Datadog Baggage API

ddtrace/tracer/baggage.go

@@ -0,0 +1,75 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package tracer
+
+import (
+	"context"
+)
+
+// baggageKey is an unexported type used as a context key. It is used to store baggage in the context.
+// We use a struct{} so it won't conflict with keys from other packages.
+type baggageKey struct{}
+
+// baggageMap retrieves the map that holds baggage from the context.
+// Returns nil if no baggage map is stored in the context.
+func baggageMap(ctx context.Context) map[string]string {
+	bm, _ := ctx.Value(baggageKey{}).(map[string]string)
+	return bm
+}
+
+// WithBaggage returns a new context with the given baggage map set.
+func WithBaggage(ctx context.Context, baggage map[string]string) context.Context {
+	return context.WithValue(ctx, baggageKey{}, baggage)
+}
+
+// SetBaggage sets or updates a single baggage key/value pair in the context.
+func SetBaggage(ctx context.Context, key, value string) context.Context {
+	bm := baggageMap(ctx)
+	if bm == nil {
+		bm = make(map[string]string)
+	}
+	bm[key] = value
+	return WithBaggage(ctx, bm)
+}
+
+// Baggage retrieves the value associated with a baggage key.
+// If the key isn't found, it returns an empty string.
+func Baggage(ctx context.Context, key string) string {
+	bm := baggageMap(ctx)
+	if bm == nil {
+		return ""
+	}
+	return bm[key]
+}
+
+// RemoveBaggage removes the specified key from the baggage (if present).
+func RemoveBaggage(ctx context.Context, key string) context.Context {
+	bm := baggageMap(ctx)
+	if bm == nil {
+		// nothing to remove
+		return ctx
+	}
+	delete(bm, key)
+	return WithBaggage(ctx, bm)
+}
+
+// AllBaggage returns a **copy** of all baggage items in the context,
+func AllBaggage(ctx context.Context) map[string]string {
+	bm := baggageMap(ctx)
+	if bm == nil {
+		return nil
+	}
+	copyMap := make(map[string]string, len(bm))
+	for k, v := range bm {
+		copyMap[k] = v
+	}
+	return copyMap
+}
+
+// ClearBaggage completely removes all baggage items from the context.
+func ClearBaggage(ctx context.Context) context.Context {
+	return WithBaggage(ctx, nil)
+}

ddtrace/tracer/baggage_test.go

@@ -0,0 +1,108 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package tracer
+
+import (
+	"context"
+	"testing"
+)
+
+func TestBaggageFunctions(t *testing.T) {
+	t.Run("SetBaggage and Baggage", func(t *testing.T) {
+		ctx := context.Background()
+
+		// Set a key/value in the baggage
+		ctx = SetBaggage(ctx, "foo", "bar")
+
+		// Retrieve that value
+		got := Baggage(ctx, "foo")
+		want := "bar"
+		if got != want {
+			t.Errorf("Baggage(ctx, \"foo\") = %q; want %q", got, want)
+		}
+
+		// Ensure retrieving a non-existent key returns an empty string
+		got = Baggage(ctx, "missingKey")
+		if got != "" {
+			t.Errorf("Baggage(ctx, \"missingKey\") = %q; want \"\"", got)
+		}
+	})
+
+	t.Run("AllBaggage", func(t *testing.T) {
+		ctx := context.Background()
+
+		// Set multiple baggage entries
+		ctx = SetBaggage(ctx, "key1", "value1")
+		ctx = SetBaggage(ctx, "key2", "value2")
+
+		// Retrieve all baggage entries
+		all := AllBaggage(ctx)
+		if len(all) != 2 {
+			t.Fatalf("Expected 2 items in baggage; got %d", len(all))
+		}
+
+		// Check each entry
+		if all["key1"] != "value1" {
+			t.Errorf("all[\"key1\"] = %q; want \"value1\"", all["key1"])
+		}
+		if all["key2"] != "value2" {
+			t.Errorf("all[\"key2\"] = %q; want \"value2\"", all["key2"])
+		}
+
+		// Confirm returned map is a copy, not the original
+		all["key1"] = "modified"
+		if Baggage(ctx, "key1") == "modified" {
+			t.Error("AllBaggage returned a map that mutates the original baggage!")
+		}
+	})
+
+	t.Run("RemoveBaggage", func(t *testing.T) {
+		ctx := context.Background()
+
+		// Add baggage to remove
+		ctx = SetBaggage(ctx, "deleteMe", "toBeRemoved")
+
+		// Remove it
+		ctx = RemoveBaggage(ctx, "deleteMe")
+
+		// Verify removal
+		got := Baggage(ctx, "deleteMe")
+		if got != "" {
+			t.Errorf("Expected empty string for removed key; got %q", got)
+		}
+	})
+
+	t.Run("ClearBaggage", func(t *testing.T) {
+		ctx := context.Background()
+
+		// Add multiple items
+		ctx = SetBaggage(ctx, "k1", "v1")
+		ctx = SetBaggage(ctx, "k2", "v2")
+
+		// Clear all baggage
+		ctx = ClearBaggage(ctx)
+
+		// Check that everything is gone
+		all := AllBaggage(ctx)
+		if len(all) != 0 {
+			t.Errorf("Expected no items after clearing baggage; got %d", len(all))
+		}
+	})
+
+	t.Run("WithBaggage", func(t *testing.T) {
+		ctx := context.Background()
+
+		// Create a map and insert into context directly
+		initialMap := map[string]string{"customKey": "customValue"}
+		ctx = WithBaggage(ctx, initialMap)
+
+		// Verify
+		got := Baggage(ctx, "customKey")
+		if got != "customValue" {
+			t.Errorf("Baggage(ctx, \"customKey\") = %q; want \"customValue\"", got)
+		}
+	})
+}