簡介

這篇文章主要是通過官方提供的 HTTP 追蹤來學(xué)習(xí)使用 Hook 的編程思想。

簡單來說 Hook 的編程思想跟事件驅(qū)動是類似的，通過預(yù)先保存一些要執(zhí)行函數(shù)或方法，在滿足某些條件的時候自動執(zhí)行。
在了解使用 Go 語言編寫 Hook 之前，最好先掌握 Context 的用法，關(guān)于 Context 的用法，請?jiān)斠奫另外一篇文章][]。

官方提供的 net/http/httptrace主要是用于追蹤客戶端的 Request 請求過程中發(fā)生的各種事件及行為，在標(biāo)準(zhǔn)庫 net/http/httptrace/trace.go 中定義了一個叫 ClientTrace 的結(jié)構(gòu)體，它包含了一系列的鉤子函數(shù) hooks 作為成員變量，如下：

// ClientTrace is a set of hooks to run at various stages of an outgoing HTTP request. 
type ClientTrace struct {
    GetConn func(hostPort string)
    GotConn func(GotConnInfo)
    PutIdleConn func(err error)

    GotFirstResponseByte func()
    Got100Continue func()

    DNSStart func(DNSStartInfo)
    DNSDone func(DNSDoneInfo)

    ConnectStart func(network, addr string)
    ConnectDone func(network, addr string, err error)

    TLSHandshakeStart func()
    TLSHandshakeDone func(tls.ConnectionState, error)
    WroteHeaders func()
    Wait100Continue func()
    WroteRequest func(WroteRequestInfo)
}

trace.go 還提供了一個 WithClientTrace() 包函數(shù)，用來把 ClientTrace 結(jié)構(gòu)體中的鉤子都保存（注冊）到 Context 中去（因?yàn)?Context 提供 key/value 存儲嘛），
key 就是一個叫 clientEventContextKey 的空結(jié)構(gòu)體，value 是 nettrace 包中的 Trace 結(jié)構(gòu)體，這個結(jié)構(gòu)體作用跟 ClientTrace 一樣，都是包含了一堆 hook 函數(shù)作為成員，
在這里它的目的只是封裝下 ClientTrace 中的 hook 函數(shù)。最終， WithClientTrace() 會返回一個 context，它保存了上述的 hook 函數(shù)。

type clientEventContextKey struct{}
func WithClientTrace(ctx context.Context, trace *ClientTrace) context.Context {
    if trace == nil {
        panic("nil trace")
    }
    old := ContextClientTrace(ctx)
    trace.compose(old)

    ctx = context.WithValue(ctx, clientEventContextKey{}, trace)
    if trace.hasNetHooks() {
        nt := &nettrace.Trace{
            ConnectStart: trace.ConnectStart,
            ConnectDone:  trace.ConnectDone,
        }
        if trace.DNSStart != nil {
            nt.DNSStart = func(name string) {
                trace.DNSStart(DNSStartInfo{Host: name})
            }
        }
        if trace.DNSDone != nil {
            ...
        }
        ctx = context.WithValue(ctx, nettrace.TraceKey{}, nt)
    }
    return ctx
}

通過 ContextClientTrace() 的函數(shù)，可以把 ClientTrace 從 Context 中取出來。

// ContextClientTrace returns the ClientTrace associated with the
// provided context. If none, it returns nil.
func ContextClientTrace(ctx context.Context) *ClientTrace {
    trace, _ := ctx.Value(clientEventContextKey{}).(*ClientTrace)
    return trace
}

現(xiàn)在，我們知道，有了 WithClientTrace()，我們就可以把鉤子函數(shù)保存在 Context 中了，現(xiàn)在，我們要把這些鉤子函數(shù)掛到 Request 中去，該怎么弄？
很簡單，通過 Request.WithContext() 把剛才賦值好的 Context 保存到 Request 中就可以了。

現(xiàn)在 Request 有了這些鉤子函數(shù)，那么什么時候會被調(diào)用呢？ 當(dāng)然會 http.Client.Do(req) 的時候啦。

接下來我們通過一段實(shí)際的代碼看看整個流程：

package main

