Go 源码深度解析之 context（四种具体实现 & 用法场景）

// cancelCtx 是可取消的上下文，支持主动取消和通知子上下文
type cancelCtx struct {
    Context  // 嵌入父上下文
    mu       sync.Mutex        // 互斥锁，保护以下字段的并发安全
    done     chan struct{}     // Done 方法返回的通道，关闭表示取消
    children map[canceler]struct{} // 存储所有子取消上下文，取消时通知它们
    err      error             // 取消原因（Canceled）
}

// canceler 接口定义了可取消的能力，cancelCtx 和 timerCtx 都实现了该接口
type canceler interface {
    cancel(removeFromParent bool, err error)
}

// Done 实现 Context 接口的 Done 方法：返回 done 通道
func (c *cancelCtx) Done() <-chan struct{} {
    c.mu.Lock()
    if c.done == nil {
        // 延迟初始化通道，避免未使用时的内存浪费
        c.done = make(chan struct{})
    }
    d := c.done
    c.mu.Unlock()
    return d
}

// Err 实现 Context 接口的 Err 方法：返回取消原因
func (c *cancelCtx) Err() error {
    c.mu.Lock()
    err := c.err
    c.mu.Unlock()
    return err
}

// cancel 是取消的核心方法：关闭通道、通知子上下文、清理父上下文关联
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    // 1. 校验取消原因非空
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    // 2. 已取消则直接返回
    if c.err != nil {
        c.mu.Unlock()
        return
    }
    // 3. 记录取消原因，关闭 done 通道（触发所有监听协程的通知）
    c.err = err
    close(c.done)
    // 4. 递归取消所有子上下文
    for child := range c.children {
        child.cancel(false, err)
    }
    // 5. 清空子上下文列表
    c.children = nil
    c.mu.Unlock()

    // 6. 若需要，从父上下文中移除自身（避免父上下文再次通知）
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

// WithCancel 基于父上下文创建可取消的子上下文
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    // 1. 校验父上下文非空
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    // 2. 创建 cancelCtx 实例
    c := newCancelCtx(parent)
    // 3. 将当前 cancelCtx 注册到父上下文的子列表中
    propagateCancel(parent, &c)
    // 4. 返回上下文和取消函数（CancelFunc 是 func() 类型）
    return &c, func() { c.cancel(true, Canceled) }
}

// newCancelCtx 初始化 cancelCtx
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

// timerCtx 是带超时/截止时间的上下文，继承 cancelCtx
type timerCtx struct {
    cancelCtx          // 嵌入 cancelCtx，继承取消能力
    mu       sync.Mutex // 保护定时器和截止时间
    timer    *time.Timer // 定时器，用于触发超时取消
    deadline time.Time  // 截止时间（若有）
}

// Deadline 实现 Context 接口的 Deadline 方法：返回截止时间
func (c *timerCtx) Deadline() (time.Time, bool) {
    c.mu.Lock()
    defer c.mu.Unlock()
    return c.deadline, !c.deadline.IsZero()
}

// cancel 重写 cancelCtx 的 cancel 方法：额外停止定时器
func (c *timerCtx) cancel(removeFromParent bool, err error) {
    // 1. 先调用父类 cancelCtx 的 cancel 方法，触发取消逻辑
    c.cancelCtx.cancel(false, err)
    // 2. 从父上下文中移除自身
    if removeFromParent {
        removeChild(c.cancelCtx.Context, c)
    }
    // 3. 停止定时器，避免资源泄漏
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

// WithTimeout 基于父上下文创建“超时自动取消”的子上下文
// timeout：超时时间（如 3*time.Second）
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    // 本质是调用 WithDeadline，截止时间为“当前时间 + 超时时间”
    return WithDeadline(parent, time.Now().Add(timeout))
}

// WithDeadline 基于父上下文创建“指定截止时间自动取消”的子上下文
// deadline：具体截止时间（如 time.Date(2025, 12, 10, 0, 0, 0, 0, time.UTC)）
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
    // 1. 校验父上下文非空
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    // 2. 若父上下文的截止时间更早，直接返回父上下文的取消能力（无需额外定时器）
    if d, ok := parent.Deadline(); ok && d.Before(deadline) {
        return WithCancel(parent)
    }
    // 3. 创建 timerCtx 实例
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  deadline,
    }
    // 4. 注册到父上下文，实现父子关联
    propagateCancel(parent, c)
    // 5. 计算当前到截止时间的差值
    dur := time.Until(deadline)
    if dur <= 0 {
        // 5.1 已超时，直接取消
        c.cancel(true, DeadlineExceeded)
        return c, func() { c.cancel(false, Canceled) }
    }
    // 5.2 未超时，启动定时器：超时后自动调用 cancel 方法
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 定时器触发时，调用 cancel 方法，原因是“超时”
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    // 6. 返回上下文和取消函数（手动取消时会停止定时器）
    return c, func() { c.cancel(true, Canceled) }
}

