Golang的context包詳解

context 包說明

說明：本文的用到的例子大部分來自context包。

概述

context 包定義了Context接口類型，它可以具有生命周期、取消/關(guān)閉的channel信號、請求域范圍的健值存儲功能。
因此可以用它來管理goroutine 的生命周期、或者與一個請求關(guān)聯(lián)，在functions之間傳遞等。

每個Context應(yīng)該視為只讀的，通過WithCancel、WithDeadline、WithTimeout和WithValue函數(shù)可以基于現(xiàn)有的一個Context（稱為父Context）派生出一個新的Context（稱為子Context）。
其中WithCancel、WithDeadline和WithTimeout函數(shù)除了返回一個派生的Context以外，還會返回一個與之關(guān)聯(lián)的CancelFunc類型的函數(shù)，用于關(guān)閉Context。

通過調(diào)用CancelFunc來關(guān)閉關(guān)聯(lián)的Context時，基于該Context所派生的Context也都會被關(guān)閉，并且會將自己從父Context中移除，停止和它相關(guān)的timer。
如果不調(diào)用CancelFunc，除非該Context的父Context調(diào)用對應(yīng)的CancelFunc，或者timer時間到，否則該Context和派生的Context就內(nèi)存泄漏了。

可以使用go vet工具來檢查所有control-flow路徑上使用的CancelFuncs。

應(yīng)用程序使用Context時，建議遵循如下規(guī)則：
1、不要將Context存儲為結(jié)構(gòu)體的字段，應(yīng)該通過函數(shù)來傳遞一個具體的Context。并且Context應(yīng)該放在第一個參數(shù)，變量名為ctx。比如

func DoSomething(ctx context.Context, arg Arg) error {
    // ... use ctx ...
}

2、即使函數(shù)允許，也不要傳遞nil。如果你不確定Context的使用，你可以傳遞context.TODO。
3、請求域參數(shù)應(yīng)該通過Context上的K/V方式傳遞，不要通過函數(shù)參數(shù)傳遞具體的請求域參數(shù)。

Context可能會在多個goroutines之間傳遞共享，它是例程安全的。

可以參考https://blog.golang.org/context 中的服務(wù)器使用例子。

包導(dǎo)入

在 go1.7 及以上版本 context 包被正式列入官方庫中，所以我們只需要import "context"就可以了，而在 go1.6 及以下版本，我們要 import "golang.org/x/net/context" 。

Context 接口

context.Context接口的定義如下：

// Context 的實現(xiàn)應(yīng)該設(shè)計為多例程安全的
type Context interface {
    // 返回代表該Context過期的時間，和表示deadline是否被設(shè)置的bool值。
    // 多次調(diào)用會返回相同的過期時間值，并不會因為時間流逝而變化
    Deadline() (deadline time.Time, ok bool)
    // 返回一個channel，關(guān)閉該channel就代表關(guān)閉該Context。返回nil代表該Context不需要被關(guān)閉。
    // 多次調(diào)用返回會返回相同的值。
    Done() <-chan struct{}

    // 如果Context未關(guān)閉，則返回nil。
    // 否則如果正常關(guān)閉,則返回Canceled，過期關(guān)閉則返回DeadlineExceeded，
    // 發(fā)生錯誤則返回對應(yīng)的error。
    // 多次調(diào)用返回相同的值。
    Err() error
    // 根據(jù)key從Context中獲取一個value，如果沒有關(guān)聯(lián)的值則返回nil。
    // 其中key是可以比較的任何類型。
    // 多次調(diào)用返回相同的值。
    Value(key interface{}) interface{}
}

空的Context的實現(xiàn)

context包內(nèi)部定義了許多Context接口的實現(xiàn)，其中最簡單的就是emptyCtx了。

// emptyCtx 不需要關(guān)閉，沒有任何鍵值對，也沒有過期時間。
type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

func (*emptyCtx) Done() <-chan struct{} {
    return nil
}

func (*emptyCtx) Err() error {
    return nil
}

func (*emptyCtx) Value(key interface{}) interface{} {
    return nil
}

func (e *emptyCtx) String() string {
    switch e {
    case background:
        return "context.Background"
    case todo:
        return "context.TODO"
    }
    return "unknown empty Context"
}

