fix sub-subsegment datarace (#61)

xray/capture_test.go

@@ -12,6 +12,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"sync"
 	"testing"
 
 	"github.com/aws/aws-xray-sdk-go/strategy/exception"
@@ -125,15 +126,20 @@ func TestNoSegmentCapture(t *testing.T) {
 }
 
 func TestCaptureAsync(t *testing.T) {
+	var mu sync.Mutex
 	ctx, root := BeginSegment(context.Background(), "Test")
 	CaptureAsync(ctx, "TestService", func(ctx1 context.Context) error {
+		mu.Lock()
+		defer mu.Unlock()
 		ctx = ctx1
 		defer root.Close(nil)
 		return nil
 	})
 
 	s, e := TestDaemon.Recv()
 	assert.NoError(t, e)
+	mu.Lock() // avoid race detection
+	defer mu.Unlock()
 	assert.Equal(t, "Test", s.Name)
 	assert.Equal(t, root.TraceID, s.TraceID)
 	assert.Equal(t, root.ID, s.ID)
@@ -143,4 +149,4 @@ func TestCaptureAsync(t *testing.T) {
 	subseg := &Segment{}
 	assert.NoError(t, json.Unmarshal(s.Subsegments[0], &subseg))
 	assert.Equal(t, "TestService", subseg.Name)
-}
+}

xray/segment.go

@@ -180,11 +180,16 @@ func BeginSubsegment(ctx context.Context, name string) (context.Context, *Segmen
 	defer seg.Unlock()
 
 	parent.Lock()
+	defer parent.Unlock()
+
 	seg.ParentSegment = parent.ParentSegment
+	if seg.ParentSegment != seg && seg.ParentSegment != parent {
+		seg.ParentSegment.Lock()
+		defer seg.ParentSegment.Unlock()
+	}
 	seg.ParentSegment.totalSubSegments++
 	parent.rawSubsegments = append(parent.rawSubsegments, seg)
 	parent.openSegments++
-	parent.Unlock()
 
 	seg.ID = NewSegmentID()
 	seg.Name = name
@@ -204,7 +209,7 @@ func NewSegmentFromHeader(ctx context.Context, name string, h *header.Header) (c
 	if h.ParentID != "" {
 		seg.ParentID = h.ParentID
 	}
-	
+
 	seg.Sampled = h.SamplingDecision == header.Sampled
 	switch h.SamplingDecision {
 	case header.Sampled:

xray/segment_test.go

@@ -10,6 +10,7 @@ package xray
 
 import (
 	"context"
+	"sync"
 	"testing"
 )
 
@@ -23,3 +24,23 @@ func TestSegmentDataRace(t *testing.T) {
 		cancel()
 	}
 }
+
+func TestSubsegmentDataRace(t *testing.T) {
+	ctx := context.Background()
+	ctx, seg := BeginSegment(ctx, "TestSegment")
+	defer seg.Close(nil)
+
+	wg := sync.WaitGroup{}
+	n := 5
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			defer wg.Done()
+			ctx, seg := BeginSubsegment(ctx, "TestSubsegment1")
+			ctx, seg2 := BeginSubsegment(ctx, "TestSubsegment2")
+			seg2.Close(nil)
+			seg.Close(nil)
+		}()
+	}
+	wg.Wait()
+}