Golang Context 內部原理一覽

context[1] 包是每個 Go 語言程序員在編程生涯中都會使用到的, 但你們真的瞭解它的內部工作原理嗎 (因爲據說真正的美麗都藏在內在)?

如果你對此感興趣, 可以查看我之前關於 sync.WaitGroup 內部工作原理的博文鏈接 [2] 。

本文旨在簡要概覽內部工作原理, 因此我會適當簡化一些細節。如果你想查看全部的實現細節, 或者你更喜歡直接閱讀源代碼而不是博客, 可以點擊這裏 [3] 。實際上, 它的實現相當簡單。

基本用法 #

你們都知道 context包的基本用法, 如果不知道, 那這篇博文可能還不太適合你。不過爲了避免突兀, 我們還是看一個簡單的使用 context包的例子吧。

我真的不知道該拿什麼好的例子來展示。經過一番思考, 我決定放一個接收 context並將其傳遞給另一個函數的函數, 因爲對大多數人來說, context就是這樣被傳來傳去的。

func main() {
    bigFunc(context.Background())
}
func bigFunc(ctx context.Context) {
    smallFunc(ctx)
}
func smallFunc(ctx context.Context) {
    // I don't know what to do with it, let' just print it
    fmt.Println(ctx)
}

如果你運行這段代碼 [4] , 將會打印出 context.Background。這是因爲 context.Background返回的值滿足了 Stringer接口, 調用 String方法時會返回這個字符串。

前戲到此爲止, 讓我們正式開始吧。

`Context`接口 #

讓我們從基礎開始。你所使用的 context.Context[5] 是一個接口, 下面是它的定義。

type Context interface {
    Deadline() (deadline time.Time, ok bool) // get the deadline time
    Done() <-chan struct{}                   // get a channel which is closed when cancelled
    Err() error                              // returns non-nil if Done channel is closed
    Value(key any) any                       // get a value from the context store
}

任何滿足這個接口的結構體都是一個有效的上下文對象。讓我們快速瀏覽一下每個方法的作用, 以防註釋沒有解釋清楚。

Deadline：返回設置爲截止時間的時間, 例如使用 context.WithDeadline創建的上下文。
Done：返回一個通道, 在上下文被取消時關閉。
Err：如果上下文已被取消, 返回非 nil 錯誤值。
Value：允許你從上下文實例中獲取存儲的值。

如果你想構建一個 "上下文", 這就是你所需要的全部。當然, 標準庫也爲我們提供了一些有用的實現。

`emptyCtx`結構體 #

這是一個滿足最基本要求的上下文結構體。下面是它的代碼:

type emptyCtx struct{}
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 any) any {
    return nil
}

正如你所見, 它什麼也不做, 但這基本上就是 context.Background和 context.TODO的全部內容。

`context.Background`和 `context.TODO` #

它們只是 emptyCtx加上一個滿足 Stringer接口的 String方法。它們爲你提供了一種創建空上下文的方式。它們唯一的區別就是名稱。

當你知道需要一個空上下文時, 你會使用 context.Background, 比如在主函數中剛開始時。而當你不確定應該使用哪個上下文或者還沒有連接好時, 你會使用 context.TODO。

你可以將 context.TODO看作是在代碼中添加 // TODO註釋。

下面是 context.Background的代碼:

type backgroundCtx struct{ emptyCtx }
func (backgroundCtx) String() string {
    return "context.Background"
}
func Background() Context {
    return backgroundCtx{}
}

以及 context.TODO的代碼:

type todoCtx struct{ emptyCtx }
func (todoCtx) String() string {
    return "context.TODO"
}
func TODO() Context {
    return todoCtx{}
}

很簡單, 對吧?

`context.WithValue` #

現在我們來看看 context包更有用的用例。如果你想使用 context傳遞一個值, 可以使用 context.WithValue。你可能見過一些日誌或 Web 框架使用這種方式。

讓我們看看它的內部實現:

type valueCtx struct {
    Context
    key, val any
}
func WithValue(parent Context, key, val any) Context {
    return &valueCtx{parent, key, val}
}

