fix: negative trace substraction when using SetTimeout (#1038)

Author: nicolasbeauvais

request.go

@@ -1241,16 +1241,31 @@ func (r *Request) TraceInfo() TraceInfo {
 		return TraceInfo{}
 	}
 
+	ct.lock.RLock()
+	defer ct.lock.RUnlock()
+
 	ti := TraceInfo{
-		DNSLookup:      ct.dnsDone.Sub(ct.dnsStart),
-		TLSHandshake:   ct.tlsHandshakeDone.Sub(ct.tlsHandshakeStart),
-		ServerTime:     ct.gotFirstResponseByte.Sub(ct.gotConn),
+		DNSLookup:      0,
+		TCPConnTime:    0,
+		ServerTime:     0,
 		IsConnReused:   ct.gotConnInfo.Reused,
 		IsConnWasIdle:  ct.gotConnInfo.WasIdle,
 		ConnIdleTime:   ct.gotConnInfo.IdleTime,
 		RequestAttempt: r.Attempt,
 	}
 
+	if !ct.dnsStart.IsZero() && !ct.dnsDone.IsZero() {
+		ti.DNSLookup = ct.dnsDone.Sub(ct.dnsStart)
+	}
+
+	if !ct.tlsHandshakeDone.IsZero() && !ct.tlsHandshakeStart.IsZero() {
+		ti.TLSHandshake = ct.tlsHandshakeDone.Sub(ct.tlsHandshakeStart)
+	}
+
+	if !ct.gotFirstResponseByte.IsZero() && !ct.gotConn.IsZero() {
+		ti.ServerTime = ct.gotFirstResponseByte.Sub(ct.gotConn)
+	}
+
 	// Calculate the total time accordingly when connection is reused,
 	// and DNS start and get conn time may be zero if the request is invalid.
 	// See issue #1016.

request_test.go

@@ -1927,6 +1927,55 @@ func TestTraceInfoOnTimeout(t *testing.T) {
 	assertEqual(t, true, tr.TotalTime == resp.Duration())
 }
 
+func TestTraceInfoOnTimeoutWithSetTimeout(t *testing.T) {
+	t.Run("timeout with very short timeout", func(t *testing.T) {
+		client := New().
+			SetTimeout(1 * time.Millisecond).
+			SetBaseURL("http://resty-nowhere.local").
+			EnableTrace()
+
+		resp, err := client.R().Get("/")
+		assertNotNil(t, err)
+		assertNotNil(t, resp)
+
+		tr := resp.Request.TraceInfo()
+
+		assertEqual(t, true, tr.DNSLookup == 0)
+		assertEqual(t, true, tr.ConnTime == 0)
+		assertEqual(t, true, tr.TLSHandshake == 0)
+		assertEqual(t, true, tr.TCPConnTime == 0)
+		assertEqual(t, true, tr.ServerTime == 0)
+		assertEqual(t, true, tr.ResponseTime == 0)
+		assertEqual(t, true, tr.TotalTime > 0)
+		assertEqual(t, true, tr.TotalTime == resp.Duration())
+	})
+
+	t.Run("successful request with SetTimeout", func(t *testing.T) {
+		ts := createGetServer(t)
+		defer ts.Close()
+
+		client := New().
+			SetTimeout(5 * time.Second).
+			SetBaseURL(ts.URL).
+			EnableTrace()
+
+		resp, err := client.R().Get("/")
+		assertNil(t, err)
+		assertNotNil(t, resp)
+
+		tr := resp.Request.TraceInfo()
+
+		assertEqual(t, true, tr.DNSLookup >= 0)
+		assertEqual(t, true, tr.ConnTime >= 0)
+		assertEqual(t, true, tr.TLSHandshake >= 0)
+		assertEqual(t, true, tr.TCPConnTime >= 0)
+		assertEqual(t, true, tr.ServerTime >= 0)
+		assertEqual(t, true, tr.ResponseTime >= 0)
+		assertEqual(t, true, tr.TotalTime > 0)
+		assertEqual(t, true, tr.TotalTime == resp.Duration())
+	})
+}
+
 func TestDebugLoggerRequestBodyTooLarge(t *testing.T) {
 	formTs := createFormPostServer(t)
 	defer formTs.Close()

trace.go

@@ -10,6 +10,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"net/http/httptrace"
+	"sync"
 	"time"
 )
 
@@ -95,6 +96,7 @@ func (ti TraceInfo) Clone() *TraceInfo {
 // with the same naming for easy understanding. Plus additional insights
 // [Request].
 type clientTrace struct {
+	lock                 sync.RWMutex
 	getConn              time.Time
 	dnsStart             time.Time
 	dnsDone              time.Time
@@ -112,37 +114,55 @@ func (t *clientTrace) createContext(ctx context.Context) context.Context {
 		ctx,
 		&httptrace.ClientTrace{
 			DNSStart: func(_ httptrace.DNSStartInfo) {
+				t.lock.Lock()
 				t.dnsStart = time.Now()
+				t.lock.Unlock()
 			},
 			DNSDone: func(_ httptrace.DNSDoneInfo) {
+				t.lock.Lock()
 				t.dnsDone = time.Now()
+				t.lock.Unlock()
 			},
 			ConnectStart: func(_, _ string) {
+				t.lock.Lock()
 				if t.dnsDone.IsZero() {
 					t.dnsDone = time.Now()
 				}
 				if t.dnsStart.IsZero() {
 					t.dnsStart = t.dnsDone
 				}
+				t.lock.Unlock()
 			},
 			ConnectDone: func(net, addr string, err error) {
+				t.lock.Lock()
 				t.connectDone = time.Now()
+				t.lock.Unlock()
 			},
 			GetConn: func(_ string) {
+				t.lock.Lock()
 				t.getConn = time.Now()
+				t.lock.Unlock()
 			},
 			GotConn: func(ci httptrace.GotConnInfo) {
+				t.lock.Lock()
 				t.gotConn = time.Now()
 				t.gotConnInfo = ci
+				t.lock.Unlock()
 			},
 			GotFirstResponseByte: func() {
+				t.lock.Lock()
 				t.gotFirstResponseByte = time.Now()
+				t.lock.Unlock()
 			},
 			TLSHandshakeStart: func() {
+				t.lock.Lock()
 				t.tlsHandshakeStart = time.Now()
+				t.lock.Unlock()
 			},
 			TLSHandshakeDone: func(_ tls.ConnectionState, _ error) {
+				t.lock.Lock()
 				t.tlsHandshakeDone = time.Now()
+				t.lock.Unlock()
 			},
 		},
 	)