golang context 超时控制

常见方法

1. context.WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) 指定时长超时结束
2. context.WithCancel(parent Context) (ctx Context, cancel CancelFunc) 手动结束
3. context.WithDeadline(parent Context, d time.Time) (Context, CancelFunc) 指定时间结束

Demo

父ctx超时，关闭所有子ctx

package main

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

func main() {
    ctx, cancel := context.WithTimeout(context.TODO(), time.Second*3)
    defer cancel() 
    go task(ctx)
    time.Sleep(time.Second * 10)
}

func task(ctx context.Context) {
    ch := make(chan struct{}, 0)
    go func() {
        // 模拟4秒耗时任务
        time.Sleep(time.Second * 4)
        ch <- struct{}{}
    }()
    select {
    case <-ch:
        fmt.Println("done")
    case <-ctx.Done():
        fmt.Println("timeout")
    }
}

源码分析

context 是定义了接口

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

常见超时控制比如withTimeout等，其实都是通过设定deadline

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

timeCtx 定义

type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        // Remove this timerCtx from its parent cancelCtx's children.
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

withDeadline 实现

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    // .....
    c := &timerCtx{ // 新建子ctx
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c) // 父ctx和子ctx间产生绑定
    // ...
    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) }
}

父子ctx间产生绑定实现：

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
    // ...
    if p, ok := parentCancelCtx(parent); ok { // 强转成功
        p.mu.Lock()
        if p.err != nil { // 父ctx出现error，子ctx也结束
            // parent has already been canceled
            child.cancel(false, p.err)
        } else {
            if p.children == nil { // 形成一个map，绑定成功
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else { // 强转失败
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

超时后回调的处理：

// 全部ctx链式cancel
// 接timeCtx的cancel调用
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()
    c.err = err
    if c.done == nil {
        c.done = closedchan
    } else {
        close(c.done)
    }
    for child := range c.children { // 这个ctx中的所有子ctx都调用cancel
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        removeChild(c.Context, c)
    }
}