fix dataraces by aws#102

logan · Jul 2, 2019 · 445cf09 · 445cf09
1 parent 52ef590
commit 445cf09
Show file tree

Hide file tree

Showing 2 changed files with 8 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")
 	}
@@ -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")
 	}

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