Golang context 的作用和实现原理

context的作用

Go语言中的context包提供了一种在进程中跨API和跨进程边界传递取消信号、超时和其他请求范围的值的方式。context的主要作用包括：

取消信号（Cancellation）：

• 当一个操作需要取消时，可以通过context传递一个取消信号。例如，一个长运行的任务可以通过检查context是否被取消来提前终止。

超时（Timeouts）：

• context可以用来设置操作的超时时间。如果操作在指定时间内没有完成，context会发送一个取消信号。

截止时间（Deadlines）：

• 类似于超时，截止时间允许你设置一个绝对时间点，超过这个时间点后，context会发送取消信号。

请求范围的值（Request-scoped values）：

• context可以在处理HTTP请求、数据库事务等过程中，存储请求相关的值，这些值可以跨API边界传递而不需要在每个函数调用中显式传递。

并发控制（Concurrency control）：

• 在并发编程中，context可以用来同步多个goroutine，确保它们在某个操作被取消时能够及时响应。

父子关系（Parent-child relationships）：

• context可以创建父子关系，子context会继承父context的取消信号。这在处理树状结构的并发任务时非常有用。

避免共享可变状态：

• 通过使用context，可以避免在多个goroutine之间共享可变状态，从而减少竞态条件和锁的使用。

API设计：

• context为API设计提供了一种标准的方式来处理取消和超时，使得API的使用更加一致和可预测。

使用context时，通常会创建一个context.Context类型的变量，然后通过context.WithCancel、context.WithDeadline、context.WithTimeout或context.WithValue等函数来创建一个新的context，这个新的context会包含额外的信息或取消功能。在实际的函数调用中，context.Context作为第一个参数传递，以确保可以在需要时访问和使用这些值或信号。

context的实现原理：

context.Context

type Context interface {
	Deadline() (deadline time.Time, ok bool)

	Done() <-chan struct{}

	Err() error

	Value(key any) any
}

Deadline()：

返回过期时间，第一个返回值为过期时间，第二个返回值代表过期时间是否存在。

Done() <-chan struct{}

监测context是否终止，返回值是一个只读的通道，只有当context终止时才能从通道读取到数据。

Err() error

返回错误，通常情况下是context过期或者被手动取消。

Value(key any) any

返回 context 中的对应 key 的值.

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 }

Canceled：

context被cancel时报错；

DeadlineExceeded：

context超时时报错

context.emptyCtx

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
}

Deadline:

返回一个默认时间和false表示不存在过期时间

Done：

返回一个nil值，用户读取时陷入阻塞状态

Err：

返回一个nil值，表示没有错误

Value：

返回一个nil值，表示没有值

context.Background&&context.TODO

type backgroundCtx struct{ emptyCtx }

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

type todoCtx struct{ emptyCtx }

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

func Background() Context {
	return backgroundCtx{}
}

func TODO() Context {
	return todoCtx{}
}

context.Background() 和 context.TODO() 方法返回的均是 emptyCtx 类型的一个实例.但是context.TODO()说明可能需要实现更加详细的context

cancelCtx

type canceler interface {
	cancel(removeFromParent bool, err, cause error)
	Done() <-chan struct{}
}

type cancelCtx struct {
	Context

	mu       sync.Mutex            
	done     atomic.Value          
	children map[canceler]struct{} 
	err      error                 
	cause    error                 
}

Context:

指向有且仅有一个的父节点

        Deadline：

                未实现该方法，不能直接调用

        Done：

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

        基于 atomic 包，读取 cancelCtx 中的 chan，倘若已存在，则直接返回，否则加互斥锁，然后再次读取，若存在直接返回，否则初始化chan储存到atomic.Value包中，并返回（懒加载：在应用程序初始化时，不会立即加载所有数据，而是等到需要使用数据时才进行加载）

        Err：

func (c *cancelCtx) Err() error {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	return err
}

        加锁，读取错误，解锁

        Value：

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

         如果key等于特定值 &cancelCtxKey，则返回cancelCtx自身的指针。否则遵循valueCtx的思路取值返回

mu：

互斥锁，用于并发保护

done：

用于反应生命周期

children:

因为map的value是struct所以，返回的是set而不是map

canceler:

子节点的信息

err：

存放错误信息

cause：

用于提供更详细的错误信息，指示context被取消的具体原因

context.WithCancel()

