context 与 goroutine 泄漏：取消信号如何穿过调用链

// leak-send/main.go
package main

import (
    "fmt"
    "runtime"
    "time"
)

func main() {
    for i := 0; i < 1000; i++ {
        ch := make(chan int)
        go func() {
            // 模拟慢查询，3 秒后才有结果
            time.Sleep(3 * time.Second)
            ch <- 42 // 关键：调用方早已超时返回，没人接收
        }()
        // 假装请求超时，立即离开作用域
        _ = ch
    }
    time.Sleep(5 * time.Second)
    fmt.Printf("goroutines: %d\n", runtime.NumGoroutine())
}

func query(ctx context.Context) (int, error) {
    ch := make(chan int, 1) // buffered，避免发送方阻塞
    go func() {
        time.Sleep(3 * time.Second)
        ch <- 42
    }()
    select {
    case v := <-ch:
        return v, nil
    case <-ctx.Done():
        return 0, ctx.Err() // context 取消后，工作 goroutine 的 send 也会因 receiver 提前 return 而解锁
    }
}

// leak-receive/main.go
package main

import (
    "fmt"
    "runtime"
)

func main() {
    ch := make(chan int)
    go func() {
        ch <- 1
        // 生产者提前 return，忘了 close(ch)
    }()
    // 消费者永远等下去
    v := <-ch
    fmt.Println(v)
    time.Sleep(2 * time.Second)
    fmt.Printf("goroutines: %d\n", runtime.NumGoroutine())
}

func producer(out chan<- int) {
    defer close(out) // 谁发谁关
    out <- 1
    out <- 2
}

func consumer(in <-chan int) {
    for v := range in { // 自动感知关闭
        fmt.Println(v)
    }
}

func watch(ctx context.Context, updates <-chan Event) {
    for {
        select {
        case e := <-updates:
            handle(e)
        // 缺一个 case <-ctx.Done(): return
        }
    }
}

func watch(ctx context.Context, updates <-chan Event) {
    for {
        select {
        case e := <-updates:
            handle(e)
        case <-ctx.Done():
            return // 收到取消信号，干净退出
        }
    }
}

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

type cancelCtx struct {
    Context                       // 父 context
    mu       sync.Mutex
    done     chan struct{}        // 关闭即广播
    children map[canceler]struct{} // 子节点集合
    err      error                // 取消原因
}

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err != nil {
        return // 已经取消过，幂等
    }
    c.err = err
    close(c.done) // 关键：关闭 done channel
    for child := range c.children {
        child.cancel(false, err) // 递归取消所有子节点
    }
    c.children = nil
    if removeFromParent {
        removeChild(c.Context, c) // 从父的 children 摘掉
    }
}

parent, _ := context.WithTimeout(context.Background(), 5*time.Second)
child, _ := context.WithTimeout(parent, 10*time.Second) // 想放更宽？没门

fmt.Println(parent.Deadline()) // 5s
fmt.Println(child.Deadline())  // 5s，被父收紧

ctx, cancel := context.WithTimeout(parent, 5*time.Second)
go doWork(ctx)
// 忘了 defer cancel() 或没把 cancel 传出去

ctx, _ := context.WithTimeout(context.Background(), time.Second) // 警告：lostcancel
doSomething(ctx)

ctx = context.WithValue(ctx, traceIDKey, "abc-123")
traceID, _ := ctx.Value(traceIDKey).(string)

package trace

type traceIDKey struct{} // 未导出类型，包外无法构造

func WithTraceID(ctx context.Context, id string) context.Context {
    return context.WithValue(ctx, traceIDKey{}, id)
}

func TraceID(ctx context.Context) string {
    id, _ := ctx.Value(traceIDKey{}).(string)
    return id
}

// pipeline/main.go
package main

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

func gen(ctx context.Context, nums ...int) <-chan int {
    out := make(chan int)
    go func() {
        defer close(out)
        for _, n := range nums {
            select {
            case out <- n:
            case <-ctx.Done():
                return // 收到取消，立即退出
            }
        }
    }()
    return out
}

func square(ctx context.Context, in <-chan int) <-chan int {
    out := make(chan int)
    go func() {
        defer close(out)
        for n := range in {
            select {
            case out <- n * n:
            case <-ctx.Done():
                return
            }
        }
    }()
    return out
}

func merge(ctx context.Context, ins ...<-chan int) <-chan int {
    out := make(chan int)
    var wg sync.WaitGroup
    wg.Add(len(ins))
    for _, in := range ins {
        go func(in <-chan int) {
            defer wg.Done()
            for n := range in {
                select {
                case out <- n:
                case <-ctx.Done():
                    return
                }
            }
        }(in)
    }
    go func() {
        wg.Wait()
        close(out) // 发送方负责关闭
    }()
    return out
}

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

    in := gen(ctx, 2, 3, 4, 5)
    c1 := square(ctx, in)
    c2 := square(ctx, in)
    out := merge(ctx, c1, c2)

    // 拿到第一个结果就取消整条 pipeline
    fmt.Println(<-out) // 4
    cancel()

    // 短暂等待，让所有 goroutine 看到 ctx.Done()
    time.Sleep(100 * time.Millisecond)
    fmt.Println("done")
}

var goroutineBaseline int

func init() {
    goroutineBaseline = runtime.NumGoroutine()
}

// 在监控里周期性上报
func report() {
    current := runtime.NumGoroutine()
    metrics.Gauge("goroutine_count", current)
    metrics.Gauge("goroutine_delta", current-goroutineBaseline)
}

curl http://localhost:6060/debug/pprof/goroutine?debug=2 > goroutine.txt

func before() { log.Println(runtime.NumGoroutine()) } // 启动后基线
func after()  { log.Println(runtime.NumGoroutine()) } // 压测 1 万请求后

func TestNoLeak(t *testing.T) {
    before()
    for i := 0; i < 10000; i++ {
        doRequest()
    }
    runtime.GC()
    time.Sleep(time.Second) // 给泄漏的 goroutine 一点时间（如果有）
    after()
    // 断言 after 接近 before
}