fix dataraces by aws#102

t-yuki · Apr 4, 2019 · f035196 · f035196
1 parent 17591b3
commit f035196
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/xray/httptrace.go b/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")
 	}

diff --git a/xray/segment.go b/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