Make SpanContext immutable (#106)

Per opentracing/opentracing.io#106 (same id, ironically)

ext/tags_test.go

@@ -35,9 +35,9 @@ func TestPeerTags(t *testing.T) {
 		"peer.port":     uint16(8080),
 		"span.kind":     ext.SpanKindRPCClientEnum,
 	}, rawSpan.Tags())
-	assert.True(t, span.Context().(*mocktracer.MockSpanContext).Sampled)
+	assert.True(t, span.Context().(mocktracer.MockSpanContext).Sampled)
 	ext.SamplingPriority.Set(span, uint16(0))
-	assert.False(t, span.Context().(*mocktracer.MockSpanContext).Sampled)
+	assert.False(t, span.Context().(mocktracer.MockSpanContext).Sampled)
 }
 
 func TestHTTPTags(t *testing.T) {
@@ -76,7 +76,7 @@ func TestMiscTags(t *testing.T) {
 func TestRPCServerOption(t *testing.T) {
 	tracer := mocktracer.New()
 	parent := tracer.StartSpan("my-trace")
-	parent.Context().SetBaggageItem("bag", "gage")
+	parent.SetBaggageItem("bag", "gage")
 
 	carrier := opentracing.HTTPHeadersCarrier{}
 	err := tracer.Inject(parent.Context(), opentracing.HTTPHeaders, carrier)
@@ -97,5 +97,5 @@ func TestRPCServerOption(t *testing.T) {
 	}, rawSpan.Tags())
 	assert.Equal(t, map[string]string{
 		"bag": "gage",
-	}, rawSpan.Context().(*mocktracer.MockSpanContext).Baggage())
+	}, rawSpan.Context().(mocktracer.MockSpanContext).Baggage)
 }

mocktracer/mockspan.go

