深入理解 Go Context：从原理到实战（基于 Go 1.26）

// 模拟：用户请求一个接口，后端派 goroutine 去查数据库
func handleRequest() {
    go queryDB("SELECT * FROM orders") // 耗时 5 秒
    go callRPC("user-service")         // 耗时 3 秒
    go fetchCache("hot-items")         // 耗时 1 秒
}

func main() {
    for i := 0; i < 100; i++ {
        handleRequest() // 每个请求派生 3 个 goroutine
        // 假设用户全部在 0.5 秒后断开
    }
    time.Sleep(1 * time.Second)
    fmt.Println("当前 goroutine 数:", runtime.NumGoroutine())
    // 输出: 当前 goroutine 数: 301
    // 300 个 goroutine 在做无用功，没人要它们的结果了
}

func handleRequest(ctx context.Context) {
    go queryDB(ctx, "SELECT * FROM orders")
    go callRPC(ctx, "user-service")
    go fetchCache(ctx, "hot-items")
}

// 用户断开 → ctx 被取消 → 三个 goroutine 收到信号 → 立即退出
// goroutine 数: 1（只剩 main）

type Context interface {
    // 返回 context 的截止时间。如果没有设置截止时间，ok 为 false
    Deadline() (deadline time.Time, ok bool)

    // 返回一个 channel，当 context 被取消或超时时，该 channel 会被关闭
    Done() <-chan struct{}

    // 返回 context 被取消的原因
    Err() error

    // 从 context 中获取 key 对应的值
    Value(key any) any
}

// 作为整棵 context 树的根节点，通常在 main、init 或顶层请求入口使用
ctx := context.Background()

// 当你还不确定该用什么 context 时的占位符
ctx := context.TODO()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

func main() {
    ctx, cancel := context.WithCancel(context.Background())

    go worker(ctx, "worker-1")
    go worker(ctx, "worker-2")

    time.Sleep(3 * time.Second)
    fmt.Println("主协程：通知所有 worker 停止")
    cancel() // 手动发出取消信号，同时释放资源

    time.Sleep(1 * time.Second) // 等待 worker 退出
}

func worker(ctx context.Context, name string) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("%s: 收到取消信号，退出。原因: %v\n", name, ctx.Err())
            return
        default:
            fmt.Printf("%s: 工作中...\n", name)
            time.Sleep(1 * time.Second)
        }
    }
}

worker-1: 工作中...
worker-2: 工作中...
worker-1: 工作中...
worker-2: 工作中...
worker-1: 工作中...
worker-2: 工作中...
主协程：通知所有 worker 停止
worker-1: 收到取消信号，退出。原因: context canceled
worker-2: 收到取消信号，退出。原因: context canceled

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

func queryDatabase(ctx context.Context) (string, error) {
    // 模拟一个可能很慢的数据库查询
    select {
    case <-time.After(5 * time.Second): // 查询需要 5 秒
        return "查询结果", nil
    case <-ctx.Done():
        return "", fmt.Errorf("查询被取消: %w", ctx.Err())
    }
}

func main() {
    // 设置 2 秒超时
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    result, err := queryDatabase(ctx)
    if err != nil {
        fmt.Println("错误:", err) // 输出: 错误: 查询被取消: context deadline exceeded
        return
    }
    fmt.Println("结果:", result)
}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)

// 标准库源码
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

func main() {
    deadline := time.Now().Add(3 * time.Second)

    ctx, cancel := context.WithDeadline(context.Background(), deadline)
    defer cancel()

    // 检查还剩多少时间
    if d, ok := ctx.Deadline(); ok {
        fmt.Printf("截止时间: %v, 剩余: %v\n", d, time.Until(d))
    }
}

func WithValue(parent Context, key, val any) Context

// 推荐使用自定义类型作为 key，避免冲突
type contextKey string

const (
    keyTraceID contextKey = "traceID"
    keyUserID  contextKey = "userID"
)

func middleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // 注入 traceID
        traceID := generateTraceID()
        ctx := context.WithValue(r.Context(), keyTraceID, traceID)

        // 注入用户信息
        userID := authenticate(r)
        ctx = context.WithValue(ctx, keyUserID, userID)

        next(w, r.WithContext(ctx))
    }
}

func handler(w http.ResponseWriter, r *http.Request) {
    traceID, _ := r.Context().Value(keyTraceID).(string)
    userID, _ := r.Context().Value(keyUserID).(string)
    fmt.Fprintf(w, "trace=%s, user=%s", traceID, userID)
}

// 场景：一个请求内，DB 和 Redis 必须在同一时刻前完成
deadline := time.Now().Add(5 * time.Second)

