fix dataraces by #102

xray/aws.go

@@ -105,7 +105,7 @@ var xRayAfterSendHandler = request.NamedHandler{
 	Fn: func(r *request.Request) {
 		curseg := GetSegment(r.HTTPRequest.Context())
 
-		if curseg.Name == "attempt" {
+		if curseg != nil && curseg.Name == "attempt" {
 			// An error could have prevented the connect subsegment from closing,
 			// so clean it up here.
 			for _, subsegment := range curseg.rawSubsegments {

xray/aws_test.go

@@ -115,3 +115,25 @@ func TestClientFailedConnection(t *testing.T) {
 	assert.NotZero(t, connectSubseg.EndTime)
 	assert.NotEmpty(t, connectSubseg.Subsegments)
 }
+
+func TestClientWithoutSegment(t *testing.T) {
+	Configure(Config{ContextMissingStrategy: &TestContextMissingStrategy{}})
+	defer ResetConfig()
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b := []byte(`{}`)
+		w.WriteHeader(http.StatusOK)
+		w.Write(b)
+	}))
+
+	svc := lambda.New(session.Must(session.NewSession(&aws.Config{
+		Endpoint:    aws.String(ts.URL),
+		Region:      aws.String("fake-moon-1"),
+		Credentials: credentials.NewStaticCredentials("akid", "secret", "noop")})))
+
+	ctx := context.Background()
+
+	AWS(svc.Client)
+
+	_, err := svc.ListFunctionsWithContext(ctx, &lambda.ListFunctionsInput{})
+	assert.NoError(t, err)
+}

xray/httptrace.go

@@ -47,6 +47,8 @@ func (xt *HTTPSubsegments) GetConn(hostPort string) {
 // DNSStart begins a dns subsegment if the HTTP operation
 // subsegment is still in progress.
 func (xt *HTTPSubsegments) DNSStart(info httptrace.DNSStartInfo) {
+	xt.mu.Lock()
+	defer xt.mu.Unlock()
 	if GetSegment(xt.opCtx).safeInProgress() && xt.connCtx != nil {
 		xt.dnsCtx, _ = BeginSubsegment(xt.connCtx, "dns")
 	}
@@ -71,6 +73,8 @@ func (xt *HTTPSubsegments) DNSDone(info httptrace.DNSDoneInfo) {
 // ConnectStart begins a dial subsegment if the HTTP operation
 // subsegment is still in progress.
 func (xt *HTTPSubsegments) ConnectStart(network, addr string) {
+	xt.mu.Lock()
+	defer xt.mu.Unlock()
 	if GetSegment(xt.opCtx).safeInProgress() && xt.connCtx != nil {
 		xt.connectCtx, _ = BeginSubsegment(xt.connCtx, "dial")
 	}

xray/segment.go

@@ -285,6 +285,10 @@ func (seg *Segment) RemoveSubsegment(remove *Segment) bool {
 			seg.rawSubsegments[len(seg.rawSubsegments)-1] = nil
 			seg.rawSubsegments = seg.rawSubsegments[:len(seg.rawSubsegments)-1]
 
+			if seg.ParentSegment != seg {
+				seg.ParentSegment.Lock()
+				defer seg.ParentSegment.Unlock()
+			}
 			seg.ParentSegment.totalSubSegments--
 			seg.openSegments--
 			return true