Save baggage in span (#153)

jaeger_span_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/opentracing/opentracing-go/ext"
 	"github.com/opentracing/opentracing-go/log"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	j "github.com/uber/jaeger-client-go/thrift-gen/jaeger"
 	"github.com/uber/jaeger-client-go/utils"
@@ -295,6 +296,32 @@ func TestBuildReferences(t *testing.T) {
 	assert.EqualValues(t, 2, spanRefs[1].TraceIdLow)
 }
 
+func TestJaegerSpanBaggageLogs(t *testing.T) {
+	tracer, closer := NewTracer("DOOP",
+		NewConstSampler(true),
+		NewNullReporter())
+	defer closer.Close()
+
+	sp := tracer.StartSpan("s1").(*Span)
+	sp.SetBaggageItem("auth.token", "token")
+	ext.SpanKindRPCServer.Set(sp)
+	sp.Finish()
+
+	jaegerSpan := buildJaegerSpan(sp)
+	require.Len(t, jaegerSpan.Logs, 1)
+	fields := jaegerSpan.Logs[0].Fields
+	require.Len(t, fields, 3)
+	assertJaegerTag(t, fields, "event", "baggage")
+	assertJaegerTag(t, fields, "key", "auth.token")
+	assertJaegerTag(t, fields, "value", "token")
+}
+
+func assertJaegerTag(t *testing.T, tags []*j.Tag, key string, value string) {
+	tag := findJaegerTag(key, tags)
+	require.NotNil(t, tag)
+	assert.Equal(t, value, tag.GetVStr())
+}
+
 func getStringPtr(s string) *string {
 	return &s
 }

propagation_test.go

@@ -86,6 +86,15 @@ func TestSpanPropagator(t *testing.T) {
 	exp, spans := spans[len(spans)-1], spans[:len(spans)-1]
 	exp.duration = time.Duration(123)
 	exp.startTime = time.Time{}.Add(1)
+	require.Len(t, exp.logs, 1) // The parent span should have baggage logs
+	fields := exp.logs[0].Fields
+	require.Len(t, fields, 3)
+	require.Equal(t, "event", fields[0].Key())
+	require.Equal(t, "baggage", fields[0].Value().(string))
+	require.Equal(t, "key", fields[1].Key())
+	require.Equal(t, "foo", fields[1].Value().(string))
+	require.Equal(t, "value", fields[2].Key())
+	require.Equal(t, "bar", fields[2].Value().(string))
 
 	if exp.context.ParentID() != 0 {
 		t.Fatalf("Root span's ParentID %d is not 0", exp.context.ParentID())
@@ -104,7 +113,7 @@ func TestSpanPropagator(t *testing.T) {
 		assert.Equal(t, exp.context, sp.context, formatName)
 		assert.Equal(t, "span.kind", sp.tags[0].key)
 		assert.Equal(t, expTags, sp.tags[1:] /*skip span.kind tag*/, formatName)
-		assert.Equal(t, exp.logs, sp.logs, formatName)
+		assert.Empty(t, sp.logs, formatName)
 		// Override collections to avoid tripping comparison on different pointers
 		sp.context = exp.context
 		sp.tags = exp.tags

span.go

@@ -66,7 +66,7 @@ type Span struct {
 	observer SpanObserver
 }
 
-// Tag a simple key value wrapper
+// Tag is a simple key value wrapper
 type Tag struct {
 	key   string
 	value interface{}
@@ -116,6 +116,11 @@ func (s *Span) LogFields(fields ...log.Field) {
 	if !s.context.IsSampled() {
 		return
 	}
+	s.logFieldsNoLocking(fields...)
+}
+
+// this function should only be called while holding a Write lock
+func (s *Span) logFieldsNoLocking(fields ...log.Field) {
 	lr := opentracing.LogRecord{
 		Fields:    fields,
 		Timestamp: time.Now(),
@@ -173,6 +178,14 @@ func (s *Span) SetBaggageItem(key, value string) opentracing.Span {
 	s.Lock()
 	defer s.Unlock()
 	s.context = s.context.WithBaggageItem(key, value)
+	if s.context.IsSampled() {
+		// If sampled, record the baggage in the span
+		s.logFieldsNoLocking(
+			log.String("event", "baggage"),
+			log.String("key", key),
+			log.String("value", value),
+		)
+	}
 	return s
 }
 

span_test.go

@@ -25,23 +25,40 @@ import (
 
 	"github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestBaggageIterator(t *testing.T) {
-	tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter())
+	service := "DOOP"
+	tracer, closer := NewTracer(service, NewConstSampler(true), NewNullReporter())
 	defer closer.Close()
 
 	sp1 := tracer.StartSpan("s1").(*Span)
 	sp1.SetBaggageItem("Some_Key", "12345")
 	sp1.SetBaggageItem("Some-other-key", "42")
 	expectedBaggage := map[string]string{"some-key": "12345", "some-other-key": "42"}
 	assertBaggage(t, sp1, expectedBaggage)
+	assertBaggageRecords(t, sp1, expectedBaggage)
 
 	b := extractBaggage(sp1, false) // break out early
 	assert.Equal(t, 1, len(b), "only one baggage item should be extracted")
 
-	sp2 := tracer.StartSpan("s2", opentracing.ChildOf(sp1.Context()))
+	sp2 := tracer.StartSpan("s2", opentracing.ChildOf(sp1.Context())).(*Span)
 	assertBaggage(t, sp2, expectedBaggage) // child inherits the same baggage
+	require.Len(t, sp2.logs, 0)            // child doesn't inherit the baggage logs
+}
+
+func assertBaggageRecords(t *testing.T, sp *Span, expected map[string]string) {
+	require.Len(t, sp.logs, len(expected))
+	for _, logRecord := range sp.logs {
+		require.Len(t, logRecord.Fields, 3)
+		require.Equal(t, "event:baggage", logRecord.Fields[0].String())
+		key := logRecord.Fields[1].Value().(string)
+		value := logRecord.Fields[2].Value().(string)
+
+		require.Contains(t, expected, key)
+		assert.Equal(t, expected[key], value)
+	}
 }
 
 func assertBaggage(t *testing.T, sp opentracing.Span, expected map[string]string) {

testutils/mock_agent_test.go

@@ -50,8 +50,14 @@ func TestMockAgentSpanServer(t *testing.T) {
 
 		err = client.EmitZipkinBatch(spans)
 		assert.NoError(t, err)
-		time.Sleep(5 * time.Millisecond)
 
+		for k := 0; k < 100; k++ {
+			time.Sleep(time.Millisecond)
+			spans = mockAgent.GetZipkinSpans()
+			if len(spans) == i {
+				break
+			}
+		}
 		spans = mockAgent.GetZipkinSpans()
 		require.Equal(t, i, len(spans))
 		for j := 0; j < i; j++ {

thrift_span_test.go

@@ -299,6 +299,21 @@ func TestSpecialTags(t *testing.T) {
 	assert.NotNil(t, findAnnotation(thriftSpan, "ss"))
 }
 
+func TestBaggageLogs(t *testing.T) {
+	tracer, closer := NewTracer("DOOP",
+		NewConstSampler(true),
+		NewNullReporter())
+	defer closer.Close()
+
+	sp := tracer.StartSpan("s1").(*Span)
+	sp.SetBaggageItem("auth.token", "token")
+	ext.SpanKindRPCServer.Set(sp)
+	sp.Finish()
+
+	thriftSpan := buildThriftSpan(sp)
+	assert.NotNil(t, findAnnotation(thriftSpan, `{"event":"baggage","key":"auth.token","value":"token"}`))
+}
+
 func findAnnotation(span *zipkincore.Span, name string) *zipkincore.Annotation {
 	for _, a := range span.Annotations {
 		if a.Value == name {