首页 > 编程语言 >Go语言Context包源码学习

Go语言Context包源码学习

时间:2024-10-22 16:59:55浏览次数:7  
标签:Context parent cancelCtx ctx cancel 源码 context Go

0前言

context包作为使用go进行server端开发的重要工具,其源码只有791行,不包含注释的话预计在500行左右,非常值得我们去深入探讨学习,于是在本篇笔记中我们一起来观察源码的实现,知其然更要知其所以然。(当前使用go版本为1.22.2)

1核心数据结构

整体的接口实现和结构体embed图

1.1Context接口

context接口定义了四个方法:

  • Deadline方法返回context是否为timerctx,以及它的结束时间
  • Done方法返回该ctx的done channel
  • Err方法返回该ctx被取消的原因
  • Value方法返回key对应的value

2emptyCtx

先来观察源码

type emptyCtx struct{}

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 any) any {
	return nil
}

emptyctx实现了context接口中的所有方法,对于每个方法返回的都是空值,它没有值、不能被取消以及没有截止时间,它只作为一个空context的载体,相当于所有ctx的祖先。

如何创建一个context?

context包提供了Background()方法和TODO()方法,都用于创建一个空的context。

func Background() Context {
	return backgroundCtx{}
}
func TODO() Context {
	return todoCtx{}
}
type backgroundCtx struct{ emptyCtx }

func (backgroundCtx) String() string {
	return "context.Background"
}

type todoCtx struct{ emptyCtx }

func (todoCtx) String() string {
	return "context.TODO"
}

它们返回的都为空的context,虽然方法名不同,但是效果是一样的,那么什么时候使用background,什么时候使用TODO呢,这是官方给出的注释

“TODO 会返回一个非空的、为空的 [Context]。 代码应该在不清楚应该使用哪个 [Context] 或者 [Context] 尚未可用(因为周围的函数尚未被扩展以接受 [Context] 参数)时使用 context.TODO。”

“background返回一个非空的、空的上下文对象。它不会被取消,没有值,也没有截止时间。它通常被主函数、初始化和测试使用,作为进入请求的顶级上下文。”

3cancelctx

先来观察cancelctx结构体的实现

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	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
}

  • cancelCtx 内嵌了 Context 作为其父 context,根据go语言的特性,cancelCtx结构体就隐式实现了context接口中的所有方法,可以被当作一个接口来被调用,但是其方法还需要被具体的实现赋值才能进行调用。并且可以得知,cancelCtx的父类一定也是一个Context。

  • mu是cancelCtx的内置锁,用来协调并发场景下的资源获取

  • done的实际类型为chan struct{},通过atomic包来实现并发安全,可以用于反应该ctx的生命周期情况,done是懒汉式创建的,只有第一次调用Done()方法时才会被创建,在下文的Done方法中会提到

  • children用于关联和子ctx的关系,当取消该ctx时,可以接连通知子ctx进行关闭,及时释放资源。

  • err用于返回ctx关闭的原因,调用的是context包定义的内置error

    var Canceled = errors.New("context canceled")
    
    var DeadlineExceeded error = deadlineExceededError{}
    
    type deadlineExceededError struct{}
    
    func (deadlineExceededError) Error() string   { return "context deadline exceeded" }
    func (deadlineExceededError) Timeout() bool   { return true }
    func (deadlineExceededError) Temporary() bool { return true }
    
  • cause返回的是该ctx失效更底层的原因,例如导致DeadlineExceeded err的具体原因是“database connection timeout”

    // Example use:
    //
    //	ctx, cancel := context.WithCancelCause(parent)
    //	cancel(myError)
    //	ctx.Err() // returns context.Canceled
    //	context.Cause(ctx) // returns myError
    

3.1Done方法的实现

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load() //获取cancelCtx的done通道
	if d != nil { //假若已经创建过了done通道,则直接返回
		return d.(chan struct{}) 
	}
	c.mu.Lock() //上锁
	defer c.mu.Unlock()
	d = c.done.Load() //Double check是否创建done通道,因为在上锁前,可能其他goroutine调用了该ctx的Done方法。
	if d == nil { //如果仍然未创建
		d = make(chan struct{}) //创建该done channel
		c.done.Store(d) //存储
	}
	return d.(chan struct{})
}

通过代码可以看见,ctx的done只有当被调用过Done方法时才会被创建,那么为什么这样子设计呢?很容易想到主要目的就是为了节省了不必要的资源浪费,提高效率,在很多情况下创建context并不需要监听done通道,只有在需要时才被创建,符合go语言的设计理念,只有需要的时候才引入

3.2value方法的实现