context包中有兩個emptyCtx的兩個實例，

var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

分別通過context.Background()和context.TODO()獲取。

Background()：通常用作初始的Context、測試、最基層的根Context。
TODO()：通常用作不清楚作用的Context，或者還未實現(xiàn)功能的場景。

WithValue

該函數(shù)的功能是，基于現(xiàn)有的一個Context，派生出一個新的Context，新的Context帶有函數(shù)指定的key和value。內(nèi)部的實現(xiàn)類型是valueCtx類型。

// key 必須是可比較的類型
func WithValue(parent Context, key, val interface{}) Context {
    if key == nil {
        panic("nil key")
    }
    if !reflect.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }
    return &valueCtx{parent, key, val}
}

type valueCtx struct {
    Context
    key, val interface{}
}

func (c *valueCtx) String() string {
    return fmt.Sprintf("%v.WithValue(%#v, %#v)", c.Context, c.key, c.val)
}

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key {
        return c.val
    }
    return c.Context.Value(key)
}

從上面的源碼可以看出，Context中的K/V存儲并不是利用map實現(xiàn)的，而是先查詢自身的一對鍵值，如果不匹配key，再向上層的Context查詢。

官方對key的使用建議：
為了避免不同包的context使用沖突，不建議直接使用string和其他內(nèi)建的類型作為key。而是自定義一個私有的key的類型，即小寫開頭的key類型。
并且最好定義一個類型安全的訪問器（返回具體類型的函數(shù)，或其他方式）。
比如：

package user

import "context"

// 定義 User 類型，準(zhǔn)備存儲到Context中
type User struct {...}

// 定義了一個未導(dǎo)出的私有key類型，避免和其他包的key沖突。
type key int

// 該key實例是為了獲取user.User而定義的，并且也是私有的。
// 用戶使用user.NewContext和user.FromContext，而避免直接使用該變量
var userKey key = 0

// 返回一個新的Context，包含了u *User值
func NewContext(ctx context.Context, u *User) context.Context {
    return context.WithValue(ctx, userKey, u)
}

// FromContext returns the User value stored in ctx, if any.
// 從Context中獲取*User值
func FromContext(ctx context.Context) (*User, bool) {
    u, ok := ctx.Value(userKey).(*User)
    return u, ok
}

為了避免分配內(nèi)存，key通常也可以基于struct{}類型定義不同的類型

 type userKeyType struct{}
 context.WithValue(ctx, userKeyType{}, u)

注意，每一個struct{}類型的變量實際上它們的地址都是一樣的，所以不會分配新的內(nèi)存，但是重新定義后的不同struct{}類型，分別賦值給interface{}后，interface{}變量將不相等，比如：

type A struct{}
type B struct{}
a, b := A{}, B{}
var ia, ib, ia2 interface{}
ia, ib, ia2 = a, b, A{}
fmt.Printf("%p, %p, %v, %v", &a, &b, ia == ib, ia == ia2)
// 0x11b3dc0, 0x11b3dc0, false, true

使用WithValue的例子：

type favContextKey string

f := func(ctx context.Context, k favContextKey) {
    if v := ctx.Value(k); v != nil {
        fmt.Println("found value:", v)
        return
    }
    fmt.Println("key not found:", k)
}

k := favContextKey("language")
ctx := context.WithValue(context.Background(), k, "Go")

f(ctx, k)
f(ctx, favContextKey("color"))

// Output:
// found value: Go
// key not found: color

可關(guān)閉的Context

context包為可關(guān)閉的Context定義了一個接口：

type canceler interface {
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}

該接口有兩個具體的實現(xiàn)，cancelCtx 和 timerCtx。

WithCancel

*cancelCtx由WithCancel返回，所以我們先看看WithCancel函數(shù)的定義：

