WIP: Datadog Baggage API

ddtrace/ddtrace.go

@@ -79,6 +79,16 @@ type Span interface {
 	// 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

@@ -80,6 +80,17 @@ func (NoopSpan) BaggageItem(_ 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

@@ -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

@@ -146,6 +146,31 @@ func TestSpanBaggageFunctions(t *testing.T) {
 		s.SetBaggageItem("a", "b")
 		assert.Equal(t, "b", s.BaggageItem("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)
+	})
 }
 
 func TestSpanContext(t *testing.T) {

ddtrace/mocktracer/mockspancontext.go

@@ -54,6 +54,38 @@ func (sc *spanContext) baggageItem(k string) string {
 	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

@@ -31,14 +31,38 @@ func TestSpanContextSetBaggage(t *testing.T) {
 	assert.Equal(t, sc.baggage["c"], "d")
 }
 
-func TestSpanContextGetBaggage(t *testing.T) {
+func TestSpanContextBaggage(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")
 }
 
+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) {
 	var sc spanContext
 	sc.setBaggageItem("a", "b")

ddtrace/tracer/civisibility_tslv.go

@@ -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

@@ -109,6 +109,22 @@ func (s *span) BaggageItem(key string) string {
 	return s.context.baggageItem(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.
 func (s *span) SetTag(key string, value interface{}) {
 	// To avoid dumping the memory address in case value is a pointer, we dereference it.

ddtrace/tracer/span_test.go

@@ -57,6 +57,42 @@ func TestSpanBaggage(t *testing.T) {
 	assert.Equal("value", span.BaggageItem("key"))
 }
 
+func TestSpanBaggageRemove(t *testing.T) {
+	assert := assert.New(t)
+
+	span := newBasicSpan("web.request")
+	span.SetBaggageItem("key", "value")
+	assert.Equal("value", span.BaggageItem("key"))
+	// test remove baggage
+	span.RemoveBaggageItem("key")
+	assert.Equal("", span.BaggageItem("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.BaggageItem("key"))
+	assert.Equal("value2", span.BaggageItem("key2"))
+	// test remove all baggage
+	span.RemoveAllBaggageItems()
+	assert.Equal("", span.BaggageItem("key"))
+	assert.Equal("", span.BaggageItem("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) {
 	assert := assert.New(t)
 

ddtrace/tracer/spancontext.go

@@ -241,6 +241,48 @@ func (c *spanContext) baggageItem(key string) string {
 	return c.baggage[key]
 }
 
+func (c *spanContext) getAllBaggageItems() map[string]string {
+	if atomic.LoadUint32(&c.hasBaggage) == 0 {
+		// Return an empty map to avoid nil issues
+		return map[string]string{}
+	}
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	copy := make(map[string]string, len(c.baggage))
+	for k, v := range c.baggage {
+		copy[k] = v
+	}
+	return copy
+}
+
+func (c *spanContext) removeBaggageItem(key string) {
+	if atomic.LoadUint32(&c.hasBaggage) == 0 {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if _, ok := c.baggage[key]; ok {
+		delete(c.baggage, key)
+		// If baggage is now empty, reset hasBaggage
+		if len(c.baggage) == 0 {
+			atomic.StoreUint32(&c.hasBaggage, 0)
+			c.baggage = nil
+		}
+	}
+}
+
+func (c *spanContext) removeAllBaggageItems() {
+	if atomic.LoadUint32(&c.hasBaggage) == 0 {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	// Clear baggage
+	c.baggage = nil
+	atomic.StoreUint32(&c.hasBaggage, 0)
+}
+
 func (c *spanContext) meta(key string) (val string, ok bool) {
 	c.span.RLock()
 	defer c.span.RUnlock()

ddtrace/tracer/spancontext_test.go

@@ -783,6 +783,14 @@ func TestSpanContextBaggage(t *testing.T) {
 	var ctx spanContext
 	ctx.setBaggageItem("key", "value")
 	assert.Equal("value", ctx.baggage["key"])
+	ctx.setBaggageItem("foo", "bar")
+	assert.Equal("bar", ctx.baggageItem("foo"))
+	ctx.removeBaggageItem("key")
+	assert.NotContains(ctx.baggage, "key")
+	ctx.setBaggageItem("key1", "value1")
+	assert.Equal(map[string]string{"foo": "bar", "key1": "value1"}, ctx.getAllBaggageItems())
+	ctx.removeAllBaggageItems()
+	assert.Empty(ctx.baggage)
 }
 
 func TestSpanContextIterator(t *testing.T) {