// 用 WithDeadline：两个操作共享截止时间，语义清晰
dbCtx, c1 := context.WithDeadline(ctx, deadline)
defer c1()
redisCtx, c2 := context.WithDeadline(ctx, deadline)
defer c2()

// 用 WithTimeout：需要各自计算剩余时间，容易出现微小偏差
dbCtx, c1 := context.WithTimeout(ctx, 5*time.Second)    // 从现在开始 5 秒
defer c1()
// ... 中间可能已经过了几毫秒
redisCtx, c2 := context.WithTimeout(ctx, 5*time.Second) // 又从现在开始 5 秒，比 DB 晚
defer c2()

// 普通版：只知道超时了
ctx, cancel := context.WithTimeout(parentCtx, 2*time.Second)
defer cancel()
// ctx.Err() → "context deadline exceeded"  但不知道是哪个环节

// Cause 版：精确定位
ctx, cancel := context.WithTimeoutCause(parentCtx, 2*time.Second,
    fmt.Errorf("调用用户服务超时"))
defer cancel()
// context.Cause(ctx) → "调用用户服务超时"  一目了然

func handler(w http.ResponseWriter, r *http.Request) {
    // 脱离请求生命周期，但给异步任务设置独立超时
    asyncCtx := context.WithoutCancel(r.Context())
    asyncCtx, cancel := context.WithTimeout(asyncCtx, 10*time.Second)

    go func() {
        defer cancel()
        writeAuditLog(asyncCtx, "user accessed resource")
    }()
}

Background (根)
├── WithCancel (请求级)
│   ├── WithTimeout (数据库查询, 2s)
│   ├── WithTimeout (Redis 查询, 1s)
│   └── WithCancel (下游 RPC 调用)
│       ├── WithTimeout (服务A, 3s)
│       └── WithTimeout (服务B, 3s)

func main() {
    // 父 context: 5 秒超时
    parent, cancelParent := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancelParent()

    // 子 context: 2 秒超时 — 实际生效 2 秒
    child1, cancel1 := context.WithTimeout(parent, 2*time.Second)
    defer cancel1()

    // 子 context: 10 秒超时 — 实际生效 5 秒（受父 context 限制）
    child2, cancel2 := context.WithTimeout(parent, 10*time.Second)
    defer cancel2()

    select {
    case <-child1.Done():
        fmt.Println("child1 超时") // 2 秒后触发
    }

    select {
    case <-child2.Done():
        fmt.Println("child2 超时") // 5 秒后触发（跟随父 context）
    }
}

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "time"
)

type contextKey string

const keyRequestID contextKey = "requestID"

// 模拟数据库查询
func fetchUser(ctx context.Context, id string) (map[string]string, error) {
    dbCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
    defer cancel()

    select {
    case <-time.After(1 * time.Second): // 正常情况 1 秒返回
        return map[string]string{"id": id, "name": "张三"}, nil
    case <-dbCtx.Done():
        return nil, fmt.Errorf("数据库查询超时: %w", dbCtx.Err())
    }
}

// 模拟调用下游服务
func fetchOrders(ctx context.Context, userID string) ([]string, error) {
    select {
    case <-time.After(500 * time.Millisecond):
        return []string{"order-001", "order-002"}, nil
    case <-ctx.Done():
        return nil, fmt.Errorf("获取订单失败: %w", ctx.Err())
    }
}

func userHandler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    reqID := ctx.Value(keyRequestID)

    log.Printf("[%s] 开始处理请求", reqID)

    // 并发查询用户信息和订单
    type userResult struct {
        data map[string]string
        err  error
    }
    type orderResult struct {
        data []string
        err  error
    }

    userCh := make(chan userResult, 1)
    orderCh := make(chan orderResult, 1)

    go func() {
        data, err := fetchUser(ctx, "123")
        userCh <- userResult{data, err}
    }()

    go func() {
        data, err := fetchOrders(ctx, "123")
        orderCh <- orderResult{data, err}
    }()

    // 等待结果，同时监听 context 取消
    var user map[string]string
    var orders []string

    for i := 0; i < 2; i++ {
        select {
        case ur := <-userCh:
            if ur.err != nil {
                http.Error(w, ur.err.Error(), http.StatusInternalServerError)
                return
            }
            user = ur.data
        case or := <-orderCh:
            if or.err != nil {
                http.Error(w, or.err.Error(), http.StatusInternalServerError)
                return
            }
            orders = or.data
        case <-ctx.Done():
            log.Printf("[%s] 请求被取消: %v", reqID, ctx.Err())
            return
        }
    }

    resp := map[string]any{"user": user, "orders": orders}
    json.NewEncoder(w).Encode(resp)
    log.Printf("[%s] 请求处理完成", reqID)
}

