Go 并发原语 - Context

// 1. 根 Context（不可取消，通常作为起点）
ctx := context.Background()  // 主函数、初始化、测试
ctx := context.TODO()        // 不确定用哪个时的占位符

// 2. 可取消的 Context
ctx, cancel := context.WithCancel(parentCtx)
defer cancel() // ✅ 必须调用，释放资源

// 3. 带超时的 Context
ctx, cancel := context.WithTimeout(parentCtx, 5*time.Second)
defer cancel()

// 4. 带截止时间的 Context
deadline := time.Now().Add(5 * time.Second)
ctx, cancel := context.WithDeadline(parentCtx, deadline)
defer cancel()

func fetchData(ctx context.Context) (string, error) {
    // 模拟一个耗时操作
    ch := make(chan string, 1)
    go func() {
        time.Sleep(3 * time.Second) // 模拟耗时
        ch <- "data"
    }()

    select {
    case data := <-ch:
        return data, nil
    case <-ctx.Done():
        return "", ctx.Err() // 超时或被取消
    }
}

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

    data, err := fetchData(ctx)
    if err != nil {
        fmt.Println("失败:", err) // context deadline exceeded
        return
    }
    fmt.Println("成功:", data)
}

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

    // 启动多个 worker
    for i := 0; i < 5; i++ {
        go worker(ctx, i)
    }

    time.Sleep(2 * time.Second)
    cancel() // 通知所有 worker 停止
    time.Sleep(time.Second) // 等待 worker 退出
}

func worker(ctx context.Context, id int) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("Worker %d: 收到取消信号，退出\n", id)
            return
        default:
            fmt.Printf("Worker %d: 工作中...\n", id)
            time.Sleep(500 * time.Millisecond)
        }
    }
}

type contextKey string

const traceIDKey contextKey = "traceID"

func withTraceID(ctx context.Context, traceID string) context.Context {
    return context.WithValue(ctx, traceIDKey, traceID)
}

func getTraceID(ctx context.Context) string {
    if v, ok := ctx.Value(traceIDKey).(string); ok {
        return v
    }
    return "unknown"
}

func handleRequest(ctx context.Context) {
    traceID := getTraceID(ctx)
    fmt.Printf("[%s] 处理请求\n", traceID)
    queryDB(ctx)
}

func queryDB(ctx context.Context) {
    traceID := getTraceID(ctx)
    fmt.Printf("[%s] 查询数据库\n", traceID)
}

ctx3 := WithValue(ctx2, "k3", "v3")
ctx2 := WithValue(ctx1, "k2", "v2")
ctx1 := WithValue(bg, "k1", "v1")

// ❌ 忘记调用 cancel，导致 goroutine 和 timer 泄漏
func handle() {
    ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
    // cancel 被丢弃了 💀
    doWork(ctx)
}

// ✅ 始终 defer cancel
func handle() {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel() // 即使操作提前完成，也释放 timer 资源
    doWork(ctx)
}

// ❌ 用 string 做 key，可能和其他包冲突
ctx = context.WithValue(ctx, "userID", 123)

// ✅ 用自定义未导出类型做 key，避免冲突
type contextKey string
const userIDKey contextKey = "userID"
ctx = context.WithValue(ctx, userIDKey, 123)

// ❌ 把业务参数放进 Context
ctx = context.WithValue(ctx, "orderID", orderID)
processOrder(ctx) // 函数签名看不出需要 orderID

// ✅ 业务参数用函数参数传递
processOrder(ctx, orderID) // 依赖关系一目了然

// ❌ 不检查 Done，cancel 调了也没用
func doWork(ctx context.Context) {
    for i := 0; i < 1000000; i++ {
        heavyComputation(i) // 死干，不看取消信号
    }
}

// ✅ 在循环中定期检查
func doWork(ctx context.Context) {
    for i := 0; i < 1000000; i++ {
        select {
        case <-ctx.Done():
            return // 收到取消信号，退出
        default:
        }
        heavyComputation(i)
    }
}

// WithoutCancel：创建一个不会被父 Context 取消的子 Context
// 适合在父请求结束后还需要继续的后台任务
ctx := context.WithoutCancel(parentCtx)

// AfterFunc：Context 取消时执行回调
stop := context.AfterFunc(ctx, func() {
    fmt.Println("ctx 被取消了，执行清理")
})
// stop() 可以阻止回调执行

// WithCancelCause（Go 1.20+）：取消时附带原因
ctx, cancel := context.WithCancelCause(parentCtx)
cancel(fmt.Errorf("用户主动取消")) // 附带原因

// 获取取消原因
err := context.Cause(ctx) // "用户主动取消"

// 超时层级示例
// HTTP handler: 10s
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
defer cancel()

// 数据库查询: 3s（小于 handler 的 10s）
dbCtx, dbCancel := context.WithTimeout(ctx, 3*time.Second)
defer dbCancel()

# 检查 Context 泄漏
go vet ./...
go test -race ./...