From e86d745752232380d8531f11f388f92d02d71955 Mon Sep 17 00:00:00 2001 From: Rachel Yang Date: Thu, 16 Jan 2025 14:52:02 -0500 Subject: [PATCH] WIP: Datadog Baggage API (#3069) Co-authored-by: Mikayla Toffler <46911781+mtoffl01@users.noreply.github.com> --- ddtrace/baggage/baggage.go | 83 +++++++++++++++++++ ddtrace/baggage/baggage_test.go | 136 ++++++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 ddtrace/baggage/baggage.go create mode 100644 ddtrace/baggage/baggage_test.go diff --git a/ddtrace/baggage/baggage.go b/ddtrace/baggage/baggage.go new file mode 100644 index 0000000000..ee6428eebb --- /dev/null +++ b/ddtrace/baggage/baggage.go @@ -0,0 +1,83 @@ +// 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 baggage + +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 returns the baggage map from the given context and a bool indicating +// whether the baggage exists or not. If the bool is false, the returned map is nil. +func baggageMap(ctx context.Context) (map[string]string, bool) { + val := ctx.Value(baggageKey{}) + bm, ok := val.(map[string]string) + if !ok { + // val was nil or not a map[string]string + return nil, false + } + return bm, true +} + +// 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) +} + +// Set sets or updates a single baggage key/value pair in the context. +// If the key already exists, this function overwrites the existing value. +func Set(ctx context.Context, key, value string) context.Context { + bm, ok := baggageMap(ctx) + if !ok { + // If there's no baggage map yet, create one + bm = make(map[string]string) + } + bm[key] = value + return withBaggage(ctx, bm) +} + +// Get retrieves the value associated with a baggage key. +// If the key isn't found, it returns an empty string. +func Get(ctx context.Context, key string) (string, bool) { + bm, ok := baggageMap(ctx) + if !ok { + return "", false + } + value, ok := bm[key] + return value, ok +} + +// Remove removes the specified key from the baggage (if present). +func Remove(ctx context.Context, key string) context.Context { + bm, ok := baggageMap(ctx) + if !ok { + // nothing to remove + return ctx + } + delete(bm, key) + return withBaggage(ctx, bm) +} + +// All returns a **copy** of all baggage items in the context, +func All(ctx context.Context) map[string]string { + bm, ok := baggageMap(ctx) + if !ok { + return nil + } + copyMap := make(map[string]string, len(bm)) + for k, v := range bm { + copyMap[k] = v + } + return copyMap +} + +// Clear completely removes all baggage items from the context. +func Clear(ctx context.Context) context.Context { + return withBaggage(ctx, nil) +} diff --git a/ddtrace/baggage/baggage_test.go b/ddtrace/baggage/baggage_test.go new file mode 100644 index 0000000000..62b5401ba5 --- /dev/null +++ b/ddtrace/baggage/baggage_test.go @@ -0,0 +1,136 @@ +// 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 baggage + +import ( + "context" + "testing" +) + +func TestBaggageFunctions(t *testing.T) { + t.Run("Set and Get", func(t *testing.T) { + ctx := context.Background() + + // Set a key/value in the baggage + ctx = Set(ctx, "foo", "bar") + + // Retrieve that value + got, ok := Get(ctx, "foo") + if !ok { + t.Error("Expected key \"foo\" to be found in baggage, got ok=false") + } + if got != "bar" { + t.Errorf("Baggage(ctx, \"foo\") = %q; want \"bar\"", got) + } + + // Ensure retrieving a non-existent key returns an empty string and false + got, ok = Get(ctx, "missingKey") + if ok { + t.Error("Expected key \"missingKey\" to not be found, got ok=true") + } + if got != "" { + t.Errorf("Baggage(ctx, \"missingKey\") = %q; want \"\"", got) + } + }) + + t.Run("All", func(t *testing.T) { + ctx := context.Background() + + // Set multiple baggage entries + ctx = Set(ctx, "key1", "value1") + ctx = Set(ctx, "key2", "value2") + + // Retrieve all baggage entries + all := All(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" + val, _ := Get(ctx, "key1") + if val == "modified" { + t.Error("AllBaggage returned a map that mutates the original baggage!") + } + }) + + t.Run("Remove", func(t *testing.T) { + ctx := context.Background() + + // Add baggage to remove + ctx = Set(ctx, "deleteMe", "toBeRemoved") + + // Remove it + ctx = Remove(ctx, "deleteMe") + + // Verify removal + got, ok := Get(ctx, "deleteMe") + if ok { + t.Error("Expected key \"deleteMe\" to be removed, got ok=true") + } + if got != "" { + t.Errorf("Expected empty string for removed key; got %q", got) + } + }) + + t.Run("Clear", func(t *testing.T) { + ctx := context.Background() + + // Add multiple items + ctx = Set(ctx, "k1", "v1") + ctx = Set(ctx, "k2", "v2") + + // Clear all baggage + ctx = Clear(ctx) + + // Check that everything is gone + all := All(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, _ := Get(ctx, "customKey") + if got != "customValue" { + t.Errorf("Baggage(ctx, \"customKey\") = %q; want \"customValue\"", got) + } + }) + + t.Run("explicitOkCheck", func(t *testing.T) { + ctx := context.Background() + + // Check an unset key + val, ok := Get(ctx, "unsetKey") + if ok { + t.Errorf("Expected unset key to return ok=false, got ok=true with val=%q", val) + } + + ctx = Set(ctx, "testKey", "testVal") + val, ok = Get(ctx, "testKey") + if !ok { + t.Error("Expected key \"testKey\" to be present, got ok=false") + } + if val != "testVal" { + t.Errorf("Expected \"testVal\"; got %q", val) + } + }) +}