Golang Context 的原理與實戰
本文讓我們一起來學習 golang Context 的使用和標準庫中的Context的實現。
golang context 包 一開始只是 Google 內部使用的一個 Golang 包，在 Golang 1.7的版本中正式被引入標準庫。下面開始學習。
簡單介紹
在學習 context 包之前，先看幾種日常開發中經常會碰到的業務場景：
- 業務需要對訪問的數據庫，RPC ，或API接口，為了防止這些依賴導致我們的服務超時，需要針對性的做超時控制。
- 為了詳細了解服務性能，記錄詳細的調用鏈Log。
上面兩種場景在web中是比較常見的，context 包就是為了方便我們應對此類場景而使用的。
接下來, 我們首先學習 context 包有哪些方法供我們使用；接着舉一些例子，使用 context 包應用在我們上述場景中去解決我們遇到的問題；最後從源碼角度學習 context 內部實現，了解 context 的實現原理。
Context 包
Context 定義
context 包中實現了多種 Context 對象。Context 是一個接口，用來描述一個程序的上下文。接口中提供了四個抽象的方法，定義如下：
type Context interface {
Deadline() (deadline time.Time, ok bool)
Done() <-chan struct{}
Err() error
Value(key interface{}) interface{}
}
- Deadline() 返回的是上下文的截至時間，如果沒有設定，ok 為 false
- Done() 當執行的上下文被取消後，Done返回的chan就會被close。如果這個上下文不會被取消，返回nil
- Err() 有幾種情況:
- 如果Done() 返回 chan 沒有關閉，返回nil
- 如果Done() 返回的chan 關閉了， Err 返回一個非nil的值，解釋為什麼會Done()
- 如果Canceled，返回 “Canceled”
- 如果超過了 Deadline，返回 “DeadlineEsceeded”
- Value(key) 返回上下文中 key 對應的 value 值
Context 構造
為了使用 Context，我們需要了解 Context 是怎麼構造的。
Context 提供了兩個方法做初始化：
func Background() Context{}
func TODO() Context {}
上面方法均會返回空的 Context，但是 Background 一般是所有 Context 的基礎，所有 Context 的源頭都應該是它。TODO 方法一般用於當傳入的方法不確定是哪種類型的 Context 時，為了避免 Context 的參數為nil而初始化的 Context。
其他的 Context 都是基於已經構造好的 Context 來實現的。一個 Context 可以派生多個子 context。基於 Context 派生新Context 的方法如下：
func WithCancel(parent Context) (ctx Context, cancel CancelFunc){}
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {}
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {}
上面三種方法比較類似，均會基於 parent Context 生成一個子 ctx，以及一個 Cancel 方法。如果調用了cancel 方法，ctx 以及基於 ctx 構造的子 context 都會被取消。不同點在於 WithCancel 必需要手動調用 cancel 方法，WithDeadline
可以設置一個時間點，WithTimeout 是設置調用的持續時間，到指定時間後，會調用 cancel 做取消操作。
除了上面的構造方式，還有一類是用來創建傳遞 traceId， token 等重要數據的 Context。
func WithValue(parent Context, key, val interface{}) Context {}
withValue 會構造一個新的context，新的context 會包含一對 Key-Value 數據，可以通過Context.Value(Key) 獲取存在 ctx 中的 Value 值。
通過上面的理解可以直到，Context 是一個樹狀結構，一個 Context 可以派生出多個不一樣的Context。我們大概可以畫一個如下的樹狀圖：
一個background，衍生出一個帶有traceId的valueCtx，然後valueCtx衍生出一個帶有cancelCtx
的context。最終在一些db查詢，http查詢，rpc沙遜等異步調用中體現。如果出現超時，直接把這些異步調用取消，減少消耗的資源，我們也可以在調用時，通過Value 方法拿到traceId，並記錄下對應請求的數據。
當然，除了上面的幾種 Context 外，我們也可以基於上述的 Context 接口實現新的Context.
使用方法
下面我們舉幾個例子，學習上面講到的方法。
超時查詢的例子
在做數據庫查詢時，需要對數據的查詢做超時控制，例如：
ctx = context.WithTimeout(context.Background(), time.Second)
rows, err := pool.QueryContext(ctx, "select * from products where id = ?", 100)
上面的代碼基於 Background 派生出一個帶有超時取消功能的ctx，傳入帶有context查詢的方法中，如果超過1s未返回結果，則取消本次的查詢。使用起來非常方便。為了了解查詢內部是如何做到超時取消的，我們看看DB內部是如何使用傳入的ctx的。
在查詢時，需要先從pool中獲取一個db的鏈接，代碼大概如下：
// src/database/sql/sql.go
// func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) *driverConn, error)
// 阻塞從req中獲取鏈接，如果超時，直接返回
select {
case <-ctx.Done():
// 獲取鏈接超時了，直接返回錯誤
// do something
return nil, ctx.Err()
case ret, ok := <-req:
// 拿到鏈接，校驗並返回
return ret.conn, ret.err
}
req 也是一個chan，是等待鏈接返回的chan，如果Done() 返回的chan 關閉後，則不再關心req的返回了，我們的查詢就超時了。
在做SQL Prepare、SQL Query 等操作時，也會有類似方法：
select {
default:
// 校驗是否已經超時，如果超時直接返回
case <-ctx.Done():
return nil, ctx.Err()
}
// 如果還沒有超時，調用驅動做查詢
return queryer.Query(query, dargs)
上面在做查詢時，首先判斷是否已經超時了，如果超時，則直接返回錯誤，否則才進行查詢。
可以看出，在派生出的帶有超時取消功能的 Context 時，內部方法在做異步操作（比如獲取鏈接，查詢等）時會先查看是否已經
Done了，如果Done，說明請求已超時，直接返回錯誤；否則繼續等待，或者做下一步工作。這裡也可以看出，要做到超時控制，需要不斷判斷 Done() 是否已關閉。
鏈路追蹤的例子
在做鏈路追蹤時，Context 也是非常重要的。（所謂鏈路追蹤，是說可以追蹤某一個請求所依賴的模塊，比如db，redis，rpc下游，接口下游等服務，從這些依賴服務中找到請求中的時間消耗）
下面舉一個鏈路追蹤的例子：
// 建議把key 類型不導出，防止被覆蓋
type traceIdKey struct{}{}
// 定義固定的Key
var TraceIdKey = traceIdKey{}
func ServeHTTP(w http.ResponseWriter, req *http.Request){
// 首先從請求中拿到traceId
// 可以把traceId 放在header里，也可以放在body中
// 還可以自己建立一個 （如果自己是請求源頭的話）
traceId := getTraceIdFromRequest(req)
// Key 存入 ctx 中
ctx := context.WithValue(req.Context(), TraceIdKey, traceId)
// 設置接口1s 超時
ctx = context.WithTimeout(ctx, time.Second)
// query RPC 時可以攜帶 traceId
repResp := RequestRPC(ctx, ...)
// query DB 時可以攜帶 traceId
dbResp := RequestDB(ctx, ...)
// ...
}
func RequestRPC(ctx context.Context, ...) interface{} {
// 獲取traceid，在調用rpc時記錄日誌
traceId, _ := ctx.Value(TraceIdKey)
// request
// do log
return
}
上述代碼中，當拿到請求後，我們通過req 獲取traceId， 並記錄在ctx中，在調用RPC，DB等時，傳入我們構造的ctx，在後續代碼中，我們可以通過ctx拿到我們存入的traceId，使用traceId 記錄請求的日誌，方便後續做問題定位。
當然，一般情況下，context 不會單純的僅僅是用於 traceId 的記錄，或者超時的控制。很有可能二者兼有之。
如何實現
知其然也需知其所以然。想要充分利用好 Context，我們還需要學習 Context 的實現。下面我們一起學習不同的 Context 是如何實現 Context 接口的，
空上下文
Background(), Empty() 均會返回一個空的 Context emptyCtx。emptyCtx 對象在方法 Deadline(), Done(), Err(), Value(interface{}) 中均會返回nil，String() 方法會返回對應的字符串。這個實現比較簡單，我們這裡暫時不討論。
有取消功能的上下文
WithCancel 構造的context 是一個cancelCtx實例，代碼如下。
type cancelCtx struct {
Context
// 互斥鎖，保證context協程安全
mu sync.Mutex
// cancel 的時候，close 這個chan
done chan struct{}
// 派生的context
children map[canceler]struct{}
err error
}
WithCancel 方法首先會基於 parent 構建一個新的 Context，代碼如下：
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
c := newCancelCtx(parent) // 新的上下文
propagateCancel(parent, &c) // 掛到parent 上
return &c, func() { c.cancel(true, Canceled) }
}
其中，propagateCancel 方法會判斷 parent 是否已經取消，如果取消，則直接調用方法取消；如果沒有取消，會在parent的children 追加一個child。這裡就可以看出，context 樹狀結構的實現。 下面是propateCancel 的實現：
// 把child 掛在到parent 下
func propagateCancel(parent Context, child canceler) {
// 如果parent 為空，則直接返回
if parent.Done() == nil {
return // parent is never canceled
}
// 獲取parent類型
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 {
// 啟動goroutine，等待parent/child Done
go func() {
select {
case <-parent.Done():
child.cancel(false, parent.Err())
case <-child.Done():
}
}()
}
}
Done() 實現比較簡單，就是返回一個chan，等待chan 關閉。可以看出 Done 操作是在調用時才會構造 chan done，done 變量是延時初始化的。
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
}
在手動取消 Context 時，會調用 cancelCtx 的 cancel 方法，代碼如下：
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
// 一些判斷,關閉 ctx.done chan
// ...
if c.done == nil {
c.done = closedchan
} else {
close(c.done)
}
// 廣播到所有的child，需要cancel goroutine 了
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()
// 然後從父context 中，刪除當前的context
if removeFromParent {
removeChild(c.Context, c)
}
}
這裡可以看到，當執行cancel時，除了會關閉當前的cancel外，還做了兩件事，① 所有的child 都調用cancel方法，② 由於該上下文已經關閉，需要從父上下文中移除當前的上下文。
定時取消功能的上下文
WithDeadline, WithTimeout 提供了實現定時功能的 Context 方法，返回一個timerCtx結構體。WithDeadline 是給定了執行截至時間，WithTimeout 是倒計時時間，WithTImeout 是基於WithDeadline實現的，因此我們僅看其中的WithDeadline
即可。WithDeadline 內部實現是基於cancelCtx 的。相對於 cancelCtx 增加了一個計時器，並記錄了 Deadline 時間點。下面是timerCtx 結構體：
type timerCtx struct {
cancelCtx
// 計時器
timer *time.Timer
// 截止時間
deadline time.Time
}
WithDeadline 的實現：
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
// 若父上下文結束時間早於child，
// 則child直接掛載在parent上下文下即可
if cur, ok := parent.Deadline(); ok && cur.Before(d) {
return WithCancel(parent)
}
// 創建個timerCtx, 設置deadline
c := &timerCtx{
cancelCtx: newCancelCtx(parent),
deadline: d,
}
// 將context掛在parent 之下
propagateCancel(parent, c)
// 計算倒計時時間
dur := time.Until(d)
if dur <= 0 {
c.cancel(true, DeadlineExceeded) // deadline has already passed
return c, func() { c.cancel(false, Canceled) }
}
c.mu.Lock()
defer c.mu.Unlock()
if c.err == nil {
// 設定一個計時器，到時調用cancel
c.timer = time.AfterFunc(dur, func() {
c.cancel(true, DeadlineExceeded)
})
}
return c, func() { c.cancel(true, Canceled) }
}
構造方法中，將新的context 掛在到parent下，並創建了倒計時器定期觸發cancel。
timerCtx 的cancel 操作，和cancelCtx 的cancel 操作是非常類似的。在cancelCtx 的基礎上，做了關閉定時器的操作
func (c *timerCtx) cancel(removeFromParent bool, err error) {
// 調用cancelCtx 的cancel 方法 關閉chan，並通知子context。
c.cancelCtx.cancel(false, err)
// 從parent 中移除
if removeFromParent {
removeChild(c.cancelCtx.Context, c)
}
c.mu.Lock()
// 關掉定時器
if c.timer != nil {
c.timer.Stop()
c.timer = nil
}
c.mu.Unlock()
}
timeCtx 的 Done 操作直接復用了cancelCtx 的 Done 操作，直接關閉 chan done 成員。
傳遞值的上下文
WithValue 構造的上下文與上面幾種有區別，其構造的context 原型如下：
type valueCtx struct {
// 保留了父節點的context
Context
key, val interface{}
}
每個context 包含了一個Key-Value組合。valueCtx 保留了父節點的Context，但沒有像cancelCtx 一樣保留子節點的Context. 下面是valueCtx的構造方法：
func WithValue(parent Context, key, val interface{}) Context {
if key == nil {
panic("nil key")
}
// key 必須是課比較的，不然無法獲取Value
if !reflect.TypeOf(key).Comparable() {
panic("key is not comparable")
}
return &valueCtx{parent, key, val}
}
直接將Key-Value賦值給struct 即可完成構造。下面是獲取Value 的方法：
func (c *valueCtx) Value(key interface{}) interface{} {
if c.key == key {
return c.val
}
// 從父context 中獲取
return c.Context.Value(key)
}
Value 的獲取是採用鏈式獲取的方法。如果當前 Context 中找不到，則從父Context中獲取。如果我們希望一個context 多放幾條數據時，可以保存一個map 數據到 context 中。這裡不建議多次構造context來存放數據。畢竟取數據的成本也是比較高的。
注意事項
最後，在使用中應該注意如下幾點：
- context.Background 用在請求進來的時候，所有其他context 來源於它。
- 在傳入的conttext 不確定使用的是那種類型的時候，傳入TODO context （不應該傳入一個nil 的context)
- context.Value 不應該傳入可選的參數，應該是每個請求都一定會自帶的一些數據。（比如說traceId，授權token 之類的）。在Value 使用時，建議把Key 定義為全局const 變量，並且key 的類型不可導出，防止數據存在衝突。
- context goroutines 安全。