// 中间件：注入 requestID + 设置总超时
func withRequestContext(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
        defer cancel()

        reqID := fmt.Sprintf("req-%d", time.Now().UnixNano()) // 示例简化，生产环境建议用 UUID
        ctx = context.WithValue(ctx, keyRequestID, reqID)

        next(w, r.WithContext(ctx))
    }
}

func main() {
    http.HandleFunc("/user", withRequestContext(userHandler))
    log.Println("服务启动在 :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

ctx, cancel := context.WithTimeout(parentCtx, 5*time.Second)
defer cancel() // 即使操作提前完成，也要释放资源

// 正确
func GetUser(ctx context.Context, id string) (*User, error)

// 错误 — 不要把 context 放在其他位置
func GetUser(id string, ctx context.Context) (*User, error)

// 错误 — 不要把 context 放在结构体里
type Service struct {
    ctx context.Context // 不推荐
}

// 错误 — 把业务参数塞进 context
ctx = context.WithValue(ctx, "userID", userID)
ctx = context.WithValue(ctx, "page", 1)
ctx = context.WithValue(ctx, "pageSize", 20)
result := queryUsers(ctx)

// 正确 — 业务参数走函数签名
result := queryUsers(ctx, userID, page, pageSize)

func processItems(ctx context.Context, items []Item) error {
    for i, item := range items {
        // 每次迭代检查是否被取消
        if ctx.Err() != nil {
            return fmt.Errorf("处理在第 %d 项中断: %w", i, ctx.Err())
        }
        process(item)
    }
    return nil
}

// 请求总超时 10 秒
reqCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()

// 数据库查询 3 秒（在请求超时内）
dbCtx, dbCancel := context.WithTimeout(reqCtx, 3*time.Second)
defer dbCancel()

// 缓存查询 1 秒（在请求超时内）
cacheCtx, cacheCancel := context.WithTimeout(reqCtx, 1*time.Second)
defer cacheCancel()

ctx, cancel := context.WithCancelCause(parentCtx)

// 取消时传入具体原因
cancel(fmt.Errorf("上游服务 %s 返回 503", serviceName))

// 获取取消原因
cause := context.Cause(ctx)
fmt.Println(cause) // "上游服务 xxx 返回 503"

// Err() 仍然返回 context.Canceled，但 Cause() 返回你传入的错误
fmt.Println(ctx.Err()) // context canceled

// WithTimeoutCause：相对时间版本
ctx, cancel := context.WithTimeoutCause(
    parentCtx,
    2*time.Second,
    fmt.Errorf("调用支付服务超时"),
)
defer cancel()

select {
case <-time.After(5 * time.Second):
    fmt.Println("完成")
case <-ctx.Done():
    fmt.Println(context.Cause(ctx)) // "调用支付服务超时"
}

// WithDeadlineCause：绝对时间版本，用法类似
// ctx, cancel := context.WithDeadlineCause(parentCtx, time.Now().Add(3*time.Second), cause)

// 场景：请求结束后仍需执行的异步任务（如写审计日志）
func handler(w http.ResponseWriter, r *http.Request) {
    // 请求结束后 r.Context() 会被取消
    // 但审计日志需要继续执行
    auditCtx := context.WithoutCancel(r.Context())
    go writeAuditLog(auditCtx, "user accessed resource")
}

ctx, cancel := context.WithCancel(parentCtx)
defer cancel()

stop := context.AfterFunc(ctx, func() {
    log.Println("context 被取消，执行清理工作")
    cleanup()
})
defer stop() // 如果不再需要回调，可以提前取消注册

func TestUserService(t *testing.T) {
    // 不再需要手动创建 context
    // t.Context() 会在测试结束时自动取消
    ctx := t.Context()

    user, err := service.GetUser(ctx, "123")
    if err != nil {
        t.Fatal(err)
    }
    // ...
}

// Go 1.26: 用 errors.AsType 替代 errors.As，更简洁
ctx, cancel := context.WithCancelCause(parentCtx)
cancel(&TimeoutError{Service: "payment", Duration: 3 * time.Second})

// 类型安全地提取 cause
if te, ok := errors.AsType[*TimeoutError](context.Cause(ctx)); ok {
    log.Printf("服务 %s 超时 %v", te.Service, te.Duration)
}

// 内存泄漏！timer 不会被释放
ctx, _ := context.WithTimeout(parentCtx, 5*time.Second)

// 危险：不同包可能用相同的 string key，导致冲突
ctx = context.WithValue(ctx, "userID", "123")

// 安全：自定义未导出类型，包级别隔离
type ctxKey struct{}
ctx = context.WithValue(ctx, ctxKey{}, "123")

ctx, cancel := context.WithTimeout(parentCtx, 1*time.Second)
defer cancel()

conn, err := db.Conn(ctx)
if err != nil {
    return err
}
// 此时 ctx 可能已经超时
// 用已取消的 ctx 执行查询会立即失败
rows, err := conn.QueryContext(ctx, "SELECT ...") // 可能立即返回错误

conn, err := db.Conn(ctx)
if err != nil {
    return err
}
queryCtx, queryCancel := context.WithTimeout(context.Background(), 3*time.Second)
defer queryCancel()
rows, err := conn.QueryContext(queryCtx, "SELECT ...")

Context (接口)
├── emptyCtx        → Background() / TODO() 返回（值类型，零分配，Done() 返回 nil）
├── cancelCtx       → WithCancel() 返回（取消机制的核心）
│   ├── timerCtx    → WithDeadline() / WithTimeout() 返回
│   └── afterFuncCtx → AfterFunc() 内部使用
├── valueCtx        → WithValue() 返回（链表式键值存储）
└── withoutCancelCtx → WithoutCancel() 返回（切断取消链，保留 Value）

// 源码
type cancelCtx struct {
    Context                          // 嵌入父 context（形成链表/树）

    mu       sync.Mutex             // 保护以下字段
    done     atomic.Value           // chan struct{}，懒创建，首次 cancel 时关闭
    children map[canceler]struct{}  // 子 context 集合，首次 cancel 时置 nil
    err      atomic.Value           // 首次 cancel 时设置
    cause    error                  // 取消原因（Cause 系列 API 使用）
}

// 源码
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 {
    if err := c.err.Load(); err != nil {
        <-c.Done()  // 确保 done channel 已关闭
        return err.(error)
    }
    return nil
}

// 源码（简化注释）
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
    if cause == nil {
        cause = err
    }
    c.mu.Lock()
    if c.err.Load() != nil {
        c.mu.Unlock()
        return // 已经取消过了，幂等
    }
    c.err.Store(err)
    c.cause = cause

    // 关闭 done channel
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan) // 复用全局已关闭 channel
    } else {
        close(d)
    }

    // 递归取消所有子 context
    for child := range c.children {
        child.cancel(false, err, cause)
    }
    c.children = nil
    c.mu.Unlock()

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