它只是返回了一個包含父上下文、鍵和值的結構體。

如果你注意到, 實例只能保存一個鍵和一個值, 但你可能在 Web 框架中看到它們從 ctx參數中取出多個值。由於你將父級嵌入到了新的上下文中, 因此可以遞歸向上搜索以獲取任何其他值。

假設你創建瞭如下內容鏈接 [6] :

bgCtx := context.Background()
v1Ctx := context.WithValue(bgCtx, "one", "uno")
v2Ctx := context.WithValue(v1Ctx, "two", "dos")

現在, 如果我們要從 v2Ctx獲取 "one" 的值, 我們可以調用 v2Ctx.Value("one")。這將首先檢查 v2Ctx中的 key是否爲 "one", 由於不是, 它將檢查父級 ( v1Ctx) 的 key是否爲 "one"。現在由於 v1Ctx中的 key是 "one", 我們返回上下文中的值。

代碼如下所示。我已經刪除了一些關於如何處理超時 / 取消值的部分。你可以在這裏 [7] 查看完整代碼。

func (c *valueCtx) Value(key any) any {
    // If it this one, just return it
    if c.key == key {
        return c.val
    }
    return value(c.Context, key)
}
func value(c Context, key any) any {
    for {
        switch ctx := c.(type) {
        case *valueCtx:
            // If the parent is a `valueCtx`, check its key
            if key == ctx.key {
                return ctx.val
            }
            c = ctx.Context
        case backgroundCtx, todoCtx:
            // If we have reached the top, ie base context
            // Return as we did not find anything
            return nil
        default:
            // If it is some other context,
            // just calls its `.Value` method
            return c.Value(key)
        }
    }
}

正如你所看到的, 代碼只是遞歸地搜索父級上下文, 看是否有任何一個與鍵匹配, 如果有則返回其值。

`context.WithCancel` #

讓我們看一些更有用的東西。你可以使用 context包創建一個 ctx, 用於向下遊函數發出取消信號。

讓我們看一個示例 [8] 來了解如何使用:

func doWork(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            return
        default:
        }
        time.Sleep(1 * time.Second)
        fmt.Println("doing work...")
    }
}
func main() {
    bgCtx := context.Background()
    innerCtx, cancel := context.WithCancel(bgCtx)
    go doWork(innerCtx) // call goroutine
    time.Sleep(3 * time.Second) // do work in main
    // well, if `doWork` is still not done, just cancel it
    cancel()
}

在這種情況下, 你可以通過在 main函數中調用 cancel來通知 doWork函數停止工作。

現在讓我們看看它是如何工作的。讓我們從函數定義開始 (我們稍後會討論結構體定義):

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := &cancelCtx{}
    c.propagateCancel(parent, c)
    return c, func() { c.cancel(true, Canceled, nil) }
}

當你調用 context.WithCancel時, 它會返回一個 cancelCtx實例和一個可以用來取消上下文的函數。僅從這一點, 我們就可以推斷出 cancelCtx是一個具有 cancel函數的上下文, 該函數可用於 "取消" 上下文。

如果你忘記了, 取消上下文只是意味着你關閉了由 Done()返回的通道。

順便說一下, propagateCancel函數在這裏的主要作用是創建一個 cancelCtx, 並確保在創建之前父級未被取消。

好的, 現在讓我們看一下結構體, 之後我們將討論它是如何工作的 ( 源碼 [9] )。

type cancelCtx struct {
    Context
    mu       sync.Mutex            // protects following fields
    done     atomic.Value          // chan struct{} for Done
    children map[canceler]struct{}
    err      error
    cause    error
}

好的, 我們在這裏有什麼:

Context: 保存父級上下文
mu sync.Mutex: 你知道這是什麼
doneatomic.Value: 保存將由 Done()函數返回的 chanstruct{}
err error: 保存導致取消的錯誤
cause error: 保存取消的原因, 即 cancel函數的最後一個參數

