baggage API

Author: rachelyangdog

ddtrace/ddtrace.go

@@ -72,13 +72,23 @@ type Span interface {
 	// a representative name for a group of spans (e.g. "grpc.server" or "http.request").
 	SetOperationName(operationName string)
 
-	// BaggageItem returns the baggage item held by the given key.
-	BaggageItem(key string) string
+	// GetBaggageItem returns the baggage item held by the given key.
+	GetBaggageItem(key string) string
 
 	// SetBaggageItem sets a new baggage item at the given key. The baggage
 	// item should propagate to all descendant spans, both in- and cross-process.
 	SetBaggageItem(key, val string)
 
+	// GetAllBaggageItems returns a copy of all baggage items.
+	// If no items exist, returns an empty map.
+	GetAllBaggageItems() map[string]string
+
+	// RemoveBaggageItem removes a single baggage item by its key.
+	RemoveBaggageItem(key string)
+
+	// RemoveAllBaggageItems removes all baggage items from the span.
+	RemoveAllBaggageItems()
+
 	// Finish finishes the current span with the given options. Finish calls should be idempotent.
 	Finish(opts ...FinishOption)
 

ddtrace/internal/globaltracer.go

@@ -75,11 +75,22 @@ func (NoopSpan) SetTag(_ string, _ interface{}) {}
 func (NoopSpan) SetOperationName(_ string) {}
 
 // BaggageItem implements ddtrace.Span.
-func (NoopSpan) BaggageItem(_ string) string { return "" }
+func (NoopSpan) GetBaggageItem(_ string) string { return "" }
 
 // SetBaggageItem implements ddtrace.Span.
 func (NoopSpan) SetBaggageItem(_, _ string) {}
 
+// RemoveBaggageItem implements ddtrace.Span.
+func (NoopSpan) RemoveBaggageItem(_ string) {}
+
+// RemoveAllBaggageItems implements ddtrace.Span.
+func (NoopSpan) RemoveAllBaggageItems() {}
+
+// GetAllBaggageItems implements ddtrace.Span.
+func (NoopSpan) GetAllBaggageItems() map[string]string {
+	return make(map[string]string)
+}
+
 // Finish implements ddtrace.Span.
 func (NoopSpan) Finish(_ ...ddtrace.FinishOption) {}
 

ddtrace/mocktracer/mockspan.go

@@ -176,9 +176,9 @@ func (s *mockspan) SetOperationName(operationName string) {
 	return
 }
 
-// BaggageItem returns the baggage item with the given key.
-func (s *mockspan) BaggageItem(key string) string {
-	return s.context.baggageItem(key)
+// GetBaggageItem returns the baggage item with the given key.
+func (s *mockspan) GetBaggageItem(key string) string {
+	return s.context.getBaggageItem(key)
 }
 
 // SetBaggageItem sets a new baggage item at the given key. The baggage
@@ -188,6 +188,22 @@ func (s *mockspan) SetBaggageItem(key, val string) {
 	return
 }
 
+// RemoveBaggageItem removes the baggage item identified by the given key.
+func (s *mockspan) RemoveBaggageItem(key string) {
+	s.context.removeBaggageItem(key)
+}
+
+// RemoveAllBaggageItems removes all baggage items from this span.
+func (s *mockspan) RemoveAllBaggageItems() {
+	s.context.removeAllBaggageItems()
+}
+
+// GetAllBaggageItems returns a copy of all baggage items.
+// If there are no items, it returns an empty map.
+func (s *mockspan) GetAllBaggageItems() map[string]string {
+	return s.context.getAllBaggageItems()
+}
+
 // Finish finishes the current span with the given options.
 func (s *mockspan) Finish(opts ...ddtrace.FinishOption) {
 	var cfg ddtrace.FinishConfig

ddtrace/mocktracer/mockspan_test.go

@@ -144,7 +144,32 @@ func TestSpanBaggageFunctions(t *testing.T) {
 	t.Run("BaggageItem", func(t *testing.T) {
 		s := basicSpan("http.request")
 		s.SetBaggageItem("a", "b")
-		assert.Equal(t, "b", s.BaggageItem("a"))
+		assert.Equal(t, "b", s.GetBaggageItem("a"))
+	})
+
+	t.Run("GetAllBaggageItems", func(t *testing.T) {
+		s := basicSpan("http.request")
+		s.SetBaggageItem("a", "b")
+		s.SetBaggageItem("c", "d")
+		baggage := s.GetAllBaggageItems()
+		assert.Len(t, baggage, 2)
+		assert.Equal(t, "b", baggage["a"])
+		assert.Equal(t, "d", baggage["c"])
+	})
+
+	t.Run("RemoveBaggageItem", func(t *testing.T) {
+		s := basicSpan("http.request")
+		s.SetBaggageItem("a", "b")
+		s.RemoveBaggageItem("a")
+		assert.Empty(t, s.context.baggage)
+	})
+
+	t.Run("RemoveAllBaggageItems", func(t *testing.T) {
+		s := basicSpan("http.request")
+		s.SetBaggageItem("a", "b")
+		s.SetBaggageItem("c", "d")
+		s.RemoveAllBaggageItems()
+		assert.Empty(t, s.context.baggage)
 	})
 }
 

ddtrace/mocktracer/mockspancontext.go

@@ -48,12 +48,44 @@ func (sc *spanContext) setBaggageItem(k, v string) {
 	sc.baggage[k] = v
 }
 
-func (sc *spanContext) baggageItem(k string) string {
+func (sc *spanContext) getBaggageItem(k string) string {
 	sc.RLock()
 	defer sc.RUnlock()
 	return sc.baggage[k]
 }
 
+func (sc *spanContext) removeBaggageItem(k string) {
+	sc.Lock()
+	defer sc.Unlock()
+	if sc.baggage == nil {
+		return
+	}
+	delete(sc.baggage, k)
+	if len(sc.baggage) == 0 {
+		sc.baggage = nil
+	}
+}
+
+func (sc *spanContext) removeAllBaggageItems() {
+	sc.Lock()
+	defer sc.Unlock()
+	sc.baggage = nil
+}
+
+func (sc *spanContext) getAllBaggageItems() map[string]string {
+	sc.RLock()
+	defer sc.RUnlock()
+	if sc.baggage == nil {
+		return make(map[string]string)
+	}
+	// Return a copy to avoid callers mutating the internal map
+	copyMap := make(map[string]string, len(sc.baggage))
+	for key, val := range sc.baggage {
+		copyMap[key] = val
+	}
+	return copyMap
+}
+
 func (sc *spanContext) setSamplingPriority(p int) {
 	sc.Lock()
 	defer sc.Unlock()

ddtrace/mocktracer/mockspancontext_test.go

@@ -35,8 +35,32 @@ func TestSpanContextGetBaggage(t *testing.T) {
 	var sc spanContext
 	sc.setBaggageItem("a", "b")
 	sc.setBaggageItem("c", "d")
-	assert.Equal(t, sc.baggageItem("a"), "b")
-	assert.Equal(t, sc.baggageItem("c"), "d")
+	assert.Equal(t, sc.getBaggageItem("a"), "b")
+	assert.Equal(t, sc.getBaggageItem("c"), "d")
+}
+
+func TestSpanContextGetAllBaggage(t *testing.T) {
+	var sc spanContext
+	sc.setBaggageItem("a", "b")
+	sc.setBaggageItem("c", "d")
+	assert.Equal(t, map[string]string{"a": "b", "c": "d"}, sc.getAllBaggageItems())
+}
+
+func TestSpanContextRemoveBaggage(t *testing.T) {
+	var sc spanContext
+	sc.setBaggageItem("a", "b")
+	sc.setBaggageItem("c", "d")
+	sc.removeBaggageItem("a")
+	assert.Equal(t, sc.baggage["c"], "d")
+	assert.Empty(t, sc.baggage["a"])
+}
+
+func TestSpanContextRemoveAllBaggage(t *testing.T) {
+	var sc spanContext
+	sc.setBaggageItem("a", "b")
+	sc.setBaggageItem("c", "d")
+	sc.removeAllBaggageItems()
+	assert.Empty(t, sc.baggage)
 }
 
 func TestSpanContextIterator(t *testing.T) {

ddtrace/mocktracer/mocktracer_test.go

@@ -278,8 +278,8 @@ func TestTracerExtract(t *testing.T) {
 		assert.Nil(err)
 		sc, ok = ctx.(*spanContext)
 		assert.True(ok)
-		assert.Equal("B", sc.baggageItem("a"))
-		assert.Equal("D", sc.baggageItem("c"))
+		assert.Equal("B", sc.getBaggageItem("a"))
+		assert.Equal("D", sc.getBaggageItem("c"))
 
 		ctx, err = mt.Extract(carry(traceHeader, "1", spanHeader, "2", priorityHeader, "-1"))
 		assert.Nil(err)
@@ -302,7 +302,7 @@ func TestTracerExtract(t *testing.T) {
 
 		assert.Equal(uint64(1), got.traceID)
 		assert.Equal(uint64(2), got.spanID)
-		assert.Equal("D", got.baggageItem("c"))
-		assert.Equal("B", got.baggageItem("a"))
+		assert.Equal("D", got.getBaggageItem("c"))
+		assert.Equal("B", got.getBaggageItem("a"))
 	})
 }

ddtrace/tracer/civisibility_tslv.go

@@ -177,7 +177,7 @@ func (e *ciVisibilityEvent) SetOperationName(operationName string) {
 	e.Content.Name = e.span.Name
 }
 
-// BaggageItem retrieves the baggage item associated with the given key from the event's span.
+// GetBaggageItem retrieves the baggage item associated with the given key from the event's span.
 //
 // Parameters:
 //
@@ -186,8 +186,8 @@ func (e *ciVisibilityEvent) SetOperationName(operationName string) {
 // Returns:
 //
 //	The baggage item value.
-func (e *ciVisibilityEvent) BaggageItem(key string) string {
-	return e.span.BaggageItem(key)
+func (e *ciVisibilityEvent) GetBaggageItem(key string) string {
+	return e.span.GetBaggageItem(key)
 }
 
 // SetBaggageItem sets a baggage item on the event's span.
@@ -200,6 +200,26 @@ func (e *ciVisibilityEvent) SetBaggageItem(key, val string) {
 	e.span.SetBaggageItem(key, val)
 }
 
+// GetAllBaggageItems returns a copy of all baggage items from the event's span.
+// If there are no items, it returns an empty map.
+func (e *ciVisibilityEvent) GetAllBaggageItems() map[string]string {
+	return e.span.GetAllBaggageItems()
+}
+
+// RemoveBaggageItem removes a baggage item from the event's span.
+//
+// Parameters:
+//
+//	key - The baggage item key.
+func (e *ciVisibilityEvent) RemoveBaggageItem(key string) {
+	e.span.RemoveBaggageItem(key)
+}
+
+// RemoveAllBaggageItems removes all baggage items from the event's span.
+func (e *ciVisibilityEvent) RemoveAllBaggageItems() {
+	e.span.RemoveAllBaggageItems()
+}
+
 // Finish completes the event's span with optional finish options.
 //
 // Parameters:

ddtrace/tracer/span.go

@@ -103,10 +103,26 @@ func (s *span) SetBaggageItem(key, val string) {
 	s.context.setBaggageItem(key, val)
 }
 
-// BaggageItem gets the value for a baggage item given its key. Returns the
+// GetBaggageItem gets the value for a baggage item given its key. Returns the
 // empty string if the value isn't found in this Span.
-func (s *span) BaggageItem(key string) string {
-	return s.context.baggageItem(key)
+func (s *span) GetBaggageItem(key string) string { // change this to GetBaggageItem
+	return s.context.getBaggageItem(key)
+}
+
+// GetAllBaggageItems returns a copy of all baggage items.
+// If no items exist, returns an empty map.
+func (s *span) GetAllBaggageItems() map[string]string {
+	return s.context.getAllBaggageItems()
+}
+
+// RemoveBaggageItem removes a single baggage item by its key.
+func (s *span) RemoveBaggageItem(key string) {
+	s.context.removeBaggageItem(key)
+}
+
+// RemoveAllBaggageItems removes all baggage items from the span.
+func (s *span) RemoveAllBaggageItems() {
+	s.context.removeAllBaggageItems()
 }
 
 // SetTag adds a set of key/value metadata to the span.

ddtrace/tracer/span_test.go

@@ -54,7 +54,43 @@ func TestSpanBaggage(t *testing.T) {
 
 	span := newBasicSpan("web.request")
 	span.SetBaggageItem("key", "value")
-	assert.Equal("value", span.BaggageItem("key"))
+	assert.Equal("value", span.GetBaggageItem("key"))
+}
+
+func TestSpanBaggageRemove(t *testing.T) {
+	assert := assert.New(t)
+
+	span := newBasicSpan("web.request")
+	span.SetBaggageItem("key", "value")
+	assert.Equal("value", span.GetBaggageItem("key"))
+	// test remove baggage
+	span.RemoveBaggageItem("key")
+	assert.Equal("", span.GetBaggageItem("key"))
+}
+
+func TestSpanBaggageRemoveAll(t *testing.T) {
+	assert := assert.New(t)
+
+	span := newBasicSpan("web.request")
+	span.SetBaggageItem("key", "value")
+	span.SetBaggageItem("key2", "value2")
+	assert.Equal("value", span.GetBaggageItem("key"))
+	assert.Equal("value2", span.GetBaggageItem("key2"))
+	// test remove all baggage
+	span.RemoveAllBaggageItems()
+	assert.Equal("", span.GetBaggageItem("key"))
+	assert.Equal("", span.GetBaggageItem("key2"))
+}
+
+func TestSpanBaggageGetAll(t *testing.T) {
+	assert := assert.New(t)
+
+	span := newBasicSpan("web.request")
+	span.SetBaggageItem("key", "value")
+	span.SetBaggageItem("key2", "value2")
+	baggage := span.GetAllBaggageItems()
+	assert.Equal("value", baggage["key"])
+	assert.Equal("value2", baggage["key2"])
 }
 
 func TestSpanContext(t *testing.T) {