Skip to content

Commit

Permalink
baggage API
Browse files Browse the repository at this point in the history
  • Loading branch information
rachelyangdog committed Dec 17, 2024
1 parent 5a8a82c commit 2ab2a42
Show file tree
Hide file tree
Showing 12 changed files with 262 additions and 22 deletions.
14 changes: 12 additions & 2 deletions ddtrace/ddtrace.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
13 changes: 12 additions & 1 deletion ddtrace/internal/globaltracer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {}

Expand Down
22 changes: 19 additions & 3 deletions ddtrace/mocktracer/mockspan.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
27 changes: 26 additions & 1 deletion ddtrace/mocktracer/mockspan_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
}

Expand Down
34 changes: 33 additions & 1 deletion ddtrace/mocktracer/mockspancontext.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
28 changes: 26 additions & 2 deletions ddtrace/mocktracer/mockspancontext_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
8 changes: 4 additions & 4 deletions ddtrace/mocktracer/mocktracer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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"))
})
}
26 changes: 23 additions & 3 deletions ddtrace/tracer/civisibility_tslv.go
Original file line number Diff line number Diff line change
Expand Up @@ -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:
//
Expand All @@ -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.
Expand All @@ -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:
Expand Down
22 changes: 19 additions & 3 deletions ddtrace/tracer/span.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
38 changes: 37 additions & 1 deletion ddtrace/tracer/span_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading

0 comments on commit 2ab2a42

Please sign in to comment.