順便說一下, 這是取消函數 ( 源碼 [10] ):

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
    if cause == nil {
        cause = err
    }
    c.err = err
    c.cause = cause
    // load the chan struct{} from atomic.Value and close it
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        // if it does not exist, store a closed channel
        c.done.Store(closedchan)
    } else {
        // if it exists, close it
        close(d)
    }
    // call cancel on all children to propagate the cancellation
    for child := range c.children {
        child.cancel(false, err, cause)
    }
    c.children = nil
    // remove itself from the parent
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

首先我們設置實例上的原因和錯誤, 然後關閉 Done返回的通道。之後它取消所有子級, 最後將自己從父級中移除。

最後, `context.WithDeadline`和 `context.WithTimeout` #

當你想創建一個在達到截止時間時自動取消自身的上下文時, 這些就很有用了。這對於強制執行諸如服務器超時之類的事情非常有用。

先把 context.WithTimeout說清楚, 它只是計算截止時間並調用 context.WithDeadline。事實上, 這就是它的全部代碼:

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

現在讓我們深入瞭解細節。正如你們中的一些人可能已經猜到的那樣, WithDeadline基本上只是一個普通的 WithCancel上下文, 但上下文包處理取消。

讓我們看看這段代碼的作用。這裏我有一個名爲 WithDeadlineCause 的函數代碼, 它是 WithDeadline 的一個變體, 但增加了傳入 "取消原因" 的能力。順便說一下, Cause 變體也可用於 context 包中的其他函數, 而像 WithDeadline 這樣的非 Cause 版本只是在調用 Cause 變體時將 cause 設置爲 nil。

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // if parent deadline earlier than child,
        // just return a cancelCtx
        return WithCancel(parent)
    }
    // create a new timerCtx
    c := &timerCtx{
        deadline: d,
    }
    c.cancelCtx.propagateCancel(parent, c)
    dur := time.Until(d)
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
        return c, func() { c.cancel(false, Canceled, nil) }
    }
    // if all good, setup a new timer and return a cancel func
    if c.err == nil {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded, cause)
        })
    }
    return c, func() { c.cancel(true, Canceled, nil) }
}

讓我們逐步瞭解它的工作原理:

如果父級截止時間早於子級, 則返回從父級創建的簡單 cancelCtx
如果不是, 則創建一個新的 timeCtx (結構定義如下)
現在檢查截止時間是否已經過期, 如果是, 則取消創建的上下文並返回
如果沒有, 我們使用 time.AfterFunc[11] 設置一個計時器來執行取消操作並返回

順便說一下, timerCtx 只是一個帶有計時器的 cancelCtx:

type timerCtx struct {
    cancelCtx
    timer *time.Timer
    deadline time.Time
}

就是這樣。正如我開始時提到的, 我已經從原始代碼中刪除了一些部分, 以便更容易解釋。實際源碼 [2] 非常易讀, 我鼓勵你去閱讀它。

← 主頁 [12]

參考鏈接

context: https://pkg.go.dev/context
鏈接: https://blog.meain.io/2024/sync-waitgroup-internals/
點擊這裏: https://cs.opensource.google/go/go/+/master:src/context/context.go
這段代碼: https://go.dev/play/p/NRn47sZ7xmk
context.Context: https://cs.opensource.google/go/go/+/refs/tags/go1.22.4:src/context/context.go;l=64-160
鏈接: https://go.dev/play/p/zD7kdqxzen1
這裏: https://cs.opensource.google/go/go/+/refs/tags/go1.22.4:src/context/context.go;l=759-790
示例: https://go.dev/play/p/nmvOh0YPmde
源碼: https://cs.opensource.google/go/go/+/refs/tags/go1.22.4:src/context/context.go;l=419-429
源碼: https://cs.opensource.google/go/go/+/refs/tags/go1.22.4:src/context/context.go;l=533-566
time.AfterFunc: https://pkg.go.dev/time#AfterFunc
← 主頁: https://blog.meain.io/

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/vIObM9YO4r2KPRat9a16Ew

基本用法 #

Context接口 #

emptyCtx結構體 #

context.Background和 context.TODO #

context.WithValue #

context.WithCancel #

最後, context.WithDeadline和 context.WithTimeout #