// 從parent上派生出一個新的Context，并返回該和一個CancelFunc類型的函數(shù)
// 調(diào)用該cancel函數(shù)會關(guān)閉該Context，該Context對應(yīng)的從Done()返回的只讀channel也會被關(guān)閉。
// parent 對應(yīng)的cancel函數(shù)如果被調(diào)用，parent派生的Context和對應(yīng)的channel也都會被關(guān)閉。
// 當(dāng)某項任務(wù)的操作完成時，應(yīng)盡快關(guān)閉Context，以便回收Context關(guān)聯(lián)的資源。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    // 創(chuàng)建一個cancelCtx類型的實例
    c := newCancelCtx(parent)
    // 關(guān)聯(lián)該實例和父Context之間的關(guān)閉/取消關(guān)系，即關(guān)閉parent也關(guān)閉基于它派生的Context。
    propagateCancel(parent, &c)
    // 返回該實例和對應(yīng)的關(guān)閉函數(shù)
    return &c, func() { c.cancel(true, Canceled) }
}

func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

其中CancelFunc的定義也非常簡單：

// 調(diào)用該函數(shù)意味著要關(guān)閉Context, 結(jié)束相關(guān)的任務(wù)。
// 第一次調(diào)用后，之后再次調(diào)用將什么都不做。
type CancelFunc func()

具體的propagateCancel實現(xiàn)和cancelCtx實現(xiàn)如下：

// 關(guān)聯(lián)child和parent之間的關(guān)閉/取消關(guān)系，即關(guān)閉parent也關(guān)閉child。
func propagateCancel(parent Context, child canceler) {
    if parent.Done() == nil {
        return // parent 不需要關(guān)閉，則不需要關(guān)聯(lián)關(guān)系
    }
    // 找到最近的cancelCtx類型的祖先Context實例
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // p 已經(jīng)關(guān)閉，所以也關(guān)閉child
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

// 找到最近的cancelCtx類型或繼承cancelCtx類型的祖先Context實例
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    for {
        switch c := parent.(type) {
        case *cancelCtx:
            return c, true
        case *timerCtx:
            return &c.cancelCtx, true
        case *valueCtx:
            parent = c.Context
        default:
            return nil, false
        }
    }
}

// 從最近的cancelCtx類型的祖先Context中移除child
func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child)
    }
    p.mu.Unlock()
}

type canceler interface {
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}

// 代表已經(jīng)關(guān)閉的channel
var closedchan = make(chan struct{})

func init() {
    close(closedchan)
}

// 實現(xiàn)可關(guān)閉的Context，關(guān)閉時，也將關(guān)閉它的子Context
type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     chan struct{}         // created lazily, closed by first cancel call
    children map[canceler]struct{} // set to nil by the first cancel call
    err      error                 // set to non-nil by the first cancel call
}

func (c *cancelCtx) Done() <-chan struct{} {
    c.mu.Lock()
    if c.done == nil {
        c.done = make(chan struct{})
    }
    d := c.done
    c.mu.Unlock()
    return d
}

func (c *cancelCtx) Err() error {
    c.mu.Lock()
    defer c.mu.Unlock()
    return c.err
}

func (c *cancelCtx) String() string {
    return fmt.Sprintf("%v.WithCancel", c.Context)
}

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    if c.done == nil {
        c.done = closedchan
    } else {
        close(c.done)
    }
    for child := range c.children {
        // NOTE: acquiring the child's lock while holding parent's lock.
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        removeChild(c.Context, c)
    }
}

WithDeadline

理解了WithCancel，再理解WithDeadline并不難。
WithDeadline返回的是*timerCtx類型的Context，timerCtx繼承 cancelCtx

