Go 上下文：context 包与取消操作

在 Go 语言中，context 是一个看似简单却蕴含深意的标准库包。它解决的问题非常明确：如何在 goroutine 之间传递取消信号、截止时间以及请求作用域内的值。

当一个请求到达服务器，服务器可能需要启动多个 goroutine 来处理不同的子任务。如果客户端提前断开连接，这些子任务应该立即停止，而不是继续浪费 CPU 和内存资源。context 就是为解决这一问题而设计的。

本文将深入剖析 context 包的实现原理，重点讲解取消操作的机制，并通过实际代码演示其使用方法。

一、context 的本质与类型层次

context 的核心是一个接口，定义了四个方法：

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key any) any
}

Deadline 返回当前上下文是否有关联的超时时间。Done 返回一个只读通道，当上下文被取消时，这个通道会被关闭。Err 返回取消的原因。Value 用于在上下文中存储和获取键值对。

标准库提供了两组实现这个接口的类型。第一组是根上下文，它们没有父上下文：context.Background() 和 context.TODO()。前者是所有上下文的起点，通常用于主函数或测试；后者用于占位，当不确定该使用什么上下文时。第二组是派生上下文，它们通过 WithCancel、WithDeadline、WithTimeout 和 WithValue 从父上下文派生而来。

理解这个层次结构至关重要。context 形成了一棵树，取消信号会从父节点传播到所有子节点。当你取消一个父上下文，所有派生的子上下文都会收到信号并自动取消。

二、取消操作的实现原理

2.1 取消机制的核心数据结构

cancelCtx 是实现取消功能的基础结构：

type cancelCtx struct {
    Context
    mu       sync.Mutex
    done     chan struct{}
    children map[canceler]struct{}
    err      error
}

children 字段是一个 map，存储所有派生自当前上下文的子上下文。当父上下文被取消时，它会遍历这个 map，依次取消每个子上下文。这种设计确保了取消操作的传播性。

done 通道是被延迟初始化的。只有当有人调用 Done() 方法时，这个通道才会被创建。这种懒加载策略避免了不必要的资源开销。

2.2 取消的传播过程

当你调用 cancel() 函数时，标准的取消流程如下：

1. 加锁保护：cancelCtx 首先获取互斥锁，确保并发安全。
2. 设置错误：将 err 字段设置为 context.Canceled。
3. 关闭通道：关闭 done 通道，触发所有等待者的 <-chan 操作。
4. 遍历子节点：遍历 children map，递归取消每个子上下文。
5. 清理引用：将自身从父上下文的 children map 中移除。

这个过程的关键在于递归。取消操作会沿着 context 树向下传播，确保每一个后代上下文都能及时收到取消信号。

在上图展示的 context 树中，如果你取消根节点 A，树中所有节点 B、C、D、E、F 都会收到取消信号。

三、创建可取消的上下文

3.1 使用 WithCancel 创建可取消上下文

WithCancel 是最基础的取消上下文创建方式：

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

返回的 cancel 函数用于手动触发取消。以下是一个完整的示例：

package main

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

func main() {
    // 从 Background 创建一个可取消的上下文
    ctx, cancel := context.WithCancel(context.Background())

    // 启动一个模拟工作的 goroutine
    go func() {
        for {
            select {
            case <-ctx.Done():
                fmt.Println("Worker: 收到取消信号，退出")
                return
            default:
                fmt.Println("Worker: 工作中...")
                time.Sleep(500 * time.Millisecond)
            }
        }
    }()

    // 模拟主程序运行 2 秒后取消
    time.Sleep(2 * time.Second)
    fmt.Println("Main: 调用 cancel()")
    cancel()

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

运行这段代码，你会看到 worker 在主程序调用 cancel() 后立即退出。ctx.Done() 返回的通道在 cancel() 被调用后会被关闭，select 语句的 <-ctx.Done() 分支因此被触发。

3.2 使用 WithTimeout 设置超时取消

WithTimeout 基于时间实现自动取消，非常适合 HTTP 请求处理等场景：

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

以下示例演示了如何在数据库查询中应用超时：

func queryDatabase(ctx context.Context, query string) ([]byte, error) {
    // 模拟耗时查询
    done := make(chan []byte, 1)

    go func() {
        time.Sleep(3 * time.Second) // 模拟数据库操作
        done <- []byte("查询结果")
    }()

    select {
    case result := <-done:
        return result, nil
    case <-ctx.Done():
        return nil, ctx.Err()
    }
}

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    result, err := queryDatabase(ctx, "SELECT * FROM users")
    if err != nil {
        fmt.Printf("查询失败: %v\n", err)
    } else {
        fmt.Printf("查询成功: %s\n", result)
    }
}

在这个例子中，如果数据库查询超过 2 秒，ctx.Done() 会先返回，queryDatabase 函数会立即返回超时错误，而不是继续等待。

3.3 使用 WithDeadline 设置绝对截止时间

WithDeadline 与 WithTimeout 类似，但它接收的是绝对时间而非持续时长：

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

这在需要精确控制截止时间的场景中非常有用，例如定时任务的调度：

deadline := time.Date(2024, 12, 31, 23, 59, 59, 0, time.UTC)
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()

四、取消操作的常见模式

4.1 防止 goroutine 泄漏

goroutine 泄漏是 Go 程序中常见的问题。当父程序不再需要某个 goroutine 的结果，而该 goroutine 又没有收到结束信号时，就会发生泄漏：

