Go語言學習之context包的用法詳解

第Go語言學習之context包的用法詳解目錄前言需求一需求二Context接口emptyCtxvalueCtx類型定義WithValuecancelCtx類型定義cancelCtxWithCanceltimerCtx類型定義WithDeadlineWithTimeout總結(jié)

前言

日常Go開發(fā)中，Context包是用的最多的一個了，幾乎所有函數(shù)的第一個參數(shù)都是ctx，那么我們?yōu)槭裁匆獋鬟fContext呢，Context又有哪些用法，底層實現(xiàn)是如何呢？相信你也一定會有探索的欲望，那么就跟著本篇文章，一起來學習吧！

需求一

開發(fā)中肯定會調(diào)用別的函數(shù)，比如A調(diào)用B，在調(diào)用過程中經(jīng)常會設(shè)置超時時間，比如超過2s就不等待B的結(jié)果了，直接返回，那么我們需要怎么做呢？

//

睡眠5s，模擬長時間操作

func

FuncB()

(interface{},

error)

{

time.Sleep(5

*

time.Second)

return

struct{}{},

nil

func

FuncA()

(interface{},

error)

{

var

res

interface{}

var

err

error

ch

:=

make(chan

interface{})

//

調(diào)用FuncB()，將結(jié)果保存至

channel

中

go

func()

{

res,

err

=

FuncB()

ch

-

res

//

設(shè)置一個2s的定時器

timer

:=

time.NewTimer(2

*

time.Second)

//

監(jiān)測是定時器先結(jié)束，還是

FuncB

先返回結(jié)果

select

{

//

超時，返回默認值

case

-timer.C:

return

"default",

err

//

FuncB

先返回結(jié)果，關(guān)閉定時器，返回

FuncB

的結(jié)果

case

r

:=

-ch:

if

!timer.Stop()

{

-timer.C

return

r,

err

func

main()

{

res,

err

:=

FuncA()

fmt.Println(res,

err)

上面我們的實現(xiàn)，可以實現(xiàn)超過等待時間后，A不等待B，但是B并沒有感受到取消信號，如果B是個計算密度型的函數(shù)，我們也希望B感知到取消信號，及時取消計算并返回，減少資源浪費。

另一種情況，如果存在多層調(diào)用，比如A調(diào)用B、C，B調(diào)用D、E，C調(diào)用E、F，在超過A的超時時間后，我們希望取消信號能夠一層層的傳遞下去，后續(xù)所有被調(diào)用到的函數(shù)都能感知到，及時返回。

需求二

在多層調(diào)用的時候，A-B-C-D，有些數(shù)據(jù)需要固定傳輸，比如LogID，通過打印相同的LogID，我們就能夠追溯某一次調(diào)用，方便問題的排查。如果每次都需要傳參的話，未免太麻煩了，我們可以使用Context來保存。通過設(shè)置一個固定的Key，打印日志時從中取出value作為LogID。

const

LogKey

=

"LogKey"

//

模擬一個日志打印，每次從

Context

中取出

LogKey

對應(yīng)的

Value

作為LogID

type

Logger

struct{}

func

(logger

*Logger)

info(ctx

context.Context,

msg

string)

{

logId,

ok

:=

ctx.Value(LogKey).(string)

if

!ok

{

logId

=

uuid.New().String()

fmt.Println(logId

+

"

"

+

msg)

var

logger

Logger

//

日志打印

并

調(diào)用

FuncB

func

FuncA(ctx

context.Context)

{

(ctx,

"FuncA")

FuncB(ctx)

func

FuncB(ctx

context.Context)

{

(ctx,

"FuncB")

//

獲取初始化的，帶有

LogID

的

Context，一般在程序入口做

func

getLogCtx(ctx

context.Context)

context.Context

{

logId,

ok

:=

ctx.Value(LogKey).(string)

if

ok

{

return

ctx

logId

=

uuid.NewString()

return

context.WithValue(ctx,

LogKey,

logId)

func

main()

{

ctx

=

getLogCtx(context.Background())

FuncA(ctx)

這利用到了本篇文章講到的valueCtx，繼續(xù)往下看，一起來學習valueCtx是怎么實現(xiàn)的吧！

Context接口

type

Context

interface

{

Deadline()

(deadline

time.Time,

ok

bool)

Done()

-chan

struct{}

Err()

error

Value(key

interface{})

interface{}

Context接口比較簡單，定義了四個方法：

Deadline()方法返回兩個值，deadline表示Context將會在什么時間點取消，ok表示是否設(shè)置了deadline。當ok=false時，表示沒有設(shè)置deadline，那么此時deadline將會是個零值。多次調(diào)用這個方法返回同樣的結(jié)果。Done()返回一個只讀的channel，類型為chanstruct{}，如果當前的Context不支持取消，Done返回nil。我們知道，如果一個channel中沒有數(shù)據(jù)，讀取數(shù)據(jù)會阻塞；而如果channel被關(guān)閉，則可以讀取到數(shù)據(jù)，因此可以監(jiān)聽Done返回的channel，來獲取Context取消的信號。Err()返回Done返回的channel被關(guān)閉的原因。當channel未被關(guān)閉時，Err()返回nil；channel被關(guān)閉時則返回相應(yīng)的值，比如Canceled、DeadlineExceeded。Err()返回一個非nil值之后，后面再次調(diào)用會返回相同的值。Value()返回Context保存的鍵值對中，key對應(yīng)的value，如果key不存在則返回nil。

Done()是一個比較常用的方法，下面是一個比較經(jīng)典的流式處理任務(wù)的示例：監(jiān)聽ctx.Done()是否被關(guān)閉來判斷任務(wù)是否需要取消，需要取消則返回相應(yīng)的原因；沒有取消則將計算的結(jié)果寫入到outchannel中。

func

Stream(ctx

context.Context,

out

chan-

Value)

error

{

for

{

//

處理數(shù)據(jù)

v,

err

:=

DoSomething(ctx)

if

err

!=

nil

{

return

err

//

ctx.Done()

讀取到數(shù)據(jù)，說明獲取到了任務(wù)取消的信號

select

{

case

-ctx.Done():

return

ctx.Err()

//

否則將結(jié)果輸出，繼續(xù)計算

case

out

-

v:

Value()也是一個比較常用的方法，用于在上下文中傳遞一些數(shù)據(jù)。使用context.WithValue()方法存入key和value，通過Value()方法則可以根據(jù)key拿到value。

func

main()

{

ctx

:=

context.Background()

c

:=

context.WithValue(ctx,

"key",

"value")

v,

ok

:=

c.Value("key").(string)

fmt.Println(v,

ok)

emptyCtx

Context接口并不需要我們自己去手動實現(xiàn)，一般我們都是直接使用context包中提供的Background()方法和TODO()方法，來獲取最基礎(chǔ)的Context。

var

(

background

=

new(emptyCtx)

todo

=

new(emptyCtx)

func

Background()

Context

{

return

background

func

TODO()

Context

{

return

todo

Background()方法一般用在main函數(shù)，或者程序的初始化方法中；在我們不知道使用哪個Context，或者上文沒有傳遞Context時，可以使用TODO()。

Background()和TODO()都是基于emptyCtx生成的，從名字可以看出來，emptyCtx是一個空的Context，沒有deadline、不能被取消、沒有鍵值對。

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外，context包中提供了功能更加豐富的Context，包括valueCtx、cancelCtx、timerCtx，下面我們就挨個來看下。

valueCtx

使用示例

我們一般使用context.WithValue()方法向Context存入鍵值對，然后通過Value()方法根據(jù)key得到value，此種功能的實現(xiàn)就依賴valueCtx。

func

main()

{

ctx

:=

context.Background()

c

:=

context.WithValue(ctx,

"myKey",

"myValue")

v1

:=

c.Value("myKey")

fmt.Println(v1.(string))

v2

:=

c.Value("hello")

fmt.Println(v2)

//

nil

類型定義

valueCtx結(jié)構(gòu)體中嵌套了Context，使用key、value來保存鍵值對：

type

valueCtx

struct

{

Context

key,

val

interface{}

WithValue

context包對外暴露了WithValue方法，基于一個parentcontext來創(chuàng)建一個valueCtx。從下面的源碼中可以看出，key必須是可比較的！

func

WithValue(parent

Context,

key,

val

interface{})

Context

{

if

parent

==

nil

{

panic("cannot

create

context

from

nil

parent")

if

key

==

nil

{

panic("nil

key")

if

!reflectlite.TypeOf(key).Comparable()

{

panic("key

is

not

comparable")

return

valueCtx{parent,

key,

val}

*valueCtx實現(xiàn)了Value()，可以根據(jù)key得到value。這是一個向上遞歸尋找的過程，如果key不在當前valueCtx中，會繼續(xù)向上找parentContext，直到找到最頂層的Context，一般最頂層的是emptyCtx，而emtpyCtx.Value()返回nil。

func

(c

*valueCtx)

Value(key

interface{})

interface{}

{

if

c.key

==

key

{

return

c.val

return

c.Context.Value(key)

cancelCtx

cancelCtx是一個用于取消任務(wù)的Context，任務(wù)通過監(jiān)聽Context是否被取消，來決定是否繼續(xù)處理任務(wù)還是直接返回。

如下示例中，我們在main函數(shù)定義了一個cancelCtx，并在2s后調(diào)用cancel()取消Context，即我們希望doSomething()在2s內(nèi)完成任務(wù)，否則就可以直接返回，不需要再繼續(xù)計算浪費資源了。

doSomething()方法內(nèi)部，我們使用select監(jiān)聽任務(wù)是否完成，以及Context是否已經(jīng)取消，哪個先到就執(zhí)行哪個分支。方法模擬了一個5s的任務(wù)，main函數(shù)等待時間是2s，因此沒有完成任務(wù)；如果main函數(shù)等待時間改為10s，則任務(wù)完成并會返回結(jié)果。

這只是一層調(diào)用，真實情況下可能會有多級調(diào)用，比如doSomething可能又會調(diào)用其他任務(wù)，一旦parentContext取消，后續(xù)的所有任務(wù)都應(yīng)該取消。

func

doSomething(ctx

context.Context)

(interface{},

error)

{

res

:=

make(chan

interface{})

go

func()

{

fmt.Println("do

something")

time.Sleep(time.Second

*

5)

res

-

"done"

select

{

case

-ctx.Done():

return

nil,

ctx.Err()

case

value

:=

-res:

return

value,

nil

func

main()

{

ctx,

cancel

:=

context.WithCancel(context.Background())

go

func()

{

time.Sleep(time.Second

*

2)

cancel()

res,

err

:=

doSomething(ctx)

fmt.Println(res,

err)

//

nil

,

context

canceled

接下來就讓我們來研究下，cancelCtx是如何實現(xiàn)取消的吧

canceler接口包含cancel()和Done()方法，*cancelCtx和*timerCtx均實現(xiàn)了這個接口。closedchan是一個被關(guān)閉的channel，可以用于后面Done()返回canceled是一個err，用于Context被取消的原因

type

canceler

interface

{

cancel(removeFromParent

bool,

err

error)

Done()

-chan

struct{}

//

closedchan

is

a

reusable

closed

channel.

var

closedchan

=

make(chan

struct{})

func

init()

{

close(closedchan)

var

Canceled

=

errors.New("context

canceled")

CancelFunc是一個函數(shù)類型定義，是一個取消函數(shù)，有如下規(guī)范：

CancelFunc告訴一個任務(wù)停止工作CancelFunc不會等待任務(wù)結(jié)束CancelFunc支持并發(fā)調(diào)用第一次調(diào)用后，后續(xù)的調(diào)用不會產(chǎn)生任何效果

type

CancelFunc

func()

cancelCtxKey是一個固定的key，用來返回cancelCtx自身

var

cancelCtxKey

int

cancelCtx

cancelCtx是可以被取消的，它嵌套了Context接口，實現(xiàn)了canceler接口。cancelCtx使用children字段保存同樣實現(xiàn)canceler接口的子節(jié)點，當cancelCtx被取消時，所有的子節(jié)點也會取消。

type

cancelCtx

struct

{

Context

mu

sync.Mutex

//

保護如下字段，保證線程安全

done

atomic.Value

//

保存

channel，懶加載，調(diào)用

cancel

方法時會關(guān)閉這個

channel

children

map[canceler]struct{}

//

保存子節(jié)點，第一次調(diào)用

cancel

方法時會置為

nil

err

error

//

保存為什么被取消，默認為nil，第一次調(diào)用

cancel

會賦值

*cancelCtx的Value()方法和*valueCtx的Value()方法類似，只不過加了個固定的key:cancelCtxKey。當key為cancelCtxKey時返回自身

func

(c

*cancelCtx)

Value(key

interface{})

interface{}

{

if

key

==

cancelCtxKey

{

return

c

return

c.Context.Value(key)

*cancelCtx的done字段是懶加載的，只有在調(diào)用Done()方法或者cancel()時才會賦值。

func

(c

*cancelCtx)

Done()

-chan

struct{}

{

d

:=

c.done.Load()

//

如果已經(jīng)有值了，直接返回

if

d

!=

nil

{

return

d.(chan

struct{})

//

沒有值，加鎖賦值

c.mu.Lock()

defer

c.mu.Unlock()

d

=

c.done.Load()

if

d

==

nil

{

d

=

make(chan

struct{})

c.done.Store(d)

return

d.(chan

struct{})

Err方法返回cancelCtx的err字段

func

(c

*cancelCtx)

Err()

error

{

c.mu.Lock()

err

:=

c.err

c.mu.Unlock()

return

err

WithCancel

那么我們?nèi)绾涡陆ㄒ粋€cancelCtx呢？context包提供了WithCancel()方法，讓我們基于一個Context來創(chuàng)建一個cancelCtx。WithCancel()方法返回兩個字段，一個是基于傳入的Context生成的cancelCtx，另一個是CancelFunc。

func

WithCancel(parent

Context)

(ctx

Context,

cancel

CancelFunc)

{

if

parent

==

nil

{

panic("cannot

create

context

from

nil

parent")

c

:=

newCancelCtx(parent)

propagateCancel(parent,

c)

return

c,

func()

{

c.cancel(true,

Canceled)

}

WithCancel調(diào)用了兩個外部方法：newCancelCtx、propagateCancel。newCancelCtx比較簡單，根據(jù)傳入的context，返回了一個cancelCtx結(jié)構(gòu)體。

func

newCancelCtx(parent

Context)

cancelCtx

{

return

cancelCtx{Context:

parent}

propagateCancel從名字可以看出，就是將cancel傳播。如果父Context支持取消，那么我們需要建立一個通知機制，這樣父節(jié)點取消的時候，通知子節(jié)點也取消，層層傳播。

在propagateCancel中，如果父Context是cancelCtx類型且未取消，會將子Context掛在它下面，形成一個樹結(jié)構(gòu)；其余情況都不會掛載。

func

propagateCancel(parent

Context,

child

canceler)

{

//

如果

parent

不支持取消，那么就不支持取消傳播，直接返回

done

:=

parent.Done()

if

done

==

nil

{

return

//

到這里說明

done

不為

nil，parent

支持取消

select

{

case

-done:

//

如果

parent

此時已經(jīng)取消了，那么直接告訴子節(jié)點也取消

child.cancel(false,

parent.Err())

return

default:

//

到這里說明此時

parent

還未取消

//

如果

parent

是未取消的

cancelCtx

if

p,

ok

:=

parentCancelCtx(parent);

ok

{

//

加鎖，防止并發(fā)更新

p.mu.Lock()

//

再次判斷，因為有可能上一個獲得鎖的進行了取消操作。

//

如果

parent

已經(jīng)取消了，那么子節(jié)點也直接取消

if

p.err

!=

nil

{

child.cancel(false,

p.err)

}

else

{

//

把子Context

掛到父節(jié)點

parent

cancelCtx

的

children字段下

//

之后

parent

cancelCtx

取消時，能通知到所有的

子Context

if

p.children

==

nil

{

p.children

=

make(map[canceler]struct{})

p.children[child]

=

struct{}{}

p.mu.Unlock()

}

else

{

//

parent

不是

cancelCtx

類型，可能是用戶自己實現(xiàn)的Context

atomic.AddInt32(goroutines,

+1)

//

啟動一個協(xié)程監(jiān)聽，如果

parent

取消了，子

Context

也取消

go

func()

{

select

{

case

-parent.Done():

child.cancel(false,

parent.Err())

case

-child.Done():

}()

cancel方法就是來取消cancelCtx，主要的工作是：關(guān)閉c.done中的channel，給err賦值，然后級聯(lián)取消所有子Context。如果removeFromParent為true，會從父節(jié)點中刪除以該節(jié)點為樹頂?shù)臉洹?/p>

cancel()方法只負責自己管轄的范圍，即自己以及自己的子節(jié)點，然后根據(jù)配置判斷是否需要從父節(jié)點中移除自己為頂點的樹。如果子節(jié)點還有子節(jié)點，那么由子節(jié)點負責處理，不用自己負責了。

propagateCancel()中有三處調(diào)用了cancel()方法，傳入的removeFromParent都為false，是因為當時根本沒有掛載，不需要移除。而WithCancel返回的CancelFunc，傳入的removeFromParent為true，是因為調(diào)用propagateCancel有可能產(chǎn)生掛載，當產(chǎn)生掛載時，調(diào)用cancel()就需要移除了。

func

(c

*cancelCtx)

cancel(removeFromParent

bool,

err

error)

{

//

err

是指取消的原因，必傳，cancelCtx

中是

errors.New("context

canceled")

if

err

==

nil

{

panic("context:

internal

error:

missing

cancel

error")

//

涉及到保護字段值的修改，都需要加鎖

c.mu.Lock()

//

如果該Context已經(jīng)取消過了，直接返回。多次調(diào)用cancel，不會產(chǎn)生額外效果

if

c.err

!=

nil

{

c.mu.Unlock()

return

//

給

err

賦值，這里

err

一定不為

nil

c.err

=

err

//

close

channel

d,

_

:=

c.done.Load().(chan

struct{})

//

因為c.done

是懶加載，有可能存在

nil

的情況

//

如果

c.done

中沒有值，直接賦值

closedchan；否則直接

close

if

d

==

nil

{

c.done.Store(closedchan)

}

else

{

close(d)

//

遍歷當前

cancelCtx

所有的子Context，讓子節(jié)點也

cancel

//

因為當前的Context

會主動把子Context移除，子Context

不用主動從parent中脫離

//

因此

child.cancel

傳入的

removeFromParent

為false

for

child

:=

range

c.children

{

child.cancel(false,

err)

//

將

children

置空，相當于移除自己的所有子Context

c.children

=

nil

c.mu.Unlock()

//

如果當前

cancelCtx

需要從上層的

cancelCtx移除，調(diào)用removeChild方法

//

c.Context

就是自己的父Context

if

removeFromParent

{

removeChild(c.Context,

c)

從propagateCancel方法中可以看到，只有parent屬于cancelCtx類型，才會將自己掛載。因此removeChild會再次判斷parent是否為cancelCtx，和之前的邏輯保持一致。找到的話，再將自己移除，需要注意的是，移除會把自己及其自己下面的所有子節(jié)點都移除。

如果上一步propagateCancel方法將自己掛載到了A上，但是在調(diào)用cancel()時，A已經(jīng)取消過了，此時parentCancelCtx()會返回false。不過這沒有關(guān)系，A取消時已經(jīng)將掛載的子節(jié)點移除了，當前的子節(jié)點不用將自己從A中移除了。

func

removeChild(parent

Context,

child

canceler)

{

//

parent

是否為未取消的

cancelCtx

p,

ok

:=

parentCancelCtx(parent)

if

!ok

{

return

//

獲取

parent

cancelCtx

的鎖，修改保護字段

children

p.mu.Lock()

//

將自己從

parent

cancelCtx

的

children

中刪除

if

p.children

!=

nil

{

delete(p.children,

child)

p.mu.Unlock()

parentCancelCtx判斷parent是否為未取消的*cancelCtx。取消與否容易判斷，難判斷的是parent是否為*cancelCtx，因為有可能其他結(jié)構(gòu)體內(nèi)嵌了cancelCtx，比如timerCtx，會通過比對channel來確定。

func

parentCancelCtx(parent

Context)

(*cancelCtx,

bool)

{

//

如果

parent

context

的

done

為

nil，

說明不支持

cancel，那么就不可能是

cancelCtx

//

如果

parent

context

的

done

為

closedchan，

說明

parent

context

已經(jīng)

cancel

了

done

:=

parent.Done()

if

done

==

closedchan

||

done

==

nil

{

return

nil,

false

//

到這里說明支持取消，且沒有被取消

//

如果

parent

context

屬于原生的

*cancelCtx

或衍生類型，需要繼續(xù)進行后續(xù)判斷

//

如果

parent

context

無法轉(zhuǎn)換到

*cancelCtx，則認為非

cancelCtx，返回

nil,fasle

p,

ok

:=

parent.Value(cancelCtxKey).(*cancelCtx)

if

!ok

{

return

nil,

false

//

經(jīng)過上面的判斷后，說明

parent

context

可以被轉(zhuǎn)換為

*cancelCtx，這時存在多種情況:

//

-

parent

context

就是

*cancelCtx

//

-

parent

context

是標準庫中的

timerCtx

//

-

parent

context

是個自己自定義包裝的

cancelCtx

//

針對這

3

種情況需要進行判斷，判斷方法就是:

//

判斷

parent

context

通過

Done()

方法獲取的

done

channel

與

Value

查找到的

context

的

done

channel

是否一致

//

一致情況說明

parent

context

為

cancelCtx

或

timerCtx

或

自定義的

cancelCtx

且未重寫

Done()，

//

這種情況下可以認為拿到了底層的

*cancelCtx

//

不一致情況說明

parent

context

是一個自定義的

cancelCtx

且重寫了

Done()

方法，并且并未返回標準

*cancelCtx

的

//

的

done

channel，這種情況需要單獨處理，故返回

nil,

false

pdone,

_

:=

p.done.Load().(chan

struct{})

if

pdone

!=

done

{

return

nil,

false

return

p,

true

timerCtx

簡介

timerCtx嵌入了cancelCtx，并新增了一個timer和deadline字段。timerCtx的取消能力是復用cancelCtx的，只是在這個基礎(chǔ)上增加了定時取消而已。

在我們的使用過程中，有可能還沒到deadline，任務(wù)就提前完成了，此時需要手動調(diào)用CancelFunc。

func

slowOperationWithTimeout(ctx

context.Context)

(Result,

error)

{

ctx,

cancel

:=

context.WithTimeout(ctx,

100*time.Millisecond)

defer

cancel()

//

如果未到截止時間，slowOperation就完成了，盡早調(diào)用

cancel()

釋放資源

return

slowOperation(ctx)

type

timerCtx

struct

{

cancelCtx

//

內(nèi)嵌

cancelCtx

timer

*time.Timer

//

受

cancelCtx.mu

互斥鎖的保護

deadline

time.Time

//

截止時間

Deadline()返回deadline字段的值

func

(c

*timerCtx)

Deadline()

(deadline

time.Time,

ok

bool)

{

return

c.deadline,

true

WithDeadline

WithDeadline基于parentContext和時間點d，返回了一個定時取消的Context，以及一個CancelFunc。返回的Context有三種情況被取消：1.到達了指定時間，就會主動取消；2.手動調(diào)用了CancelFunc；3.父Context取消，導致該Context被取消。這三種情況哪種先到，就會首次觸發(fā)取消操作，后續(xù)的再次取消不會產(chǎn)生任何效果。

如果傳入parentContext的deadline比指定的時間d還要早，此時d就沒用處了，直接依賴parent取消傳播就可以了。

func

WithDeadline(parent

Context,

d

time.Time)

(Context,

CancelFunc)

{

//

傳入的

parent

不能為

nil

if

parent

==

nil

{

panic("cannot

create

context

from

nil

parent")

//

parent

也有

deadline，并且比

d

還要早，直接依賴

parent

的取消傳播即可

if

cur,

ok

:=

parent.Deadline();

ok

cur.Before(d)

{

//

The

current

deadline

is

already

sooner

than

the

new

one.

return

WithCancel(parent)

//

定義

timerCtx

接口

c

:=

timerCtx{

cancelCtx:

newCancelCtx(parent),

deadline:

d,

//

設(shè)置傳播，如果parent

屬于

cancelC