feat: stat HTTP request latencies (DNSLookup, TCP Connection and so on)

hrp/internal/httpstat/main.go

@@ -0,0 +1,169 @@
+// Package httpstat traces HTTP latency infomation (DNSLookup, TCP Connection and so on) on any golang HTTP request.
+// It uses `httptrace` package.
+// Forked from https://github.com/tcnksm/go-httpstat
+package httpstat
+
+import (
+	"context"
+	"crypto/tls"
+	"net/http/httptrace"
+	"time"
+)
+
+// Stat stores httpstat info.
+type Stat struct {
+	// The following are duration for each phase
+	// DNSLookup => TCPConnection => TLSHandshake => ServerProcessing => ContentTransfer
+	DNSLookup        time.Duration
+	TCPConnection    time.Duration
+	TLSHandshake     time.Duration
+	ServerProcessing time.Duration
+	ContentTransfer  time.Duration // from the first response byte to tansfer done.
+
+	// The followings are timeline of request
+	NameLookup    time.Duration // = DNSLookup
+	Connect       time.Duration // = DNSLookup + TCPConnection
+	Pretransfer   time.Duration // = DNSLookup + TCPConnection + TLSHandshake
+	StartTransfer time.Duration // = DNSLookup + TCPConnection + TLSHandshake + ServerProcessing
+	Total         time.Duration // = DNSLookup + TCPConnection + TLSHandshake + ServerProcessing + ContentTransfer
+
+	// internal timelines, including start and finish timestamps of each phase
+	dnsStart      time.Time
+	dnsDone       time.Time
+	tcpStart      time.Time
+	tcpDone       time.Time
+	tlsStart      time.Time
+	tlsDone       time.Time
+	serverStart   time.Time
+	serverDone    time.Time
+	transferStart time.Time
+	transferDone  time.Time // need to be provided from outside
+
+	// isTLS is true when connection seems to use TLS
+	isTLS bool
+
+	// isReused is true when connection is reused (keep-alive)
+	isReused bool
+}
+
+// Finish sets the time when reading response is done.
+// This must be called after reading response body.
+func (s *Stat) Finish() {
+	s.transferDone = time.Now()
+
+	// This means result is empty (it does nothing).
+	// Skip setting value (contentTransfer and total will be zero).
+	if s.dnsStart.IsZero() {
+		return
+	}
+
+	s.ContentTransfer = s.transferDone.Sub(s.transferStart)
+	s.Total = s.transferDone.Sub(s.dnsStart)
+}
+
+// Durations returns all durations and timelines of request latencies
+func (s *Stat) Durations() map[string]time.Duration {
+	return map[string]time.Duration{
+		"DNSLookup":        s.DNSLookup / time.Millisecond,
+		"TCPConnection":    s.TCPConnection / time.Millisecond,
+		"TLSHandshake":     s.TLSHandshake / time.Millisecond,
+		"ServerProcessing": s.ServerProcessing / time.Millisecond,
+		"ContentTransfer":  s.ContentTransfer / time.Millisecond,
+		"NameLookup":       s.NameLookup / time.Millisecond,
+		"Connect":          s.Connect / time.Millisecond,
+		"Pretransfer":      s.Connect / time.Millisecond,
+		"StartTransfer":    s.StartTransfer / time.Millisecond,
+		"Total":            s.Total / time.Millisecond,
+	}
+}
+
+// WithHTTPStat is a wrapper of httptrace.WithClientTrace.
+// It records the time of each httptrace hooks.
+func WithHTTPStat(ctx context.Context, s *Stat) context.Context {
+	return httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
+		DNSStart: func(i httptrace.DNSStartInfo) {
+			s.dnsStart = time.Now()
+		},
+
+		DNSDone: func(i httptrace.DNSDoneInfo) {
+			s.dnsDone = time.Now()
+
+			s.DNSLookup = s.dnsDone.Sub(s.dnsStart)
+			s.NameLookup = s.dnsDone.Sub(s.dnsStart)
+		},
+
+		ConnectStart: func(_, _ string) {
+			s.tcpStart = time.Now()
+
+			// When connecting to IP (When no DNS lookup)
+			if s.dnsStart.IsZero() {
+				s.dnsStart = s.tcpStart
+				s.dnsDone = s.tcpStart
+			}
+		},
+
+		ConnectDone: func(network, addr string, err error) {
+			s.tcpDone = time.Now()
+			s.TCPConnection = s.tcpDone.Sub(s.tcpStart)
+			s.Connect = s.tcpDone.Sub(s.dnsStart)
+		},
+
+		TLSHandshakeStart: func() {
+			s.isTLS = true
+			s.tlsStart = time.Now()
+		},
+
+		TLSHandshakeDone: func(_ tls.ConnectionState, _ error) {
+			s.tlsDone = time.Now()
+			s.TLSHandshake = s.tlsDone.Sub(s.tlsStart)
+			s.Pretransfer = s.tlsDone.Sub(s.dnsStart)
+		},
+
+		GotConn: func(i httptrace.GotConnInfo) {
+			// Handle when keep alive is used and connection is reused.
+			// DNSStart(Done) and ConnectStart(Done) is skipped
+			if i.Reused {
+				s.isReused = true
+			}
+		},
+
+		WroteRequest: func(info httptrace.WroteRequestInfo) {
+			s.serverStart = time.Now()
+
+			// When client doesn't use DialContext or using old (before go1.7) `net`
+			// package, DNS/TCP/TLS hook is not called.
+			if s.dnsStart.IsZero() && s.tcpStart.IsZero() {
+				now := s.serverStart
+				s.dnsStart = now
+				s.dnsDone = now
+				s.tcpStart = now
+				s.tcpDone = now
+			}
+
+			// When connection is re-used, DNS/TCP/TLS hook is not called.
+			if s.isReused {
+				now := s.serverStart
+				s.dnsStart = now
+				s.dnsDone = now
+				s.tcpStart = now
+				s.tcpDone = now
+				s.tlsStart = now
+				s.tlsDone = now
+			}
+
+			if s.isTLS {
+				return
+			}
+
+			s.TLSHandshake = s.tcpDone.Sub(s.tcpDone)
+			s.Pretransfer = s.Connect
+		},
+
+		GotFirstResponseByte: func() {
+			s.serverDone = time.Now()
+			s.ServerProcessing = s.serverDone.Sub(s.serverStart)
+			s.StartTransfer = s.serverDone.Sub(s.dnsStart)
+			s.transferStart = s.serverDone
+		},
+	})
+}