func (c *cancelCtx) Value(key any) any {
	if key == &cancelCtxKey {
		return c
	}
	return value(c.Context, key)
}

若参数key与cancelCtxKey相符,则返回当前ctx本身

否则,就向父层 层层寻找

源码表现的非常晦涩,为了具体知道这个Value方法是做什么的,我们先来看对于cancelCtxKey的定义

// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int

cancelCtxKey是一个私有的、唯一的标识符,它用于返回cancelCtx它本身。

对于该cancelCtxKey具体使用场景,下面还会讲到

3.3创建cancelCtx的WithCancel

//WithCancel返回一个带有新的Done通道的父context的副本。只有当父ctx被关闭,或者返回的cancel方法被调用时,该ctx的Done通道才会被关闭。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

该方法返回了一个cancelCtx,以及关闭它的cancel方法。

在1.20版本中,新增了一个WithCancelCause方法,该方法返回了一个cancelctx和它的CancelCauseFunc,我们也来看一下

func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
	c := withCancel(parent)
	return c, func(cause error) { c.cancel(true, Canceled, cause) }
}

在代码方面基本和withCancel方法一致,但是返回的CancelCauseFunc可以用于给用户自定义ctx被取消的原因,例如

ctx, cancel := context.WithCancelCause(parent)
cancel(myError)
ctx.Err() // returns context.Canceled
context.Cause(ctx) // returns myError

接着来看withCancel

func withCancel(parent Context) *cancelCtx {
	if parent == nil { //若父ctx是nil,那么不能创建
		panic("cannot create context from nil parent")
	}
	c := &cancelCtx{}
	c.propagateCancel(parent, c) //传递cancel给新创建的ctx
	return c
}

主要来看propagateCancel做了什么:

//该方法主要用于建立父context和子context的联系,如果父context也是一个cancelCtx,它需要保证父context被取消时,子context也能跟着被取消。
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent //将父context内嵌入子context中

	done := parent.Done() //获取父ctx的Done通道
	if done == nil { //如果通道不存在,那么说明父context不是cancelCtx,不需要为它们两个之间建立联系,因为父context永远不会关闭。
		return // parent is never canceled
	}

	select {
	case <-done: //非堵塞地获取父done通道的状态,如果done通道以及被closed,那么这里会接受到一个零值,如果没有被closed,会执行default后面的语句。
		//接受到零值,说明done通道以及被关闭了,父context已经被取消,此时应该立即调用子context的cancel方法,取消子context。
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
	//否则,执行以下代码
    //此处判断父context是否是一个context包实现的cancelCtx,如何理解会在下文讲述该方法时继续说明
	if p, ok := parentCancelCtx(parent); ok {
		// 如果父context是一个cancelCtx,或者是从某个cancelCtx衍生出来的context
		p.mu.Lock() //加锁
		if p.err != nil {
			// 如果存在err,说明已经被取消,此时也应该调用child的cancel方法取消子ctx
			child.cancel(false, p.err, p.cause)
		} else {
			if p.children == nil { //如果这是p的第一个子cancelCtx,需要初始化map的内存
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}//绑定子ctx到父ctx的map中,用于当父ctx取消时,能通知所有的child跟着取消。
		}
		p.mu.Unlock()
		return
	}
	//如果父ctx不是cancelctx并且实现了AfterFuncer接口,即实现了AfterFunc方法(该方法会在ctx被取消后唯一一次调用),那么就需要为父ctx再设置一个afterFunc方法,用于取消child并且传递err和cause
	if a, ok := parent.(afterFuncer); ok {
		// parent implements an AfterFunc method.
		c.mu.Lock()
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})
		c.Context = stopCtx{
			Context: parent,
			stop:    stop, //这里的stop方法可以用于当父ctx调用afterFunc的时候,取消父ctx对cancel函数的调用(看需求)
		}
		c.mu.Unlock()
		return
	}
	//下面的情况为增加一个goroutine,监听父ctx自己实现的Done channel(用户自定义的)
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done()://如果监听到父ctx自定义实现的Done channel关闭时,就关闭child
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done(): //如果child先关闭,那么就立即释放协程,避免协程泄露
		}
	}()
}

propagateCancel方法主要的作用是,保证了对于父ctx被取消时,为了能及时取消子ctx,避免不必要的资源浪费,建立父ctx和子ctx之间的联系。使用流程图来表示该方法如下(省略了检查afterFunc)

在该方法中,比较难以理解的地方是第二个if,"if p, ok := parentCancelCtx(parent); ok"究竟做了什么,为此我们跟进源码查看:

