Go并发控制Context源码

86 阅读2分钟

Context接口

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

Context本身是一个接口,意为上下文,具有不同的实现,是在go1.7以后添加到go标准库中,目的在于但不限于标准化并发控制。

它强大之处在于对协程的关联关闭,协程在被创建后就与它的创建者脱离关系独立存在,有了Context之后,开启协程 之后传入ctx,在某一级ctx关闭后由它开启的协程也会关联关闭,避免造成协程泄露。

当然channel也具有上述能力,时间上Context也是通过channel来控制协程,但由于开发者千变万化的代码风格, 每个人的并发控制逻辑存在差异,不利于维护,且在代码会有冗余的并发控制逻辑,为了跟专注于业务和代码逻辑, Context减少了冗余代码,也规范化了并发控制。

方法解析:

Deadline:返回上下文过期时间

Done: 返回一个只读channel

Err:返回导致Context Done的错误,对于不同的接口实现有不同

Value:返回上下文中的事先存入的值,有点ThreadLocal的意思

获取Context实例的方法

Context接口总共提拱了多种方法来实例化Context对象,常用的方法就是调用WithCancel,WithTimeout和WithDeadLine函数

这些实例化的方法都需要传入父context对象,context.go中提供了两个用于作为根Context的对象, background和todo,通过Context.BackGround和Context.TODO方法获取,本质没有区别都是emptyCtx, 只不过各有其语义。

emptyCtx是一个空实现,无法被取消,无值,无过期时间,专门作根Context

它们的底层实现大差不差,核心实现还是WithCancel函数,所以主要解析这个函数

具体解析:

WithCancel(parent Context, d time.Time) (Context, CancelFunc)

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := withCancel(parent)
    return c, func() { c.cancel(true, Canceled, nil) }
}

func withCancel(parent Context) *cancelCtx {
    if parent == nil {
		panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent)
    propagateCancel(parent, c)
    return c
}

实际返回的是一个cancelCtx类型的实例,

type cancelCtx struct {
	//父Context,这里不是继承,是隐式声名,类型和变量名都是Context
	Context

	mu       sync.Mutex //互斥锁,为下面字段的修改服务
	done     atomic.Value //来加载时v为chan struct{}
	//type atomic.Value struct {
    //  v any
    //}
	children map[canceler]struct{} //存储子Context
	//cancel是一个接口cancelCtx和timerCtx都实现了它
	//type canceler interface {
    //  cancel(removeFromParent bool, err, cause error)
    //  Done() <-chan struct{}
    //}
	err      error //此Context Done时创建的错误
	cause    error //导致Context Done的错误            
}

能够将协程关联起来的主要逻辑在propagateCancel(parent, c)中

func propagateCancel(parent Context, child canceler) {
	//获取父Context的只读channel,只有根Context的Done方法返回nil
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	//如果父Ctx已经Done,直接Done
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	//这里是返回父Ctx的实际实现,返回nil,false才是context.go自己的实现
	//所以跳过为true状况,看下一个else
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err, p.cause)
		} else {
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		//这是为了基准测试准备的,忽略
		goroutines.Add(1)
		//启动一个协程,父Ctx Done,就调用子Ctx的cancel方法,这里ctx的实际实现是cancelCtx
		//子Context主动Done就结束监听
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err(), Cause(parent))
			case <-child.Done():
			}
		}()
	}
}

cancelCtx的cancel实现

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	//必须要有传入错误
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	//引起错误的错误
	if cause == nil {
		cause = err
	}
	//加锁
	c.mu.Lock()
	//如果已取消,解锁并返回
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	//赋值错误
	c.err = err
	c.cause = cause
	//获取cancelCtx的done变量
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		//nil就赋值为全局的closedchan,是一个已经关闭的channel
		c.done.Store(closedchan)
	} else {
		//非nil就关闭它
		close(d)
	}
	//以同样逻辑,逐个取消子Context
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err, cause)
	}
	//置空子Context的map,以便垃圾回收
	c.children = nil
	//解锁
	c.mu.Unlock()

	//将本身从父Context的子Context列表移除,以便垃圾回收
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

总结:

将协程之间关联起来的主要逻辑就是,创建新Context时会创建一个协程,监听父Context的Done和其本身的Done, 父Context Done和本身主动DOne都会结束监听

//启动一个协程,父Ctx Done,就调用子Ctx的cancel方法,这里ctx的实际实现是cancelCtx
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err(), Cause(parent))
			case <-child.Done():
			}
		}()