Golang的context包詳解

context 包說明

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

概述

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

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

通過調用CancelFunc來關閉關聯的Context時，基于該Context所派生的Context也都會被關閉，并且會將自己從父Context中移除，停止和它相關的timer。
如果不調用CancelFunc，除非該Context的父Context調用對應的CancelFunc，或者timer時間到，否則該Context和派生的Context就內存泄漏了。

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

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

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

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

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

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

包導入

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

Context 接口

context.Context接口的定義如下：

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

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

空的Context的實現

context包內部定義了許多Context接口的實現，其中最簡單的就是emptyCtx了。

// emptyCtx 不需要關閉，沒有任何鍵值對，也沒有過期時間。
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，或者還未實現功能的場景。

WithValue

該函數的功能是，基于現有的一個Context，派生出一個新的Context，新的Context帶有函數指定的key和value。內部的實現類型是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實現的，而是先查詢自身的一對鍵值，如果不匹配key，再向上層的Context查詢。

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

package user

import "context"

// 定義 User 類型，準備存儲到Context中
type User struct {...}

// 定義了一個未導出的私有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
}

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

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

注意，每一個struct{}類型的變量實際上它們的地址都是一樣的，所以不會分配新的內存，但是重新定義后的不同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

可關閉的Context

context包為可關閉的Context定義了一個接口：

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

該接口有兩個具體的實現，cancelCtx 和 timerCtx。

WithCancel

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

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

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

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

// 調用該函數意味著要關閉Context, 結束相關的任務。
// 第一次調用后，之后再次調用將什么都不做。
type CancelFunc func()

具體的propagateCancel實現和cancelCtx實現如下：

// 關聯child和parent之間的關閉/取消關系，即關閉parent也關閉child。
func propagateCancel(parent Context, child canceler) {
    if parent.Done() == nil {
        return // parent 不需要關閉，則不需要關聯關系
    }
    // 找到最近的cancelCtx類型的祖先Context實例
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // p 已經關閉，所以也關閉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{}
}

// 代表已經關閉的channel
var closedchan = make(chan struct{})

func init() {
    close(closedchan)
}

// 實現可關閉的Context，關閉時，也將關閉它的子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 根據parent和deadline返回一個派生的Context。
// 如果parent存在過期時間，且已過期，則返回一個語義上等同于parent的派生Context。
// 當到達過期時間、或者調用CancelFunc函數關閉、或者關閉parent會使該函數返回的派生Context關閉。
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
    if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
        // parent 已經過期，返回一個語義上等同于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，并且定義了過期時間。
// 當關閉 timerCtx時會關閉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在超時之前完成，則需要調用cancel關閉Context，以便回收Context相關聯的資源
    return slowOperation(ctx)
}

使用例子

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

WithCancel的例子

如何使用Context讓goroutine退出，避免goroutine內存泄漏。

// 在獨立的goroutine中生成整數，通過channel傳遞出去。
// 一旦context關閉，該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() // 當完成整數生產時，關閉Context

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

WithDeadline的例子

當某項任務需要在規定的時間內完成，如果未完成則需要立即取消任務，并且返回錯誤的情況，可以使用WithDeadline。

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

    // 即使ctx會因為過期而關閉，我們也應該在最后調用cancel，因為任務可能會在規定時間內完成，這種情況需要主動調用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在前面的介紹中已經舉過一個例子，我們再舉一個http.Request相關的Context的例子。
在 Golang1.7 中，"net/http"原生支持將Context嵌入到 *http.Request中，并且提供了http.Request.Conext() 和 http.Request.WithContext(context.Context)這兩個函數。
http.Request.Conext()函數返回或新建一個 context。
http.Request.WithContext(context.Context)函數返回一個新的Request，并且將傳入的context與新的Reuest實例關聯。

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)
}

結束

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