// parentCancelCtx 返回父级对象的底层 cancelCtx。
// 它通过查找 parent.Value(&cancelCtxKey) 来找到最内层的 enclosing cancelCtx,然后检查 parent.Done() 是否与该 cancelCtx 匹配。(如果不匹配,则该 cancelCtx 可能已被封装在提供了不同 done 通道的自定义实现中,在这种情况下,我们不应该绕过它。)
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

三个if分别做了什么:

  • 第一个if:获取parentCtx的done channel,并且查看情况,若已经关闭或者父Ctx不可取消,此时返回false,回到propagate方法中最终会cancel掉child

  • 第二个if:可以理解为通过Value方法找到离parentCtx“最近的”cancelCtx p,一般情况下,如果parentCtx就是一个CancelCtx,这时候就是parentCtx它本身。如果p不存在,也就是没有可以取消的ctx,此时也会返回false。

  • 第三个if:找到了p后,还读取了p的done channel,这时候一般情况,pdone 当然会 == done,因此最终会返回p和true,那么什么时候会不相等呢?为什么会不相等呢?为此,看下方的层次图来理解这个if

在这个情况假设下,ParentCtx2是从ParentCtx1衍生出来的,ParentCtx1是一个标准的CancelCtx,而ParentCtx2是一个用户自定义了ctx,它内层继承了ParentCtx1并且自己实现了Done()方法,这时候代码中的“p”找到的就是ParentCtx1,而parent是ParentCtx2,此时的done和pdone就是两个不同的channel了,这时候cihldctx应该监听哪一个done channel呢?答案是监听用户自定义实现的Ctx的chennel,因为我们不应该绕过用户实现的Done channel,这更加符合ctx到层次逻辑。假如这时候不去判断pdone == done,直接返回的指针就是ParentCtx1的指针了。

3.4Cancel方法实现

接下来我们来看cancel方法是如何实现的。

//cancel 关闭 c.done,cancel 每一个c的children。如果removeFromParent为true,将会把c从parentCtx的child中移除。
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	if err == nil { //必须存在err
		panic("context: internal error: missing cancel error")
	}
	if cause == nil { //没有自定义设置cause,默认为err
		cause = err
	}
	c.mu.Lock() //上锁
	if c.err != nil { //double check
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	c.cause = cause
	d, _ := c.done.Load().(chan struct{})
	if d == nil { //done没有被创建,直接存储context包内以及创建了的关闭的channel,不需要再次创建
		c.done.Store(closedchan)
	} else {
		close(d) //关闭done
	}
	for child := range c.children { //关闭每一个子context
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err, cause)
	}
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
		removeChild(c.Context, c) //从父child中移除自己
	}
}

代码比较简单,问题在于什么时候removFromParent为true呢?为什么为true?

回到withCancel方法中,我们可以看到返回的cancel方法中,此时removeFromParent就为true

return c, func() { c.cancel(true, Canceled, nil) }

当用户主动调用cancel()时,就会将子ctx从父ctx中的child删除。因为此时没有必要再在父ctx中接受父或祖先的cancel通知。而当调用cancel函数内部,对child执行的cancel就为false,这是因为后面设置了c.children = nil,这时候是从父ctx的方向关闭了子ctx对其的链接。

4afterFuncCtx

afterFuncCtx是1.20版本后引入的新ctx,它的作用是当ctx被取消后,能执行一次自定义的F函数,一般用于回收资源等。

type afterFuncCtx struct {
	cancelCtx
	once sync.Once // either starts running f or stops f from running
	f    func()
}

可以看到afterFuncCtx embed了cancelCtx,在此基础上添加了once和f。once保证了f只会被执行一次。接着我们来看如何实现一个afterFuncCtx

4.1AfterFunc

func AfterFunc(ctx Context, f func()) (stop func() bool) {
	a := &afterFuncCtx{ //创建局部变量afterFuncCtx,记录cancel时需要被执行的func
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
	return func() bool { //返回stop方法,如果需要不执行func,则调用stop
		stopped := false
		a.once.Do(func() { //尝试stop func函数,如果成功则stop会被设置为true
			stopped = true
		})
		if stopped { //stop后,a没有存在的意义,进行取消。
			a.cancel(true, Canceled, nil)
		}
		return stopped //true为成功取消,false表示f已经被执行或者正在被执行。
	}
}

这是一个闭包实现,闭包是指在函数内部定义的函数(如这里返回的 stop 函数),它会“捕获”并保存定义时可访问的所有外部变量。在 AfterFunc 方法中,虽然 a 是局部变量,但返回的 stop 方法引用了 a,形成了一个闭包,闭包会将 a 的内存保留在堆上,即使 AfterFunc 方法返回后,a 依然存在。所以当虽然没有返回a,但是返回的stop方法任然能调用a,a的生命周期超出了afterFunc方法。

以下是一个示例

package main

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

func main() {
	// 创建一个 2 秒后超时的 context
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	// 注册 AfterFunc,context 完成时将调用清理操作
	stop := context.AfterFunc(ctx, func() {
		fmt.Println("清理操作正在执行...")
	})

	// 等待 3 秒
	time.Sleep(3 * time.Second)

	// 尝试停止清理操作
	stopped := stop()
	if stopped {
		fmt.Println("成功停止清理操作")
	} else {
		fmt.Println("清理操作已经开始或已停止")
	}

	// 等待 context 超时
	<-ctx.Done()

	fmt.Println("程序结束")
}

4timerCtx

接下来来看timerCtx

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

	deadline time.Time
}