WithCancel

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 := &cancelCtx{}
	c.propagateCancel(parent, c)
	return c
}

        首先校验父context不为空，在 propagateCancel 方法内启动一个守护协程，以保证父 context 终止时，该 cancelCtx 也会被终止；返回cancelCtx以及用于终止cancelCtx的闭包函数。

propagateCancel

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent

	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	if p, ok := parentCancelCtx(parent); ok {
		// parent is a *cancelCtx, or derives from one.
		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()
		return
	}

	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,
		}
		c.mu.Unlock()
		return
	}

	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done():
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done():
		}
	}()
}

        首先将父context注入子context，如果父context是不会被cancel的类型，则直接返回。如果父context已经被cancel，则直接终止子context，并以父context的err作为子context的err。如果父context是cancelCtx的类型则加锁，并将子context添加到parent的children map当中。如果父context不是cancelCtx的类型，但是具备cancel的能力，则开启一个协程通过多路复用的方式监控 父contxet状态，倘若其终止，则同时终止子context，并透传父contxet的err。

parentCancelCtx

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
}

如果传入的父context是不能被cancel的节点则返回false，如果用特定的cancelCtxKey取值得到的不是本身，说明不是cancelCtx的类型，返回false（详见之前的Value）

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
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	for child := range c.children {
		// 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)
	}
}

        方法有两个入参，第一个表示表示当前 context 是否需要从父 context 的 children set 中删除，第二个表示要返回的具体的错误。

        首先校验是否存在err，若为空则 panic，然后检验cause是否为空，若为空则将传入的err赋值给cause，然后加锁判断当前cancelCtx的err是否为空，若不为空说明已经被cancel，直接返回。否则将传入的err赋值给当前cancelCtx的err，将cause赋值给当前cancelCtx的cause，处理当前cancelCtx的channel，若 channel 此前未初始化则直接注入一个 closedChan，否则关闭channel。遍历当前 cancelCtx 的 children set，依次将 children context 都进行cancel，解锁，根据传入的 removeFromParent flag 判断是否需要手动把 cancelCtx 从 parent 的 children set 中移除。

removeChild

func removeChild(parent Context, child canceler) {
	if s, ok := parent.(stopCtx); ok {
		s.stop()
		return
	}
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

timerCtx

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

	deadline time.Time
}

在cancelCtx的基础上又进行了一层封装。新增了一个 time.Timer 用于定时终止 context；另外新增了一个 deadline 字段用于字段 timerCtx 的过期时间.

timerCtx.Deadline

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
	return c.deadline, true
}

用于展示过期时间

timerCtx.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()
}

复用继承的 cancelCtx 的 cancel 能力，进行 cancel 处理，加锁，停止定时终止context，解锁。

context.WithTimeout & context.WithDeadline

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

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

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) {
		// 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 {
		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() {
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

WithTimeout传入过期所用的时间，WithDeadline传入过期时的时间。

首先检验父context是否为空，检验父context的过期时间是否早于自己，如果早于自己，则直接构造一个cancelCtx即可，否则构造一个timerCtx， 传入过期时间，启动守护方法，同步 parent 的 cancel 事件到子 context，判断过期时间是否已到，如果已到，直接cancel，返DeadlineExceeded错误。加锁，如果当前context的err为空，启动time.Timer，设定一个延时时间，即达到过期时间后会终止该timerCtx，并返回DeadlineExceeded错误，解锁，返回timerCtx，以及封装了cancel逻辑的闭包cancel函数

valueCtx

type valueCtx struct {
	Context
	key, val any
}

一个 valueCtx 中仅有一组 kv 对

valueCtx.Value

func (c *valueCtx) Value(key any) any {
	if c.key == key {
		return c.val
	}
	return value(c.Context, key)
}

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			if key == ctx.key {
				return ctx.val
			}
			c = ctx.Context
		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
		default:
			return c.Value(key)
		}
	}
}

假如当前 valueCtx 的 key 等于用户传入的 key，则直接返回其 value，否则从父context 中依次向上寻找，启动一个 for 循环，由下而上，由子及父，依次对 key 进行匹配，找到匹配的 key，则将该组 value 进行返回。

context.WithValue

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() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

如果父context为空panic，如果key为空panic，如果key的类型无法比较panic，包含父context和kv对，创建一个新的valueCtx并返回。