[Advanced Go - Concurrent Programming] Context

Context is a commonly used concurrency control technology for Go application development. The biggest difference between it and WaitGroup is that Context has stronger control over derived goroutines, and it can control multi-level goroutines.

Although there are many controversies, it is very convenient to use Context in many scenarios, so now it has spread in the Go ecosystem, including many web application frameworks, which have switched to the Context of the standard library. Context is used in the database/sql, os/exec, net, net/http and other packages in the standard library. Moreover, if you encounter some of the following scenarios, you can also consider using Context:

• Contextual information passing, such as processing http requests, passing information on request processing links;
• Control the operation of child goroutines;
• Method calls controlled by timeout;
• Cancellable method calls.

Implementation Principle

Interface Definition

The package context defines the Context interface. The specific implementation of Context includes four methods, namely Deadline, Done, Err and Value, as follows:

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

The Deadline method returns the deadline when this Context was cancelled. If no expiration date is set, the value of ok is false. Each subsequent call to the object's Deadline method will return the same result as the first call.

The Done method returns a Channel object, which is basically used in select statements. When the Context is canceled, the Channel will be closed, if not canceled, it may return nil. When Done is closed, the error information can be obtained through ctx.Err.

About the Err method, you must remember the knowledge points: If Done is not closed, the Err method returns nil; if Done is closed, the Err method returns the reason why Done was closed.

Value Returns the value associated with the specified key in this ctx.

Two commonly used methods for generating top-level Context are implemented in Context:

• context.Background(): Returns a non-nil, empty Context with no value, no cancellation, no timeout, and no deadline. Generally used in the main function, initialization, testing and creating the root Context.
• context.TODO(): Returns a non-nil, empty Context with no value, no cancellation, no timeout, and no deadline. You can use this method when you don't know whether to use Context, or when you don't know what context information to pass.

In fact, the two underlying implementations are exactly the same. In most cases, context.Background can be used directly.

When using Context, there are some conventional rules:

1. When a general function uses Context, it will put this parameter in the position of the first parameter.
2. Never use nil as a parameter value of type Context, you can use context.Background() to create an empty context object, and don't use nil.
3. Context is only used for temporary context transmission between functions, it cannot persist Context or store Context for a long time. It is a wrong usage to persist Context to database, local file or global variable, cache.
4. The type of key is not recommended as string type or other built-in types, otherwise it is easy to cause conflicts when using Context between packages. When using WithValue, the type of key should be the type defined by yourself.
5. Often the type of key is defined using struct{} as the underlying type. For static types of exported keys, this is usually an interface or a pointer. This minimizes memory allocations.

The structs in the context package that implement the Context interface include cancelCtx, timerCtx and valueCtx in addition to emptyCtx for context.Background().

cancelCtx

type cancelCtx struct {
    Context

    mu       sync.Mutex            // 互斥锁
    done     atomic.Value          // 调用 cancel 时会关闭的 channel
    children map[canceler]struct{} // 记录了由此 Context 派生的所有 child，此 Context 被 cancel 时会同时 cancel 所有 child
    err      error                 // 错误信息
}

WithCancel

cancelCtx is generated by the WithCancel method. We often create this type of Context when some long-term tasks need to be actively canceled, and then pass this Context to the goroutine that executes the task for a long time. When the task needs to be aborted, we can cancel the Context, so that the goroutine that executes the task for a long time can check the Context and know that the Context has been canceled.

The second value in the return value of WithCancel is a cancel function. Remember, cancel is not called only if you want to give up halfway. As long as your task is completed normally, you need to call cancel, so that the Context can release its resources (notify its children to handle cancel, from its parent remove yourself, or even free the associated goroutine).

Take a look at the core source code:

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

func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

func propagateCancel(parent Context, child canceler) {
    done := parent.Done()

    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent 已经取消了，直接取消子 Context
            child.cancel(false, p.err)
        } else {
            // 将 child 添加到 parent 的 children 切片
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        atomic.AddInt32(&goroutines, +1)
        // 没有 parent 可以“挂载”，启动一个 goroutine 监听 parent 的 cancel，同时 cancel 自身
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

The propagateCancel method called in the code will search up the parent path until it finds a cancelCtx, or is nil. If it is not empty, add yourself to the child of this cancelCtx, so that you can notify yourself when this cancelCtx is canceled. If it is empty, a new goroutine will be started, which will monitor whether the parent's Done has been closed.

When the cancel function of the cancelCtx is called, or the Done of the parent is closed, the Done of the cancelCtx will be closed.

Cancel is passed down. If a Context generated by WithCancel is canceled, if its child Context (which may also be a grandchild, or lower, depending on the type of the child) is also of type cancelCtx, it will be canceled.

cancel

The function of the cancel method is to close the done channel of itself and its descendants, so as to achieve the purpose of notification cancellation. The second return value cancel of the WithCancel method is this function. Let's take a look at the main code implementation:

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()
    // 设置 cancel 的原因
    c.err = err 
    // 关闭自身的 done 通道
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan)
    } else {
        close(d)
    }
    // 遍历所有 children，逐个调用 cancel 方法
    for child := range c.children {
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    // 正常情况下，需要将自身从 parent 的 children 切片中删除
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

timerCtx

type timerCtx struct {
    cancelCtx
    timer *time.Timer 

    deadline time.Time
}

timerCtx adds deadline to cancelCtx to mark the final time of automatic cancel, and timer is a timer that triggers automatic cancel. timerCtx can be generated by WithDeadline and WithTimeout. WithTimeout actually calls WithDeadline. The implementation principle of the two is the same, but the usage context is different: WithDeadline specifies the deadline, and WithTimeout specifies the maximum survival time.

WithDeadline

Take a look at the implementation of the WithDeadline method:

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    // 如果parent的截止时间更早，直接返回一个cancelCtx即可
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c) // 同cancelCtx的处理逻辑
    dur := time.Until(d)
    if dur <= 0 { //当前时间已经超过了截止时间，直接cancel
        c.cancel(true, DeadlineExceeded)
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 设置一个定时器，到截止时间后取消
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

WithDeadline returns a copy of parent with a deadline no later than parameter d, of type timerCtx (or cancelCtx).

If its deadline is later than the parent's deadline, then the parent's deadline shall prevail, and a Context of type cancelCtx is returned, because the parent's deadline is up, the cancelCtx will be cancelled. If the current time has exceeded the deadline, it will directly return a canceled timerCtx. Otherwise, a timer will be started, and the timerCtx will be cancelled when the deadline is reached.

To sum up, the Done of timerCtx is Closed, which is mainly triggered by one of the following events:

• deadline is up;
• cancel function is called;
• Done of parent is closed.

Like cancelCtx, the cancel returned by WithDeadline (WithTimeout) must be called, and it must be called as early as possible, so that resources can be released as soon as possible, and do not simply rely on the deadline for passive cancellation.

valueCtx

type valueCtx struct {
    Context
    key, val interface{}
}

valueCtx just adds a key-value pair on the basis of Context, which is used to transfer some data between coroutines at all levels.

WithValue generates a new valueCtx based on the parent Context, saving a key-value pair. valueCtx overrides the Value method, and first checks the key from its own storage, and if it does not exist, it will continue to check from the parent.