go context 实现分析

接口定义

type Context interface {
	// Deadline 返回在此上下文中代表工作应该被取消的时间。当没有设置截止时间时，Deadline 返回 ok==false。连续调用 Deadline 返回相同的结果。
	Deadline() (deadline time.Time, ok bool)

	// Done 返回一个通道，在此上下文中代表工作应该被取消时会关闭。如果此上下文永远不会被取消，Done 可能返回 nil。连续调用 Done 返回相同的值。
	// Done 通道的关闭可能是异步的，在取消函数返回后发生。
	//
	// WithCancel 会在调用 cancel 时将 Done 关闭；
	// WithDeadline 会在截止时间过期时将 Done 关闭；
	// WithTimeout 会在超时时间过去时将 Done 关闭。
	//
	// Done 用于在 select 语句中使用：
	//
	//  // Stream 通过 DoSomething 生成值并发送到 out，直到 DoSomething 返回错误或 ctx.Done 被关闭。
	//  func Stream(ctx context.Context, out chan<- Value) error {
	//  	for {
	//  		v, err := DoSomething(ctx)
	//  		if err != nil {
	//  			return err
	//  		}
	//  		select {
	//  		case <-ctx.Done():
	//  			return ctx.Err()
	//  		case out <- v:
	//  		}
	//  	}
	//  }
	//
	// 更多如何使用 Done 通道进行取消的示例请参见 https://blog.golang.org/pipelines 。
	Done() <-chan struct{}

	// 如果 Done 还没有关闭，则 Err 返回 nil。
	// 如果 Done 关闭了，则 Err 返回一个非 nil 错误解释原因：
	// 如果上下文被取消则返回 Canceled；
	// 如果上下文的截止时间过期则返回 DeadlineExceeded。
	// 在 Err 返回非 nil 错误后，连续的 Err 调用会返回相同的错误。
	Err() error

	// Value 返回与此上下文关联的键的值，如果该键没有关联值则返回 nil。连续对于相同键的 Value 调用会返回相同的结果。
	//
	// 仅用于在进程和 API 边界传递请求范围的数据中使用上下文值，而不是将可选参数传递给函数。
	//
	// 一个键标识上下文中的特定值。希望将值存储在上下文中的函数通常会在全局变量中分配一个键，然后将该键作为参数传递给 context.WithValue 和 Context.Value。键可以是支持相等性的任何类型；包应该将键定义为未导出类型，以避免冲突。
	//
	// 定义上下文键的包应该为使用该键存储的值提供类型安全的访问器：
	//
	// 	// Package user 定义了存储在上下文中的 User 类型。
	// 	package user
	//
	// 	import "context"
	//
	// 	// User 是存储在上下文中的值的类型。
	// 	type User struct {...}
	//
	// 	// key 是在此包中定义的键的未导出类型。
	// 	// 这样可以避免与其他包中定义的键冲突。
	// 	type key int
	//
	// 	// userKey 是上下文中 user.User 值的键。它是未导出的；客户端使用 user.NewContext 和 user.FromContext，而不是直接使用此键。
	// 	var userKey key
	//
	// 	// NewContext 返回一个带有值 u 的新上下文。
	// 	func NewContext(ctx context.Context, u *User) context.Context {
	// 		return context.WithValue(ctx, userKey, u)
	// 	}
	//
	// 	// FromContext 返回存储在 ctx 中的 User 值（如果有）。
	// 	func FromContext(ctx context.Context) (*User, bool) {
	// 		u, ok := ctx.Value(userKey).(*User)
	// 		return u, ok
	// 	}
	Value(key any) any
}

这段代码是 Go 语言标准库中的 Context 接口的定义，提供了一种用于控制多个 goroutine 之间的取消、超时和请求范围值传递的机制，可以通过它实现对 goroutine 生命周期的管理、取消、超时控制，以及传递请求范围的值。

公共方法

下面是几个context提供的公共方法

// WithCancel 返回带有新的 Done 通道的父上下文的副本。
// 当调用返回的取消函数时，或者当父上下文的 Done 通道关闭时（以先发生者为准），
// 返回的上下文的 Done 通道将被关闭。
//
// 取消此上下文将释放与其关联的资源，因此代码应在此上下文中的操作完成后立即调用取消函数。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

