Golang Context 包詳解

0. 引言

在 Go 語言編寫的伺服器程式中，伺服器通常要為每個 HTTP 請求建立一個 goroutine 以併發地處理業務。同時，這個 goroutine 也可能會建立更多的 goroutine 來訪問資料庫或者 RPC 服務。
當這個請求超時或者被終止的時候，需要優雅地退出所有衍生的 goroutine，並釋放資源。因此，我們需要一種機制來通知衍生 goroutine 請求已被取消。 比如以下例子，sleepRandom_1 的結束就無法通知到 sleepRandom_2。

package main

import (
    "fmt"
    "time"
)

func sleepRandom_1() {
    i := 0
    for {
        time.Sleep(1 * time.Second)
        fmt.Printf("This is sleep Random 1: %d\n", i)

        i++
        if i == 5 {
            fmt.Println("cancel sleep random 1")
            break
        }
    }
}

func sleepRandom_2() {
    i := 0
    for {
        time.Sleep(1 * time.Second)
        fmt.Printf("This is sleep Random 2: %d\n", i)
        i++
    }
}

func main() {

    go sleepRandom_1() // 迴圈 5 次後退出
    go sleepRandom_2() // 會一直列印 This is sleep Random 2

    for {
        time.Sleep(1 * time.Second)
        fmt.Println("Continue...")
    }
}
 

1. Context

Context 包提供上下文機制在 goroutine 之間傳遞 deadline、取消訊號（cancellation signals）或者其他請求相關的資訊。使用方法是：

1. 首先，伺服器程式為每個接受的請求建立一個 Context 例項（稱為根 context，通過 context.Background() 方法建立）；
2. 之後的 goroutine 接受根 context 的一個派生 Context 物件。比如通過呼叫根 context 的 WithCancel 方法，建立子 context；
3. goroutine 通過 context.Done() 方法監聽取消訊號。func Done() <-chan struct{}
   是一個通訊操作，會阻塞 goroutine，直到收到取消訊號接觸阻塞。
   （可以藉助 select 語句，如果收到取消訊號，就退出 goroutine；否則，預設子句是繼續執行 goroutine）；
4. 當一個 Context 被取消（比如執行了 cancelFunc()），那麼該 context 派生出來的 context 也會被取消。

1.1 Context 型別

// A Context carries a deadline, a cancelation signal, and other values across
// API boundaries.
//
// Context's methods may be called by multiple goroutines simultaneously.
type Context interface {

    Done() <-chan struct{}

    Deadline() (deadline time.Time, ok bool)
    
    Err() error
    
    Value(key interface{}) interface{}
}
 

Done() <-chan struct{}

Done 方法返回一個 channel，阻塞當前執行的程式碼，直到以下條件之一發生時，channel 才會被關閉，進而解除阻塞：

1. WithCancel 建立的 context，cancelFunc 被呼叫。該 context 以及派生子 context 的 Done channel 都會收到取消訊號;
2. WithDeadline 建立的 context，deadline 到期。
3. WithTimeout 建立的 context，timeout 到期

Done 要配合 select 語句使用:

// DoSomething 生產資料併發送給通道 out
// 但如果 DoSomething 返回一個則退出函式，
// 或者 ctx.Done 被關閉時也會退出函式.
func Stream(ctx context.Context, out chan<- Value) error {
    for {
        v, err := DoSomething(ctx)
        if err != nil {
            return err
        }
        select {
        case <-ctx.Done():
            return ctx.Err()
        case out <- v:
        }
    }
}

Deadline() (deadline time.Time, ok bool)

WithDeadline 方法會給 context 設定 deadline，到期自動傳送取消訊號。呼叫 Deadline() 返回 deadline 的值。如果沒設定，ok 返回 false。
該方法可用於確定當前時間是否臨近 deadline。

Err() error

如果 Done 的 channel 被關閉了, Err 函式會返回一個 error，說明錯誤原因:

1. 如果 channel 是因為被取消而關閉，列印 canceled；
2. 如果 channel 是因為 deadline 到時了，列印 deadline exceeded。

重複呼叫，返回相同值。