package main

import (
    "context"
    "fmt"
    "net/http"
    "time"

    "github.com/google/uuid"
)

// 生成全局唯一的 traceID
func generateTraceID() string {
    return uuid.NewString()
}

// 模拟数据库查询（内部协程，需要 traceID 和超时控制）
func queryDB(ctx context.Context, sql string) (string, error) {
    // 从上下文获取 traceID
    traceID, _ := ctx.Value("traceID").(string)
    fmt.Printf("traceID: %s, 执行 SQL: %s\n", traceID, sql)

    // 模拟查询耗时（3 秒）
    select {
    case <-time.After(3 * time.Second):
        return "查询结果：用户信息", nil
    case <-ctx.Done():
        // 感知超时或取消，返回错误
        return "", ctx.Err()
    }
}

// HTTP 处理函数
func handler(w http.ResponseWriter, r *http.Request) {
    // 1. 创建根上下文（emptyCtx），设置超时 5 秒（timerCtx）
    rootCtx := context.Background()
    ctx, cancel := context.WithTimeout(rootCtx, 5*time.Second)
    defer cancel() // 函数结束时手动取消，释放资源

    // 2. 传递 traceID（valueCtx）
    traceID := generateTraceID()
    ctx = context.WithValue(ctx, "traceID", traceID)

    // 3. 启动协程执行数据库查询
    resultChan := make(chan string, 1)
    errChan := make(chan error, 1)
    go func() {
        res, err := queryDB(ctx, "SELECT * FROM users WHERE id = 1")
        resultChan <- res
        errChan <- err
    }()

    // 4. 等待结果或超时
    select {
    case res := <-resultChan:
        err := <-errChan
        if err != nil {
            w.WriteHeader(http.StatusInternalServerError)
            w.Write([]byte(fmt.Sprintf("错误：%v", err)))
            return
        }
        w.Write([]byte(res))
    case <-ctx.Done():
        w.WriteHeader(http.StatusRequestTimeout)
        w.Write([]byte(fmt.Sprintf("超时：%v", ctx.Err())))
    }
}

func main() {
    http.HandleFunc("/user", handler)
    fmt.Println("服务启动：http://localhost:8080/user")
    http.ListenAndServe(":8080", nil)
}

package main

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

// 模拟任务处理（每个任务耗时不同）
func processTask(ctx context.Context, taskID int, wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Printf("任务 %d 开始执行\n", taskID)

    select {
    case <-time.After(time.Duration(taskID) * time.Second):
        // 任务完成：打印结果并返回
        fmt.Printf("任务 %d 执行完成\n", taskID)
    case <-ctx.Done():
        // 收到取消信号：打印取消信息
        fmt.Printf("任务 %d 被取消，原因：%v\n", taskID, ctx.Err())
    }
}

func main() {
    // 1. 创建可取消上下文（cancelCtx）
    rootCtx := context.Background()
    ctx, cancel := context.WithCancel(rootCtx)
    defer cancel()

    var wg sync.WaitGroup
    taskCount := 3

    // 2. 启动 3 个协程处理任务
    for i := 1; i <= taskCount; i++ {
        wg.Add(1)
        go processTask(ctx, i, &wg)
    }

    // 3. 监听任务完成：任一任务完成后取消所有任务
    doneChan := make(chan struct{})
    go func() {
        wg.Wait()
        close(doneChan)
    }()

    // 等待第一个任务完成（耗时 1 秒），然后取消其他任务
    select {
    case <-doneChan:
        // 所有任务完成（理论上不会走这里，因为第一个任务 1 秒完成）
    case <-time.After(1 * time.Second):
        fmt.Println("\n检测到任务完成，取消剩余任务")
        cancel() // 手动取消所有子协程
    }

    // 等待所有协程退出
    wg.Wait()
    fmt.Println("所有任务处理结束")
}