timer直接内嵌了cancelCtx以实现Done和Err,并且新添了timer和deadline字段,deadline用于查看ctx的截止时间,timer用于完成过时取消ctx。

4.1WithDeadline

先来看WithTimeout方法

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

WithTimeout接受一个存活时间来创建一个timerCtx,可以看到最终都是调用了WithDeadline

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	return WithDeadlineCause(parent, d, nil)
}

WithDeadline又调用了WithDeadlineCause,返回了timerCtx和一个CancelFunc。

4.2WithDeadlineCause

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if cur, ok := parent.Deadline(); ok && cur.Before(d) { //如果parent的deadline更早,则直接返回parent的副本,不需要再创建timer
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	c := &timerCtx{
		deadline: d,
	}
	c.cancelCtx.propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {//检测是否在创建过程中已经过了ddl
		c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
		return c, func() { c.cancel(false, Canceled, nil) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() { //设置超时自动cancel
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }//返回timerCtx和cancelfunc
}

比较好理解,所以接着往下看cancel方法

4.3.cancel

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
	c.cancelCtx.cancel(false, err, cause)
	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()
}

这里值得注意的地方只有第二行,它调用的是c.cancelCtx.cancel,取消的并非是父ctx,而是它本身,这里的作用是为了区分cancel方法的实现,c.cancelCtx是父ctx的一个副本,并不是父ctx,所以真正的parent是c.cancelCtx.Context。

5valueCtx

type valueCtx struct {
	Context
	key, val any
}

valueCtx内嵌了Context接口所以拥有该接口的所有方法,以及添加了k-v pair。

5.1WithValue

// 提供的键必须可比较,并且不应是字符串或任何其他内置类型,以避免在使用上下文时与其他包发生冲突。使用WithValue的用户应为其键定义自己的类型。为了避免在将值赋给接口{}时分配内存,上下文键通常具有具体的类型struct{}。或者,应将导出的上下文键变量的静态类型设置为指针或接口。
func WithValue(parent Context, key, val any) Context {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if key == nil {
		panic("nil key")
	}
	if !reflectlite.TypeOf(key).Comparable() { //必须保证key是可比较的
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

5.2value

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) { //对context进行类型断言
		case *valueCtx:
			if key == ctx.key { //如果key就是当前ctx的key,则直接返回val
				return ctx.val
			}
			c = ctx.Context //否则向父ctx查询
		case *cancelCtx:
			if key == &cancelCtxKey {
				return c
			}
			c = ctx.Context
		case withoutCancelCtx:
			if key == &cancelCtxKey {
				// This implements Cause(ctx) == nil
				// when ctx is created using WithoutCancel.
				return nil
			}
			c = ctx.c
		case *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case backgroundCtx, todoCtx:
			return nil //不存在key value
		default:
			return c.Value(key) //返回用户实现的Value
		}
	}
}

可以看见,valueCtx查询value的过程,类似于链表查询,它是自底向上的,并且时间复杂度为O(N),它并不适用于存放大量的kv数据,原因有以下:

  • 线性时间复杂度O(N),耗时太长
  • 一个 valueCtx 实例只能存一个 kv 对,因此 n 个 kv 对会嵌套 n 个 valueCtx,造成空间浪费
  • 不支持基于 k 的去重,相同 k 可能重复存在,并基于起点的不同,返回不同的 v. 由此得知,valueContext 的定位类似于请求头,只适合存放少量作用域较大的全局 meta 数据.

感谢观看,参考博客:

Golang context 实现原理 (qq.com)

深入Go:Context-腾讯云开发者社区-腾讯云 (tencent.com)

Go context的使用和源码分析_&cancelctxkey-CSDN博客

标签:Context,parent,cancelCtx,ctx,cancel,源码,context,Go
From: https://www.cnblogs.com/MelonTe/p/18493295

相关文章