如何修复流水线取消时的Goroutine泄漏问题

你遇到的核心问题是：当下游流水线阶段（比如sum）提前退出或context超时后，上游阶段的worker goroutine因阻塞在结果通道的发送操作上，无法响应context的取消信号，导致残留运行。以下是标准的解决思路和实现：

关键问题根源 #

1. 发送阻塞无法响应取消：当下游不再接收结果通道的数据时，上游worker会阻塞在result <- do(val)操作上，此时即使context取消，select也无法切换到<-ctx.Done()分支。
2. 取消信号未及时传播：仅使用顶层context时，下游阶段完成后无法主动通知上游停止工作，导致上游继续处理无用任务。
3. 生成器未响应取消：原始generate函数不会监听context，会持续发送数据，浪费资源。

解决方案实现 #

1. 改进生成器，使其响应context #

让数据生成阶段在context取消时立即停止发送并关闭通道：

func generate(ctx context.Context, amount int) <-chan int {
    result := make(chan int)
    go func() {
        defer close(result)
        for i := 0; i < amount; i++ {
            select {
            case <-ctx.Done():
                return
            case result <- i:
            }
        }
    }()
    return result
}

2. 优化process函数，避免阻塞并及时响应取消 #

在任务执行、结果发送前都检查context状态，同时给结果通道设置缓冲减少阻塞概率：

func process[T any, R any](ctx context.Context, workers int, input <-chan T, do func(T) R) <-chan R {
    wg := new(sync.WaitGroup)
    // 缓冲大小设为worker数量，降低发送阻塞概率
    result := make(chan R, workers)

    for i := 0; i < workers; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            for {
                select {
                case <-ctx.Done():
                    return
                case val, ok := <-input:
                    if !ok {
                        return
                    }
                    // 先检查取消信号，避免执行无用任务
                    select {
                    case <-ctx.Done():
                        return
                    default:
                    }
                    res := do(val)
                    // 发送结果时监听取消，防止永久阻塞
                    select {
                    case <-ctx.Done():
                        return
                    case result <- res:
                    }
                }
            }
        }()
    }

    go func() {
        defer close(result)
        wg.Wait()
    }()
    return result
}

3. 传播取消信号到整个流水线 #

创建派生context，让下游阶段完成时主动取消整个流水线，确保所有上游worker及时退出：

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

    // 派生流水线专用context，sum完成或超时都触发取消
    pipelineCtx, pipelineCancel := context.WithCancel(ctx)
    defer pipelineCancel()

    input := generate(pipelineCtx, 1000)
    multiplied := process(pipelineCtx, 15, input, func(val int) int {
        time.Sleep(time.Second)
        return val * 2
    })
    increased := process(pipelineCtx, 15, multiplied, func(val int) int {
        return val + 10
    })

    var wg sync.WaitGroup
    wg.Add(1)
    var finalResult int
    go func() {
        defer wg.Done()
        finalResult = sum(increased)
        // sum完成后立即取消流水线，停止所有上游工作
        pipelineCancel()
    }()

    // 等待sum完成或超时
    select {
    case <-ctx.Done():
        pipelineCancel()
    case <-wg.Wait():
    }

    wg.Wait()
    fmt.Println("Result: ", finalResult)
    fmt.Println("Num goroutine: ", runtime.NumGoroutine())
}

核心优化点总结 #

• 取消信号全链路传播：通过派生context实现下游触发上游取消，确保流水线整体终止。
• 避免无用任务执行：在任务处理前检查context状态，减少资源浪费。
• 防止发送阻塞：结果通道加缓冲，发送时监听取消信号，避免worker永久阻塞。
• 生成器响应取消：停止无用数据生产，快速关闭上游通道。

内容的提问来源于stack exchange，提问作者onekram