@@ -18,12 +18,10 @@ import (
 // By default all spans have Sampled=true flag, unless {"sampling.priority": 0}
 // tag is set.
 type MockSpanContext struct {
-	sync.RWMutex
-
 	TraceID int
 	SpanID  int
 	Sampled bool
-	baggage map[string]string
+	Baggage map[string]string
 }
 
 var mockIDSource = uint32(42)
@@ -33,71 +31,44 @@ func nextMockID() int {
 	return int(atomic.LoadUint32(&mockIDSource))
 }
 
-func newMockSpanContext(traceID int, spanID int, sampled bool, baggage map[string]string) *MockSpanContext {
-	baggageCopy := make(map[string]string)
-	for k, v := range baggage {
-		baggageCopy[k] = v
-	}
-	return &MockSpanContext{
-		TraceID: traceID,
-		SpanID:  spanID,
-		Sampled: sampled,
-		baggage: baggageCopy,
-	}
-}
-
-// SetBaggageItem belongs to the SpanContext interface
-func (s *MockSpanContext) SetBaggageItem(key, val string) opentracing.SpanContext {
-	s.Lock()
-	defer s.Unlock()
-	if s.baggage == nil {
-		s.baggage = make(map[string]string)
-	}
-	s.baggage[key] = val
-	return s
-}
-
-// BaggageItem belongs to the SpanContext interface
-func (s *MockSpanContext) BaggageItem(key string) string {
-	s.RLock()
-	defer s.RUnlock()
-	return s.baggage[key]
-}
-
 // ForeachBaggageItem belongs to the SpanContext interface
-func (s *MockSpanContext) ForeachBaggageItem(handler func(k, v string) bool) {
-	s.RLock()
-	defer s.RUnlock()
-	for k, v := range s.baggage {
+func (c MockSpanContext) ForeachBaggageItem(handler func(k, v string) bool) {
+	for k, v := range c.Baggage {
 		if !handler(k, v) {
 			break
 		}
 	}
 }
 
-// Baggage returns a copy of baggage items in the span
-func (s *MockSpanContext) Baggage() map[string]string {
-	s.RLock()
-	defer s.RUnlock()
-	baggage := make(map[string]string)
-	for k, v := range s.baggage {
-		baggage[k] = v
+// WithBaggageItem creates a new context with an extra baggage item.
+func (c MockSpanContext) WithBaggageItem(key, value string) MockSpanContext {
+	var newBaggage map[string]string
+	if c.Baggage == nil {
+		newBaggage = map[string]string{key: value}
+	} else {
+		newBaggage = make(map[string]string, len(c.Baggage)+1)
+		for k, v := range c.Baggage {
+			newBaggage[k] = v
+		}
+		newBaggage[key] = value
 	}
-	return baggage
+	// Use positional parameters so the compiler will help catch new fields.
+	return MockSpanContext{c.TraceID, c.SpanID, c.Sampled, newBaggage}
 }
 
 // MockSpan is an opentracing.Span implementation that exports its internal
 // state for testing purposes.
 type MockSpan struct {
+	sync.RWMutex
+
 	ParentID int
 
 	OperationName string
 	StartTime     time.Time
 	FinishTime    time.Time
 
-	// All of the below (including SpanContext) are protected by SpanContext's
-	// embedded RWMutex.
-	SpanContext *MockSpanContext
+	// All of the below are protected by the embedded RWMutex.
+	SpanContext MockSpanContext
 	tags        map[string]interface{}
 	logs        []opentracing.LogData
 	tracer      *MockTracer
@@ -113,12 +84,12 @@ func newMockSpan(t *MockTracer, name string, opts opentracing.StartSpanOptions)
 	var baggage map[string]string
 	sampled := true
 	if len(opts.References) > 0 {
-		traceID = opts.References[0].ReferencedContext.(*MockSpanContext).TraceID
-		parentID = opts.References[0].ReferencedContext.(*MockSpanContext).SpanID
-		sampled = opts.References[0].ReferencedContext.(*MockSpanContext).Sampled
-		baggage = opts.References[0].ReferencedContext.(*MockSpanContext).Baggage()
+		traceID = opts.References[0].ReferencedContext.(MockSpanContext).TraceID
+		parentID = opts.References[0].ReferencedContext.(MockSpanContext).SpanID
+		sampled = opts.References[0].ReferencedContext.(MockSpanContext).Sampled
+		baggage = opts.References[0].ReferencedContext.(MockSpanContext).Baggage
 	}
-	spanContext := newMockSpanContext(traceID, nextMockID(), sampled, baggage)
+	spanContext := MockSpanContext{traceID, nextMockID(), sampled, baggage}
 	startTime := opts.StartTime
 	if startTime.IsZero() {
 		startTime = time.Now()
@@ -137,8 +108,8 @@ func newMockSpan(t *MockTracer, name string, opts opentracing.StartSpanOptions)
 
 // Tags returns a copy of tags accumulated by the span so far
 func (s *MockSpan) Tags() map[string]interface{} {
-	s.SpanContext.RLock()
-	defer s.SpanContext.RUnlock()
+	s.RLock()
+	defer s.RUnlock()
 	tags := make(map[string]interface{})
 	for k, v := range s.tags {
 		tags[k] = v
@@ -148,15 +119,15 @@ func (s *MockSpan) Tags() map[string]interface{} {
 
 // Tag returns a single tag
 func (s *MockSpan) Tag(k string) interface{} {
-	s.SpanContext.RLock()
-	defer s.SpanContext.RUnlock()
+	s.RLock()
+	defer s.RUnlock()
 	return s.tags[k]
 }
 
 // Logs returns a copy of logs accumulated in the span so far
 func (s *MockSpan) Logs() []opentracing.LogData {
-	s.SpanContext.RLock()
-	defer s.SpanContext.RUnlock()
+	s.RLock()
+	defer s.RUnlock()
 	logs := make([]opentracing.LogData, len(s.logs))
 	copy(logs, s.logs)
 	return logs
@@ -169,8 +140,8 @@ func (s *MockSpan) Context() opentracing.SpanContext {
 
 // SetTag belongs to the Span interface
 func (s *MockSpan) SetTag(key string, value interface{}) opentracing.Span {
-	s.SpanContext.Lock()
-	defer s.SpanContext.Unlock()
+	s.Lock()
+	defer s.Unlock()
 	if key == string(ext.SamplingPriority) {
 		if v, ok := value.(uint16); ok {
 			s.SpanContext.Sampled = v > 0
@@ -185,20 +156,35 @@ func (s *MockSpan) SetTag(key string, value interface{}) opentracing.Span {
 	return s
 }
 
+// SetBaggageItem belongs to the Span interface
+func (s *MockSpan) SetBaggageItem(key, val string) opentracing.Span {
+	s.Lock()
+	defer s.Unlock()
+	s.SpanContext = s.SpanContext.WithBaggageItem(key, val)
+	return s
+}
+
+// BaggageItem belongs to the Span interface
+func (s *MockSpan) BaggageItem(key string) string {
+	s.RLock()
+	defer s.RUnlock()
+	return s.SpanContext.Baggage[key]
+}
+
 // Finish belongs to the Span interface
 func (s *MockSpan) Finish() {
-	s.SpanContext.Lock()
+	s.Lock()
 	s.FinishTime = time.Now()
-	s.SpanContext.Unlock()
+	s.Unlock()
 	s.tracer.recordSpan(s)
 }
 
 // FinishWithOptions belongs to the Span interface
 func (s *MockSpan) FinishWithOptions(opts opentracing.FinishOptions) {
-	s.SpanContext.Lock()
+	s.Lock()
 	s.FinishTime = opts.FinishTime
 	s.logs = append(s.logs, opts.BulkLogData...)
-	s.SpanContext.Unlock()
+	s.Unlock()
 	s.tracer.recordSpan(s)
 }
 
@@ -227,15 +213,15 @@ func (s *MockSpan) LogEventWithPayload(event string, payload interface{}) {
 
 // Log belongs to the Span interface
 func (s *MockSpan) Log(data opentracing.LogData) {
-	s.SpanContext.Lock()
-	defer s.SpanContext.Unlock()
+	s.Lock()
+	defer s.Unlock()
 	s.logs = append(s.logs, data)
 }
 
 // SetOperationName belongs to the Span interface
 func (s *MockSpan) SetOperationName(operationName string) opentracing.Span {
-	s.SpanContext.Lock()
-	defer s.SpanContext.Unlock()
+	s.Lock()
+	defer s.Unlock()
 	s.OperationName = operationName
 	return s
 }

mocktracer/mocktracer.go

@@ -76,7 +76,7 @@ func (t *MockTracer) RegisterExtractor(format interface{}, extractor Extractor)
 
 // Inject belongs to the Tracer interface.
 func (t *MockTracer) Inject(sm opentracing.SpanContext, format interface{}, carrier interface{}) error {
-	spanContext, ok := sm.(*MockSpanContext)
+	spanContext, ok := sm.(MockSpanContext)
 	if !ok {
 		return opentracing.ErrInvalidCarrier
 	}

mocktracer/mocktracer_test.go

@@ -26,7 +26,7 @@ func TestMockTracer_StartSpan(t *testing.T) {
 	parent := spans[1]
 	child := spans[0]
 	assert.Equal(t, map[string]interface{}{"x": "y"}, parent.Tags())
-	assert.Equal(t, child.ParentID, parent.Context().(*MockSpanContext).SpanID)
+	assert.Equal(t, child.ParentID, parent.Context().(MockSpanContext).SpanID)
 }
 
 func TestMockSpan_SetOperationName(t *testing.T) {
@@ -39,9 +39,9 @@ func TestMockSpan_SetOperationName(t *testing.T) {
 func TestMockSpanContext_Baggage(t *testing.T) {
 	tracer := New()
 	span := tracer.StartSpan("x")
-	span.Context().SetBaggageItem("x", "y")
-	assert.Equal(t, "y", span.Context().BaggageItem("x"))
-	assert.Equal(t, map[string]string{"x": "y"}, span.Context().(*MockSpanContext).Baggage())
+	span.SetBaggageItem("x", "y")
+	assert.Equal(t, "y", span.BaggageItem("x"))
+	assert.Equal(t, map[string]string{"x": "y"}, span.Context().(MockSpanContext).Baggage)
 
 	baggage := make(map[string]string)
 	span.Context().ForeachBaggageItem(func(k, v string) bool {
@@ -50,13 +50,13 @@ func TestMockSpanContext_Baggage(t *testing.T) {
 	})
 	assert.Equal(t, map[string]string{"x": "y"}, baggage)
 
-	span.Context().SetBaggageItem("a", "b")
+	span.SetBaggageItem("a", "b")
 	baggage = make(map[string]string)
 	span.Context().ForeachBaggageItem(func(k, v string) bool {
 		baggage[k] = v
 		return false // exit early
 	})
-	assert.Equal(t, 2, len(span.Context().(*MockSpanContext).Baggage()))
+	assert.Equal(t, 2, len(span.Context().(MockSpanContext).Baggage))
 	assert.Equal(t, 1, len(baggage))
 }
 
@@ -116,7 +116,7 @@ func TestMockTracer_Propagation(t *testing.T) {
 	for _, test := range tests {
 		tracer := New()
 		span := tracer.StartSpan("x")
-		span.Context().SetBaggageItem("x", "y")
+		span.SetBaggageItem("x", "y")
 		if !test.sampled {
 			ext.SamplingPriority.Set(span, 0)
 		}
@@ -140,9 +140,9 @@ func TestMockTracer_Propagation(t *testing.T) {
 
 		extractedContext, err := tracer.Extract(opentracing.TextMap, opentracing.TextMapCarrier(carrier))
 		require.NoError(t, err)
-		assert.Equal(t, mSpan.SpanContext.TraceID, extractedContext.(*MockSpanContext).TraceID)
-		assert.Equal(t, mSpan.SpanContext.SpanID, extractedContext.(*MockSpanContext).SpanID)
-		assert.Equal(t, test.sampled, extractedContext.(*MockSpanContext).Sampled)
-		assert.Equal(t, "y", extractedContext.BaggageItem("x"))
+		assert.Equal(t, mSpan.SpanContext.TraceID, extractedContext.(MockSpanContext).TraceID)
+		assert.Equal(t, mSpan.SpanContext.SpanID, extractedContext.(MockSpanContext).SpanID)
+		assert.Equal(t, test.sampled, extractedContext.(MockSpanContext).Sampled)
+		assert.Equal(t, "y", extractedContext.(MockSpanContext).Baggage["x"])
 	}
 }