func leakyFunction() {
    ctx, _ := context.WithCancel(context.Background())

    // 启动一个永远不会退出的 goroutine
    go func() {
        for {
            select {
            case <-ctx.Done():
                return
            default:
                // 模拟工作
                time.Sleep(time.Hour)
            }
        }
    }()

    // 函数返回后，ctx 没有被取消
    // 上面的 goroutine 将永远运行
}

正确的做法是确保在不再需要上下文时取消它：

func correctFunction() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel() // 确保在函数返回前取消

    go func() {
        for {
            select {
            case <-ctx.Done():
                return
            default:
                time.Sleep(time.Hour)
            }
        }
    }()
}

4.2 多任务并发取消

当一个请求需要并发执行多个子任务时，可以使用 context 统一管理所有子任务的取消：

func processMultipleTasks(ctx context.Context) error {
    var wg sync.WaitGroup
    results := make(chan string, 3)
    errors := make(chan error, 3)

    tasks := []func(ctx context.Context) error{
        task1,
        task2,
        task3,
    }

    for _, task := range tasks {
        wg.Add(1)
        go func(t func(ctx context.Context) error) {
            defer wg.Done()
            if err := t(ctx); err != nil {
                errors <- err
            } else {
                results <- "success"
            }
        }(task)
    }

    // 等待所有任务完成或任意一个出错
    go func() {
        wg.Wait()
        close(results)
        close(errors)
    }()

    for {
        select {
        case err := <-errors:
            return err
        case <-ctx.Done():
            return ctx.Err()
        case result := <-results:
            // 处理成功的结果
            fmt.Println("Task completed:", result)
        }
    }
}

func task1(ctx context.Context) error {
    time.Sleep(100 * time.Millisecond)
    return nil
}

func task2(ctx context.Context) error {
    time.Sleep(200 * time.Millisecond)
    return nil
}

func task3(ctx context.Context) error {
    time.Sleep(300 * time.Millisecond)
    return nil
}

4.3 在 HTTP Handler 中使用 Context

HTTP 服务器是使用 context 的典型场景。Go 的 net/http 包已经在 http.Request 中包含了上下文：

func handleRequest(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()

    // 检查请求是否已经被取消
    select {
    case <-ctx.Done():
        http.Error(w, "请求已取消", http.StatusRequestTimeout)
        return
    default:
    }

    // 从上下文获取请求 ID
    requestID := ctx.Value("request_id")

    // 执行实际的业务逻辑
    result, err := processWithContext(ctx, requestID)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }

    w.Write(result)
}

func processWithContext(ctx context.Context, requestID interface{}) ([]byte, error) {
    // 模拟处理
    select {
    case <-time.After(5 * time.Second):
        return []byte("处理完成"), nil
    case <-ctx.Done():
        return nil, ctx.Err()
    }
}

中间件可以方便地为每个请求添加超时和取消功能：

func timeoutMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
        defer cancel()

        // 使用新的上下文创建新的请求
        r = r.WithContext(ctx)

        next.ServeHTTP(w, r)
    })
}

五、传递请求作用域的值

5.1 使用 WithValue 存储上下文数据

WithValue 允许在上下文中存储键值对，这些值可以在整个请求处理链中传递：

func main() {
    // 创建带值的上下文
    ctx := context.WithValue(context.Background(), "user_id", 12345)
    ctx = context.WithValue(ctx, "request_id", "req-abc-123")

    // 在 handler 中获取值
    userID := ctx.Value("user_id")
    requestID := ctx.Value("request_id")

    fmt.Printf("User ID: %d, Request ID: %s\n", userID, requestID)
}

5.2 值的获取与类型断言

从上下文中获取值时需要进行类型断言：

func getUserFromContext(ctx context.Context) (int, bool) {
    val := ctx.Value("user_id")
    if val == nil {
        return 0, false
    }

    userID, ok := val.(int)
    return userID, ok
}

需要注意的是，上下文值的键应该使用自定义类型，以避免不同包之间的键冲突：

type contextKey string

const (
    userIDKey    contextKey = "user_id"
    requestIDKey contextKey = "request_id"
)

func processRequest(ctx context.Context) {
    userID := ctx.Value(userIDKey)
}

六、最佳实践与注意事项

6.1 Context 的传递规则

6.2 常见错误与避免方法

错误一：在 HTTP handler 中忘记使用 Context

// 错误示例
func handler(w http.ResponseWriter, r *http.Request) {
    db.Query("...") // 没有传递 context
}

应该使用 r.Context() 获取请求的上下文并传递：

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    rows, err := db.QueryContext(ctx, "...")
}

错误二：取消后仍然继续执行

ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()

result, err := doWork(ctx)
if err != nil {
    return err
}
// 错误：即使出错也继续执行后续操作
processResult(result)

错误三：混用不同的取消机制

尽量统一使用 context 的取消机制，避免在同一个调用链中混用 channel 和 context。

6.3 性能考量

Context 的设计非常高效，但仍需注意几点：Done() 方法返回的通道是延迟初始化的，只有调用该方法时才会创建。Value 的查找会沿着上下文树向上遍历，因此应该避免在热代码路径中频繁访问上下文值。

七、总结

context 包是 Go 并发编程的核心工具。它提供了优雅的取消传播机制，使得管理多个 goroutine 的生命周期变得简单可靠。

核心要点回顾：context.Context 是通过接口定义的；WithCancel、WithTimeout、WithDeadline 都可以创建可取消的上下文；取消信号通过关闭通道传播；上下文形成树形结构，取消会从父节点传播到所有子节点。

在实际开发中，始终将 context 作为函数的第一个参数传递，使用 defer 确保上下文被正确取消，并在 HTTP handler 中使用请求的上下文。