context基本用法(2)

124 阅读3分钟

cancelCtx的实现

路径:/go1.20/src/context/context.go:396

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
    //done是个interface{}类型,存储结束信号
	done     atomic.Value          // of chan struct{}, created lazily, closed by first cancel call
	children map[canceler]struct{} // set to nil by the first cancel call
	err      error                 // set to non-nil by the first cancel call
	cause    error                 // set to non-nil by the first cancel call
}

典型的装饰器模式,在Context的基础上增加了取消功能。

实现:

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load() 
    //懒汉式加载
	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{})
}

第2行:此处其实应该放的是读锁,但是为了提高性能,采用原子操作

第2,3行和8,9行通过类似double-check的机制,防止多协程进来,重复处理or覆盖值,这种原子操作和锁结合的用法比较罕见。

解析:

  1. cancelCtx.children
func TestContext_WithCancelV2(t *testing.T) {
	ctx := context.Background()
	pctx, cancel := context.WithCancel(ctx)
	// 用完 ctx 再去调用
	//defer cancel()
	cctx, _ := context.WithCancel(pctx)
	go func() {
		time.Sleep(3 * time.Second)
		cancel()
	}()
	// 用 ctx
	<-pctx.Done()
	<-cctx.Done()
	t.Log("儿子收到了结束信号")
	t.Log("hello, cancel: ", ctx.Err())
}

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

23-> context/context.go:295
// newCancelCtx returns an initialized cancelCtx. 原始对象是父亲
func newCancelCtx(parent Context) *cancelCtx {
	return &cancelCtx{Context: parent}
}

24 ->context/context.go:303
// goroutines counts the number of goroutines ever created; for testing.
var goroutines atomic.Int32
// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	done := parent.Done()  // 首先判断 parent 能不能被取消 -> done的注释看到nil的含义
	if done == nil { 
		return // parent is never canceled
	}
//看一下 parent 是不是已经被取消了,已经被取消的情况下直接取消子 context
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

    //找到最近的类型是cancelCtx的祖先,把child加进去 
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
        //p就是一个cancelCtx, err不为空说明有取消了,把取消原因什么的传给儿子
		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 {
        //找不到就监听父亲的取消信号,收到后,就调用子的cancel,来源于cancel or 超时
		goroutines.Add(1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err(), Cause(parent))
			case <-child.Done():
			}
		}()
	}
}

上面的53行 : context/context.go:345  
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    //先判断传入的是不是永远不可取消,是就返回
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
    //判断传入的父亲的类型是否是cancelCtx,因为
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    //如果不是该对象,就返回
	if !ok {
		return nil, false
	}
    //判断传入的和找到的done是不是一样的,防止是自定义的done()方法
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

对应上面的88// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int //标记是cancelCtx的

func (c *cancelCtx) Value(key any) any {
    //88行传入的key始终是相等的,所以会返回自身
	if key == &cancelCtxKey {
		return c
	}
	return value(c.Context, key)
}

//全部取消:核心在于context/context.go:453的cancel -> 用户主动调
// 遍历所有的children
// 关闭done这个channel:谁创建谁关闭
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
    // 由于 cancel context 的 done 是懒加载的,所以有可能存在还没有初始化的情况
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
    // 循环的将所有的子 context 取消掉
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}
    // 将所有的子 context 和当前 context 关系解除
	c.children = nil
	c.mu.Unlock()

    // 如果需要将当前 context 从 parent context 移除,就移除掉
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

在第3行创建父亲,第6行创建儿子,会发现父亲的children会有值

综上:children核心是儿子把自己加入到父亲的children字段,但是因为 Context 里面存在非常多的层级所以父亲不一定是 cancelCtx,因此本质上是找最近属于 cancelCtx 类型的祖先,然后儿子把自己加进去。cancel 就是遍历 children,挨个调用cancel。然后儿子调用孙子的 cancel,子子孙孙无穷匮也。

timerCtx的实现

/go1.20/src/context/context.go525
type timerCtx struct {
	*cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

timerCtx 也是装饰器模式:在已有 cancelCtx 的基础上增加了超时的功能

实现要点:

WithTimeout 和 WithDeadline 本质一样,WithDeadline 里面,在创建 timerCtx 的时候利用 time.AfterFunc 来实现超时

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}

   	// 会先判断 parent context 的过期时间,如果过期时间比当前传入的时间要早的话,就没有必要再设置过期时间了
    // 只需要返回 WithCancel 就可以了,因为在 parent 过期的时候,子 context 也会被取消掉
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}

    // 构造相关结构体
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}

    // 和 WithCancel 中的逻辑相同,构建上下文关系
	propagateCancel(parent, c)

    // 判断传入的时间是不是已经过期,如果已经过期了就 cancel 掉然后再返回
	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()

    // 这里是超时取消的逻辑,启动 timer 时间到了之后就会调用取消方法
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

注意事项:

  1. 一般只用做方法参数,而且是作为第一个参数;
  2. 所有公共方法,除非是 util,helper 之类的方法,否则都加上 context 参数->链路
  3. 不要用作结构体字段,除非你的结构体本身也是表达一个上下文的概念

面试:

  1. context.Context 使用场景:上下文传递和控制
  2. context.Context 原理: 父亲如何控制儿子: 通过儿子主动加入到父亲的 children 里面,父亲只需要遍历就可以
  3. valueCtx 和 timeCtx 的原理