Value(key interface{}) interface{}

返回由 WithValue 關聯到 context 的值。

1.2 建立根 Context

有兩種方法建立根 Context:

1. context.Background()
2. context.TODO()

根 context 不會被 cancel。這兩個方法只能用在最外層程式碼中，比如 main 函式裡。一般使用 Background() 方法建立根 context。
TODO() 用於當前不確定使用何種 context，留待以後調整。

1.3 派生 Context

一個 Context 被 cancel，那麼它的派生 context 都會收到取消訊號（表現為 context.Done() 返回的 channel 收到值）。
有四種方法派生 context ：

1. func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

2. func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)

3. func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

4. func WithValue(parent Context, key, val interface{}) Context

WithCancel

最常用的派生 context 方法。該方法接受一個父 context。父 context 可以是一個 background context 或其他 context。
返回的 cancelFunc，如果被呼叫，會導致 Done channel 關閉。因此，絕不要把 cancelFunc 傳給其他方法。

WithDeadline

該方法會建立一個帶有 deadline 的 context。當 deadline 到期後，該 context 以及該 context 的可能子 context 會受到 cancel 通知。
另外，如果 deadline 前呼叫 cancelFunc 則會提前傳送取消通知。

WithTimeout

與 WithDeadline 類似。建立一個帶有超時機制的 context。

WithValue

WithValue 方法建立一個攜帶資訊的 context，可以是 user 資訊、認證 token等。該 context 與其派生的子 context 都會攜帶這些資訊。

WithValue 方法的第二個引數是資訊的唯一 key。該 key 型別不應對外暴露，為了避免與其他包可能的 key 型別衝突。所以使用 WithValue 也
應像下面例子的方式間接呼叫 WithValue。

WithValue 方法的第三個引數即是真正要存到 context 中的值。

使用 WithValue 的例子：

package user

import "context"

// User 型別物件會被儲存到 Context 中
type User struct {
    // ...
}

// key 不應該暴露出來。這樣避免與包中其他 key 型別衝突
type key int

// userKey 是 user 的 key，不應暴露; 
// 通過 user.NewContext 和 user.FromContext 間接使用 key
var userKey key

// NewContext 返回攜帶 u 作為 value 的 Context
func NewContext(ctx context.Context, u *User) context.Context {
    return context.WithValue(ctx, userKey, u)
}

// FromContext 返回關聯到 context 的 User型別的 value 的值
func FromContext(ctx context.Context) (*User, bool) {
    u, ok := ctx.Value(userKey).(*User)
    return u, ok
}

2. 例子

改進引子裡的例子。 sleepRandom_1 結束後，會觸發 cancelParent() 被呼叫。所以 sleepRandom_2 中的 ctx.Done() 會被關閉。
sleepRandom_2 執行退出。

package main

import (
    "context"
    "fmt"
    "time"
)

func sleepRandom_1(stopChan chan struct{}) {
    i := 0
    for {
        time.Sleep(1 * time.Second)
        fmt.Printf("This is sleep Random 1: %d\n", i)

        i++
        if i == 5 {
            fmt.Println("cancel sleep random 1")
            stopChan <- struct{}{}
            break
        }
    }
}

func sleepRandom_2(ctx context.Context) {
    i := 0
    for {
        time.Sleep(1 * time.Second)
        fmt.Printf("This is sleep Random 2: %d\n", i)
        i++

        select {
        case <-ctx.Done():
            fmt.Printf("Why? %s\n", ctx.Err())
            fmt.Println("cancel sleep random 2")
            return
        default:
        }
    }
}

func main() {
    
    ctxParent, cancelParent := context.WithCancel(context.Background())
    ctxChild, _ := context.WithCancel(ctxParent)
    
    stopChan := make(chan struct{})

    go sleepRandom_1(stopChan)
    go sleepRandom_2(ctxChild)

    select {
    case <- stopChan:
        fmt.Println("stopChan received")
    }
    cancelParent()
    
    for {
        time.Sleep(1 * time.Second)
        fmt.Println("Continue...")
    }
}

3. 參考文件

Go Concurrency Patterns: Context

Understanding the context package in gol