Go context源碼解析

這是一個建立於 的文章，其中的資訊可能已經有所發展或是發生改變。

在上一篇文章 golang context初探 中，已經初步瞭解了context的用法以及應用的情境。那麼接下來深入到源碼中來學習一下context是怎麼實現的。

emptyCtx

context包的代碼很少，一個context.go檔案，總共才480行代碼，其中還包括大量的注釋。context包首先定義了一個Context介面：

type Context interface {    Deadline() (deadline time.Time, ok bool)    Done() <-chan struct{}    Err() error    Value(key interface{}) interface{}}

接下來定義了一個emptyCtx類型：

// An emptyCtx is never canceled, has no values, and has no deadline. It is not// struct{}, since vars of this type must have distinct addresses.type emptyCtx int

為什麼叫做emptyCtx呢？注釋說了emptyCtx不能被取消，沒有值，也沒有deadline。同時，emptyCtx也不是一個struct{}，因為這種類型地變數需要有不同的地址。
這個emptyCtx實現了Context介面：

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

看了上面這段代碼就知道為什麼emptyCtx不能被取消，沒有值，也沒有deadline了，因為上面的實現都是直接return。那麼，這個emptyCtx有什麼用呢？還記得Background()和TODO()函數嗎？對的，它們的內部就是直接返回emptyCtx類型的指標：

var (    background = new(emptyCtx)    todo       = new(emptyCtx))func Background() Context {    return background}func TODO() Context {    return todo}

所以這兩個函數一般是用在main函數、初始化、測試以及頂層請求的Context。OK，繼續往下看。
既然emptyCtx類型什麼都不做，那麼應該有其他的類型來實現相關的功能才對，也就是cancelCtx，timerCtx，valueCtx三種類型。下面來講一下這三種類型。

Cancel

在之前的文章中我們知道，WithCancel，WithTimeout，WithDeadline這三個方法會返回一個CancelFunc類型的函數，在Context內部就定義了canceler介面：

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

canceler是一種可以直接被取消的context類型，待會繼續往下看我們會發現，cancelCtx和timerCtx不單單實現了Context介面（通過匿名成員變數），也實現了canceler介面。

cancelCtx

cancelCtx的結構體定義：

type cancelCtx struct {    Context    done chan struct{} // closed by the first cancel call.    mu       sync.Mutex    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{} {    return c.done}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    // 關閉c的done channel，所有監聽c.Done()的goroutine都會收到訊息    close(c.done)    // 取消child，由於是map結構，所以取消的順序是不固定的    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()    // 從c.children中移除取消過的child    if removeFromParent {        removeChild(c.Context, c)    }}

當我們在調用WithCancel的時候，實際上返回的就是一個cancelCtx指標和cancel()方法：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {    c := newCancelCtx(parent)    propagateCancel(parent, &c)    return &c, func() { c.cancel(true, Canceled) }}// newCancelCtx returns an initialized cancelCtx.func newCancelCtx(parent Context) cancelCtx {    return cancelCtx{        Context: parent,        done:    make(chan struct{}),    }}

那麼，propagateCancel函數又是幹什麼的呢？

// 向上找到最近的可以被取消的父context，將子context放入parent.Children中// propagateCancel arranges for child to be canceled when parent is.func propagateCancel(parent Context, child canceler) {    if parent.Done() == nil {        return // parent is never canceled    }    // 判斷返回的parent是否是cancelCtx    if p, ok := parentCancelCtx(parent); ok {        p.mu.Lock()        if p.err != nil {            // parent has already been canceled            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():            }        }()    }}// parentCancelCtx follows a chain of parent references until it finds a// *cancelCtx. This function understands how each of the concrete types in this// package represents its parent.// 不停地向上尋找最近的可取消的父contextfunc 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        }    }}

Timer

WithTimeout和WithDeadline其實是差不多的，只是源碼內部幫我們封裝了一下：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {    return WithDeadline(parent, time.Now().Add(timeout))}func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {    // 當前的deadline比新的deadline還要早，直接返回    if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {        // The current deadline is already sooner than the new one.        return WithCancel(parent)    }    c := &timerCtx{        cancelCtx: newCancelCtx(parent),        deadline:  deadline,    }    propagateCancel(parent, c)    d := time.Until(deadline)    // 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 {        // 設定d時間後執行取消方法        c.timer = time.AfterFunc(d, func() {            c.cancel(true, DeadlineExceeded)        })    }    return c, func() { c.cancel(true, Canceled) }}

timerCtx的代碼也實現地比較簡潔：

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 {        // Remove this timerCtx from its parent cancelCtx's children.        removeChild(c.cancelCtx.Context, c)    }    c.mu.Lock()    if c.timer != nil {        c.timer.Stop()        c.timer = nil    }    c.mu.Unlock()}

這裡需要注意的是，timerCtx並沒有直接實現了canceler介面，而是使用了匿名成員變數，這樣就可以不用重頭實現一遍Context介面，而是按需只實現了Deadline方法。timerCtx的cancel方法先調用了cancelCtx的cancel方法，然後再去停止定時器。

Value

先來看看valueCtx的定義：

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實現了，在匿名變數Context之外，還增加了兩個key，value變數來儲存值。看看WithValue的實現：

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

需要注意的是，key不能是nil並且必須是可比較的，否則就會導致panic！可比較的意思是key不能為函數類型或者NaN之類的，具體可以看一下reflect包，這裡就不細說了。

最後

整個context看下來，對context是怎麼實現的也有了清晰的瞭解，而且context包有大量的測試代碼，非常棒！最近在看go的源碼，發現真的是一個很好的學習材料，對於如何寫出簡潔明了的代碼是很有協助的。