// 源码（简化）
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
    c.Context = parent

    done := parent.Done()
    if done == nil {
        return // 父 context 永远不会取消（如 Background），无需传播
    }

    select {
    case <-done:
        child.cancel(false, parent.Err(), Cause(parent))
        return // 父已取消，立即取消子
    default:
    }

    if p, ok := parentCancelCtx(parent); ok {
        // 父是标准 cancelCtx：直接加入 children map（高效）
        p.mu.Lock()
        if err := p.err.Load(); err != nil {
            child.cancel(false, err.(error), p.cause)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else if a, ok := parent.(afterFuncer); ok {
        // 父实现了 AfterFunc 接口：用 AfterFunc 注册取消回调
        stop := a.AfterFunc(func() {
            child.cancel(false, parent.Err(), Cause(parent))
        })
        c.Context = stopCtx{Context: parent, stop: stop}
    } else {
        // 父是自定义 Context 实现：启动 goroutine 监听
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err(), Cause(parent))
            case <-child.Done():
            }
        }()
    }
}

// 源码
type timerCtx struct {
    cancelCtx
    timer    *time.Timer
    deadline time.Time
}

// 源码（简化）
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
    // 关键优化：如果父的 deadline 更早，直接退化为 WithCancel
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        return WithCancel(parent)
    }

    c := &timerCtx{deadline: d}
    c.cancelCtx.propagateCancel(parent, c)

    dur := time.Until(d)
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded, cause) // 已过期，立即取消
        return c, func() { c.cancel(false, Canceled, nil) }
    }

    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err.Load() == nil {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded, cause)
        })
    }
    return c, func() { c.cancel(true, Canceled, nil) }
}

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) // 向上查找
}

type withoutCancelCtx struct {
    c Context  // 注意：不是嵌入
}

func (withoutCancelCtx) Done() <-chan struct{} { return nil }  // 永不取消
func (withoutCancelCtx) Err() error            { return nil }
func (c withoutCancelCtx) Value(key any) any   { return value(c, key) } // 仍继承值