From c91452099472403cbc86a0cf7cf435981a687fc6 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Wed, 31 Aug 2022 18:36:00 -0500
Subject: [PATCH] Add mutex for writing to result

Fixes tcnksm/go-httpstat#21.
---
 go18.go     | 18 ++++++++++++++++++
 httpstat.go |  3 +++
 2 files changed, 21 insertions(+)

diff --git a/go18.go b/go18.go
index 3b5df3c..9c20709 100644
--- a/go18.go
+++ b/go18.go
@@ -41,10 +41,14 @@ func (r *Result) Total(t time.Time) time.Duration {
 func withClientTrace(ctx context.Context, r *Result) context.Context {
 	return httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
 		DNSStart: func(i httptrace.DNSStartInfo) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.dnsStart = time.Now()
 		},
 
 		DNSDone: func(i httptrace.DNSDoneInfo) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.dnsDone = time.Now()
 
 			r.DNSLookup = r.dnsDone.Sub(r.dnsStart)
@@ -52,6 +56,8 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		ConnectStart: func(_, _ string) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.tcpStart = time.Now()
 
 			// When connecting to IP (When no DNS lookup)
@@ -62,6 +68,8 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		ConnectDone: func(network, addr string, err error) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.tcpDone = time.Now()
 
 			r.TCPConnection = r.tcpDone.Sub(r.tcpStart)
@@ -69,11 +77,15 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		TLSHandshakeStart: func() {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.isTLS = true
 			r.tlsStart = time.Now()
 		},
 
 		TLSHandshakeDone: func(_ tls.ConnectionState, _ error) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.tlsDone = time.Now()
 
 			r.TLSHandshake = r.tlsDone.Sub(r.tlsStart)
@@ -81,6 +93,8 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		GotConn: func(i httptrace.GotConnInfo) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			// Handle when keep alive is used and connection is reused.
 			// DNSStart(Done) and ConnectStart(Done) is skipped
 			if i.Reused {
@@ -89,6 +103,8 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		WroteRequest: func(info httptrace.WroteRequestInfo) {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.serverStart = time.Now()
 
 			// When client doesn't use DialContext or using old (before go1.7) `net`
@@ -123,6 +139,8 @@ func withClientTrace(ctx context.Context, r *Result) context.Context {
 		},
 
 		GotFirstResponseByte: func() {
+			r.mutex.Lock()
+			defer r.mutex.Unlock()
 			r.serverDone = time.Now()
 
 			r.ServerProcessing = r.serverDone.Sub(r.serverStart)
diff --git a/httpstat.go b/httpstat.go
index 15ce78c..ececabb 100644
--- a/httpstat.go
+++ b/httpstat.go
@@ -8,11 +8,14 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"sync"
 	"time"
 )
 
 // Result stores httpstat info.
 type Result struct {
+	mutex sync.Mutex
+
 	// The following are duration for each phase
 	DNSLookup        time.Duration
 	TCPConnection    time.Duration