Files
httprunner/hrp/pkg/httpstat/main.go
2022-10-16 22:36:49 +08:00

277 lines
8.3 KiB
Go

// Package httpstat traces HTTP latency infomation (DNSLookup, TCP Connection and so on) on any golang HTTP request.
// It uses `httptrace` package.
// Inspired by https://github.com/tcnksm/go-httpstat and https://github.com/davecheney/httpstat
package httpstat
import (
"context"
"crypto/tls"
"fmt"
"net/http"
"net/http/httptrace"
"strconv"
"strings"
"sync"
"time"
"github.com/fatih/color"
"github.com/rs/zerolog/log"
)
const (
httpsTemplate = "\n" +
` DNS Lookup TCP Connection TLS Handshake Server Processing Content Transfer` + "\n" +
`[%s | %s | %s | %s | %s ]` + "\n" +
` | | | | |` + "\n" +
` namelookup:%s | | | |` + "\n" +
` connect:%s | | |` + "\n" +
` pretransfer:%s | |` + "\n" +
` starttransfer:%s |` + "\n" +
` total:%s` + "\n\n"
httpTemplate = "\n" +
` DNS Lookup TCP Connection Server Processing Content Transfer` + "\n" +
`[ %s | %s | %s | %s ]` + "\n" +
` | | | |` + "\n" +
` namelookup:%s | | |` + "\n" +
` connect:%s | |` + "\n" +
` starttransfer:%s |` + "\n" +
` total:%s` + "\n\n"
)
func fmta(d time.Duration) string {
return color.YellowString("%7dms", int(d.Milliseconds()))
}
func fmtb(d time.Duration) string {
return color.RedString("%-9s", strconv.Itoa(int(d.Milliseconds()))+"ms")
}
func grayscale(code color.Attribute) func(string, ...interface{}) string {
return color.New(code + 232).SprintfFunc()
}
func colorize(s string) string {
v := strings.Split(s, "\n")
v[0] = grayscale(16)(v[0])
return strings.Join(v, "\n")
}
func printf(format string, a ...interface{}) (n int, err error) {
return fmt.Fprintf(color.Output, format, a...)
}
// 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
// https or http
schema string
// connected network info
network, addr string
mux *sync.RWMutex // avoid data race
}
// 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]int64 {
return map[string]int64{
"DNSLookup": s.DNSLookup.Milliseconds(),
"TCPConnection": s.TCPConnection.Milliseconds(),
"TLSHandshake": s.TLSHandshake.Milliseconds(),
"ServerProcessing": s.ServerProcessing.Milliseconds(),
"ContentTransfer": s.ContentTransfer.Milliseconds(),
"NameLookup": s.NameLookup.Milliseconds(),
"Connect": s.Connect.Milliseconds(),
"Pretransfer": s.Pretransfer.Milliseconds(),
"StartTransfer": s.StartTransfer.Milliseconds(),
"Total": s.Total.Milliseconds(),
}
}
func (s *Stat) Print() {
if s.network != "" && s.addr != "" {
printf("\n%s %s: %s\n",
color.CyanString("Connected to"),
color.MagentaString(s.network),
color.BlueString(s.addr),
)
}
switch s.schema {
case "https":
printf(colorize(httpsTemplate),
fmta(s.DNSLookup), // dns lookup
fmta(s.TCPConnection), // tcp connection
fmta(s.TLSHandshake), // tls handshake
fmta(s.ServerProcessing), // server processing
fmta(s.ContentTransfer), // content transfer
fmtb(s.NameLookup), // namelookup
fmtb(s.Connect), // connect
fmtb(s.Pretransfer), // pretransfer
fmtb(s.StartTransfer), // starttransfer
fmtb(s.Total), // total
)
case "http":
printf(colorize(httpTemplate),
fmta(s.DNSLookup), // dns lookup
fmta(s.TCPConnection), // tcp connection
fmta(s.ServerProcessing), // server processing
fmta(s.ContentTransfer), // content transfer
fmtb(s.NameLookup), // namelookup
fmtb(s.Connect), // connect
fmtb(s.StartTransfer), // starttransfer
fmtb(s.Total), // total
)
}
log.Info().
Interface("httpstat(ms)", s.Durations()).
Msg("HTTP latency statistics")
}
// WithHTTPStat is a wrapper of httptrace.WithClientTrace.
// It records the time of each httptrace hooks.
func WithHTTPStat(req *http.Request, s *Stat) context.Context {
s.mux = new(sync.RWMutex)
s.schema = req.URL.Scheme
return httptrace.WithClientTrace(req.Context(), &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.DNSLookup
},
ConnectStart: func(network, addr string) {
s.network = network
s.addr = addr
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.mux.Lock()
defer s.mux.Unlock()
now := time.Now()
s.serverStart = now
// When client doesn't use DialContext, DNS/TCP/TLS hook is not called.
if s.dnsStart.IsZero() && s.tcpStart.IsZero() {
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 {
s.dnsStart = now
s.dnsDone = now
s.tcpStart = now
s.tcpDone = now
s.tlsStart = now
s.tlsDone = now
}
if s.isTLS { // https
return
}
// http
s.TLSHandshake = 0
s.Pretransfer = s.Connect
},
GotFirstResponseByte: func() {
s.mux.Lock()
defer s.mux.Unlock()
s.serverDone = time.Now()
s.ServerProcessing = s.serverDone.Sub(s.serverStart)
s.StartTransfer = s.serverDone.Sub(s.dnsStart)
s.transferStart = s.serverDone
},
})
}