// WithDeadline 返回父上下文的副本，并将截止时间调整为不晚于 d。
// 如果父上下文的截止时间已经比 d 更早，则 WithDeadline(parent, d) 在语义上等同于 parent。
// 返回的上下文的 Done 通道在截止时间到期时关闭，或者在调用返回的取消函数时，或者在父上下文的 Done 通道关闭时，以先发生者为准。
//
// 取消此上下文将释放与其关联的资源，因此代码应在此上下文中的操作完成后立即调用取消函数。
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// 当前截止时间已经比新截止时间更早。
		return WithCancel(parent)
	}
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, nil) // 截止时间已经过期
		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, nil)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

// WithValue 返回父上下文的副本，其中与键关联的值为 val。
//
// 仅用于传输进程和API的请求范围数据，而不用于将可选参数传递给函数。
//
// 提供的键必须是可比较的，并且不应为字符串或任何其他内置类型，以避免上下文使用的包之间的冲突。
// WithValue 的用户应为键定义自己的类型。为了在分配给 interface{} 时避免分配，上下文键通常具有具体类型 struct{}。
// 或者，导出的上下文键变量的静态类型应为指针或接口。
func WithValue(parent Context, key, val any) Context {
	if parent == nil {
		panic("无法从空父上下文创建上下文")
	}
	if key == nil {
		panic("空键")
	}
	if !reflectlite.TypeOf(key).Comparable() {
		panic("键不可比较")
	}
	return &valueCtx{parent, key, val}
}


// WithTimeout 返回 WithDeadline(parent, time.Now().Add(timeout)) 的结果。
//
// 取消此上下文会释放与之关联的资源，因此代码应在此上下文中运行的操作完成后尽快调用 cancel：
//
//	func slowOperationWithTimeout(ctx context.Context) (Result, error) {
//		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
//		defer cancel()  // 如果 slowOperation 在超时时间到达之前完成，则释放资源
//		return slowOperation(ctx)
//	}
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

WithCancel

首先看WithCancel，返回cancelCtx，和取消的方法，cancelCtx通过done来进行通信，进行不同ctx的取消同步，通过children来存储所有要通知的子ctx，用context来存储父ctx。

type cancelCtx struct {
	Context

	mu       sync.Mutex            // 保护以下字段
	done     atomic.Value          // chan struct{} 类型，延迟创建，在第一次调用 cancel 时关闭
	children map[canceler]struct{} // 在第一次调用 cancel 时设置为 nil
	err      error                 // 在第一次调用 cancel 时设置为非 nil
	cause    error                 // 在第一次调用 cancel 时设置为非 nil
}

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

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

func (c *cancelCtx) Err() error {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	return err
}
// cancel 方法关闭 c.done，并取消 c 的每个子上下文，如果 removeFromParent 为 true，则从其父上下文的子列表中移除 c。
// 如果这是第一次取消 c，则将 cause 设置为 err。
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 // 已经取消了
	}
	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 {
		// 注意：在持有父上下文锁的同时获取子上下文的锁。
		child.cancel(false, err, cause)
	}
	c.children = nil
	c.mu.Unlock()

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

WithTimeout和WithDeadline

WithTimeout用WithDeadline来进行实现返回timerCtx，timerCtx 包含一个定时器和一个截止时间。它嵌入了一个 cancelCtx 来实现 Done 和 Err。它通过停止定时器然后委托给 cancelCtx.cancel 来实现取消。

// timerCtx 包含一个定时器和一个截止时间。它嵌入了一个 cancelCtx 来实现 Done 和 Err。
// 它通过停止定时器然后委托给 cancelCtx.cancel 来实现取消。
type timerCtx struct {
	*cancelCtx
	timer *time.Timer // 在 cancelCtx.mu 下。

	deadline time.Time
}

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

func (c *timerCtx) String() string {
	return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
		c.deadline.String() + " [" +
		time.Until(c.deadline).String() + "])"
}

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

WithValue

valueCtx 携带一个键值对。它为该键实现了 Value，并将所有其他调用委托给嵌入的 Context。

// valueCtx 携带一个键值对。它为该键实现了 Value，并将所有其他调用委托给嵌入的 Context。
type valueCtx struct {
	Context
	key, val any
}

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

使用方式

package main

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()
	context.WithCancel(context.Background())
	context.WithDeadline(context.Background(), time.Now().Add(20*time.Second))
	ctx2 := context.WithValue(ctx, "hhh", "vvvv")
	c := ctx2.Value("hhh")
	fmt.Println(c)
	go handle(ctx, 500*time.Millisecond)
	select {
	case <-ctx.Done():
		fmt.Println("main", ctx.Err())
	}
}

func handle(ctx context.Context, duration time.Duration) {
	select {
	case <-ctx.Done():
		fmt.Println("handle", ctx.Err())
	case <-time.After(duration):
		fmt.Println("process request with", duration)
	}
}