Go context 传递实战：超时、取消和元数据的最佳实践指南

sequenceDiagram
    participant Client as HTTP 客户端
    participant Handler as HTTP Handler
    participant Service as 业务服务层
    participant DB as 数据库
    participant LLM as 大模型 API
    Client->>Handler: GET /api/chat?q=xxx
    Handler->>Handler: ctx = context.WithTimeout(5s)
    Handler->>Handler: ctx = WithTraceID(ctx, "abc123")
    Handler->>Service: ProcessMessage(ctx, msg)
    Service->>Service: ctx, cancel = WithTimeout(3s)
    Service->>DB: QueryHistory(ctx, userID)
    alt 数据库查询超时
        DB-->>Service: context.DeadlineExceeded
        Service->>Service: cancel()
        Service-->>Handler: Err: timeout
        Handler-->>Client: 504 Gateway Timeout
    else 查询成功
        DB-->>Service: results
        Service->>LLM: CallModel(ctx, prompt)
        alt 客户端断开连接
            Handler->>Handler: ctx.Done()
            Service->>LLM: (cancel 传播)
            LLM-->>Service: context.Canceled
        end
    end

package service
import (
    "context"
    "database/sql"
    "errors"
    "fmt"
    "net/http"
    "time"
)
// ========== 模式 1：多层超时控制 ==========
// ProcessChatMessage 处理聊天消息，展示多层超时的正确用法
func ProcessChatMessage(ctx context.Context, userID, message string) (string, error) {
    // 第一层：整体请求超时（由 HTTP Handler 设置）
    // 这里继承上游的超时，不再重复设置
    // 第二层：数据库查询子任务超时
    dbCtx, dbCancel := context.WithTimeout(ctx, 2*time.Second)
    defer dbCancel() // 关键：用完立即释放资源
    history, err := queryChatHistory(dbCtx, userID)
    if err != nil {
        if errors.Is(err, context.DeadlineExceeded) {
            return "", fmt.Errorf("查询历史记录超时，请稍后重试")
        }
        return "", fmt.Errorf("查询历史失败: %w", err)
    }
    // 第三层：模型推理子任务超时
    llmCtx, llmCancel := context.WithTimeout(ctx, 5*time.Second)
    defer llmCancel()
    resp, err := callLLM(llmCtx, history, message)
    if err != nil {
        return "", fmt.Errorf("模型推理失败: %w", err)
    }
    return resp, nil
}
// ========== 模式 2：Fan-out 并发控制 ==========
// BatchInference 批量推理，任意一个子任务失败则取消所有
func BatchInference(ctx context.Context, prompts []string) ([]string, error) {
    // 创建可取消的子 context，任一子任务失败时取消其余
    ctx, cancel := context.WithCancel(ctx)
    defer cancel()
    type result struct {
        idx  int
        text string
        err  error
    }
    resultCh := make(chan result, len(prompts))
    for i, p := range prompts {
        go func(idx int, prompt string) {
            r := result{idx: idx}
            // 在 goroutine 内部再次包装超时
            taskCtx, taskCancel := context.WithTimeout(ctx, 3*time.Second)
            defer taskCancel()
            resp, err := callLLM(taskCtx, nil, prompt)
            if err != nil {
                r.err = err
                cancel() // 任一失败，通知其他 goroutine 停止
            } else {
                r.text = resp
            }
            resultCh <- r
        }(i, p)
    }
    // 收集结果
    results := make([]string, len(prompts))
    var firstErr error
    for i := 0; i < len(prompts); i++ {
        r := <-resultCh
        if r.err != nil && firstErr == nil {
            firstErr = r.err
        }
        results[r.idx] = r.text
    }
    if firstErr != nil {
        return nil, firstErr
    }
    return results, nil
}
// ========== 模式 3：元数据传递 ==========
type contextKey string
const (
    KeyTraceID   contextKey = "trace_id"
    KeyUserID    contextKey = "user_id"
    KeyRequestID contextKey = "request_id"
)
// WithTraceID 往 context 中注入 TraceID
func WithTraceID(ctx context.Context, traceID string) context.Context {
    return context.WithValue(ctx, KeyTraceID, traceID)
}
// TraceIDFromContext 从 context 中提取 TraceID
func TraceIDFromContext(ctx context.Context) string {
    if v, ok := ctx.Value(KeyTraceID).(string); ok {
        return v
    }
    return "unknown"
}
// HTTP 中间件：从请求头提取并注入元数据
func TraceMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ctx := r.Context()
        // 优先使用上游传入的 TraceID
        traceID := r.Header.Get("X-Trace-ID")
        if traceID == "" {
            traceID = generateTraceID()
        }
        ctx = WithTraceID(ctx, traceID)
        // 设置整体请求超时
        ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
        defer cancel()
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}
// ========== 辅助函数 ==========
type Message struct {
    Role    string
    Content string
}
func queryChatHistory(ctx context.Context, userID string) ([]Message, error) {
    // 检查 context 是否已取消，避免无意义查询
    select {
    case <-ctx.Done():
        return nil, ctx.Err()
    default:
    }
    // 实际数据库查询...
    _ = userID
    return nil, nil
}
func callLLM(ctx context.Context, history []Message, prompt string) (string, error) {
    // 检查是否已取消
    select {
    case <-ctx.Done():
        return "", ctx.Err()
    default:
    }
    // 实际 LLM 调用...
    return "response", nil
}
func generateTraceID() string {
    return fmt.Sprintf("%d", time.Now().UnixNano())
}