import (
    "fmt"
    "log"
    "net/http"
    "net/http/httptrace"
)

// transport is an http.RoundTripper that keeps track of the in-flight
// request and implements hooks to report HTTP tracing events.
type transport struct {
    current *http.Request
}

// RoundTrip wraps http.DefaultTransport.RoundTrip to keep track
// of the current request.
func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
    t.current = req
    return http.DefaultTransport.RoundTrip(req)
}

// GotConn prints whether the connection has been used previously
// for the current request.
func (t *transport) GotConn(info httptrace.GotConnInfo) {
    fmt.Printf("Connection reused for %v? %v\n", t.current.URL, info.Reused)
}

func main() {
    t := &transport{}

    req, _ := http.NewRequest("GET", "https://google.com", nil)
    trace := &httptrace.ClientTrace{
        GotConn: t.GotConn,
    }
    req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))

    client := &http.Client{Transport: t}
    if _, err := client.Do(req); err != nil {
        log.Fatal(err)
    }
}

所有的鉤子的調(diào)用，最終都會在 client.Do(req) 里面執(zhí)行，我們看看是怎么執(zhí)行的。

注意到這里的 transport 結(jié)構(gòu)體，它其實(shí)是 RoundTripper 接口類型（在 client.go 中聲明）的一個 implementer，這個 RoundTripper 實(shí)際只有一個方法：

// RoundTripper is an interface representing the ability to execute a
// single HTTP transaction, obtaining the Response for a given Request.
type RoundTripper interface {
    // RoundTrip executes a single HTTP transaction, returning
    // a Response for the provided Request.
    RoundTrip(*Request) (*Response, error)
}

在 client.Do() 中，會調(diào)用 client.send()，如下：

 resp, didTimeout, err = c.send(req, deadline)

c.send() 內(nèi)部：

// didTimeout is non-nil only if err != nil.
func (c *Client) send(req *Request, deadline time.Time) (resp *Response, didTimeout func() bool, err error) {
    ...
    resp, didTimeout, err = send(req, c.transport(), deadline)
    ...
    return resp, nil, nil
}

send() 內(nèi)部：

func send(ireq *Request, rt RoundTripper, deadline time.Time) (resp *Response, didTimeout func() bool, err error) {
    ...
    resp, err = rt.RoundTrip(req)
    ...
}

可見，最終調(diào)用了 rt.RoundTrip() 函數(shù)。也就是上述 main.go 中 transport 實(shí)現(xiàn)的 RoundTrip() 函數(shù)。

在 rt.RoundTrip() 里面，把 req 賦給了 DefaultTransport.RoundTrip(req)，
這個 DefaultTransport 是包提供的一個 RoundTripper 的默認(rèn)實(shí)現(xiàn)，

var DefaultTransport RoundTripper = &Transport{
    Proxy: ProxyFromEnvironment,
    DialContext: (&net.Dialer{
        Timeout:   30 * time.Second,
        KeepAlive: 30 * time.Second,
        DualStack: true,
    }).DialContext,
    MaxIdleConns:          100,
    IdleConnTimeout:       90 * time.Second,
    TLSHandshakeTimeout:   10 * time.Second,
    ExpectContinueTimeout: 1 * time.Second,
}

然后，在它的 RoundTrip() 函數(shù)里面最終會調(diào)用上述的鉤子函數(shù)。

// RoundTrip implements the RoundTripper interface.
//
// For higher-level HTTP client support (such as handling of cookies
// and redirects), see Get, Post, and the Client type.
func (t *Transport) RoundTrip(req *Request) (*Response, error) {
    t.nextProtoOnce.Do(t.onceSetNextProtoDefaults)
    ctx := req.Context()
    trace := httptrace.ContextClientTrace(ctx)
    
    for {
        treq := &transportRequest{Request: req, trace: trace}
        cm, err := t.connectMethodForRequest(treq)
        ...
        pconn, err := t.getConn(treq, cm)
    }
}

解析：

通過調(diào)用 httptrace.ContextClientTrace(ctx) 把 context 中的鉤子函數(shù)都取出來，再在 t.getConn() 中調(diào)用鉤子函數(shù)，如下：