// WithDeadline 根據(jù)parent和deadline返回一個派生的Context。
// 如果parent存在過期時間，且已過期，則返回一個語義上等同于parent的派生Context。
// 當(dāng)?shù)竭_(dá)過期時間、或者調(diào)用CancelFunc函數(shù)關(guān)閉、或者關(guān)閉parent會使該函數(shù)返回的派生Context關(guān)閉。
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
    if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
        // parent 已經(jīng)過期，返回一個語義上等同于parent的派生Context
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  deadline,
    }
    propagateCancel(parent, c)
    d := time.Until(deadline)
    if d <= 0 {
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(true, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        c.timer = time.AfterFunc(d, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

// timerCtx繼承 cancelCtx，并且定義了過期時間。
// 當(dāng)關(guān)閉 timerCtx時會關(guān)閉timer和cancelCtx
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}

func (c *timerCtx) String() string {
    return fmt.Sprintf("%v.WithDeadline(%s [%s])", c.cancelCtx.Context, c.deadline, time.Until(c.deadline))
}

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        // 從父Context的children字段中移除自己
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

WithTimeout

與WithDeadline稍由不同，WithTimeout傳遞的是一個超時時間間隔，而WithDeadline傳遞的是一個具體的過期時間。一般情況下WithTimeout比WithDeadline更常用：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

例子：

func slowOperationWithTimeout(ctx context.Context) (Result, error) {
    ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
    defer cancel()  // 如果slowOperation在超時之前完成，則需要調(diào)用cancel關(guān)閉Context，以便回收Context相關(guān)聯(lián)的資源
    return slowOperation(ctx)
}

使用例子

最后，我們再舉幾個例子來加深印象。

WithCancel的例子

如何使用Context讓goroutine退出，避免goroutine內(nèi)存泄漏。

// 在獨立的goroutine中生成整數(shù)，通過channel傳遞出去。
// 一旦context關(guān)閉，該goroutine也將安全退出。
gen := func(ctx context.Context) <-chan int {
    dst := make(chan int)
    n := 1
    go func() {
        for {
            select {
            case <-ctx.Done():
                return // returning not to leak the goroutine
            case dst <- n:
                n++
            }
        }
    }()
    return dst
}

ctx, cancel := context.WithCancel(context.Background())
defer cancel() // 當(dāng)完成整數(shù)生產(chǎn)時，關(guān)閉Context

for n := range gen(ctx) {
    fmt.Println(n)
    if n == 5 {
        break
    }
}
// Output:
// 1
// 2
// 3
// 4
// 5

WithDeadline的例子

當(dāng)某項任務(wù)需要在規(guī)定的時間內(nèi)完成，如果未完成則需要立即取消任務(wù)，并且返回錯誤的情況，可以使用WithDeadline。

    d := time.Now().Add(50 * time.Millisecond)
    ctx, cancel := context.WithDeadline(context.Background(), d)

    // 即使ctx會因為過期而關(guān)閉，我們也應(yīng)該在最后調(diào)用cancel，因為任務(wù)可能會在規(guī)定時間內(nèi)完成，這種情況需要主動調(diào)用cancel來盡快釋放Context資源
    defer cancel()

    select {
    case <-time.After(1 * time.Second):
        fmt.Println("overslept")
    case <-ctx.Done():
        fmt.Println(ctx.Err())
    }

    // Output:
    // context deadline exceeded

WithTimeout例子

同WithDeadline，只是傳遞的是一個time.Duration

ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
    defer cancel()

    select {
    case <-time.After(1 * time.Second):
        fmt.Println("overslept")
    case <-ctx.Done():
        fmt.Println(ctx.Err()) // prints "context deadline exceeded"
    }

    // Output:
    // context deadline exceeded

WithValue的例子

WithValue在前面的介紹中已經(jīng)舉過一個例子，我們再舉一個http.Request相關(guān)的Context的例子。
在 Golang1.7 中，"net/http"原生支持將Context嵌入到 *http.Request中，并且提供了http.Request.Conext() 和 http.Request.WithContext(context.Context)這兩個函數(shù)。
http.Request.Conext()函數(shù)返回或新建一個 context。
http.Request.WithContext(context.Context)函數(shù)返回一個新的Request，并且將傳入的context與新的Reuest實例關(guān)聯(lián)。

type userKey string

const userIDKey userKey = "uid"

func requestFilter(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
        cook, err := req.Cookie("USERID")
        uid := ""
        if err == nil {
            uid = cook.Value
        }
        ctx := context.WithValue(req.Context(), userIDKey, uid)
        next.ServeHTTP(w, req.WithContext(ctx))
    })
}
func userIdFromContext(ctx context.Context) string {
    return ctx.Value(userIDKey).(string)
}
func process(w http.ResponseWriter, req *http.Request) {
    uid := userIdFromContext(req.Context())
    fmt.Fprintln(w, "user ID is ", uid)
    return
}
func main() {
    http.Handle("/", requestFilter(http.HandlerFunc(process)))
    http.ListenAndServe(":8080", nil)
}

結(jié)束

本文詳解了context包源碼的實現(xiàn)，并結(jié)合了一些例子來說明用法。相信讀者對context包已經(jīng)有一定的理解了，還可以再看看源碼包中完整的代碼，來加深印象哦。