如何在取消Go协程前等待正在处理的操作完成

要实现消费者优雅停机，确保正在处理的消息不被中断，需要从跟踪正在处理的任务和等待任务完成后再退出两个核心点入手。以下是修改后的完整代码及说明：

修改后的消费者代码 #

package main

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

type Consumer struct {
    wg sync.WaitGroup // 跟踪正在处理的消息任务
}

func (c *Consumer) Start(ctx context.Context) {
    fmt.Println("Consumer is starting")
    defer func() {
        c.wg.Wait() // 等待所有正在处理的消息完成后再退出
        fmt.Println("Consumer is stopping")
    }()

    i := 0
    for {
        select {
        case <-ctx.Done():
            fmt.Println("Context is done, waiting for current message to finish")
            return // 退出循环，触发defer中的等待逻辑
        default:
            // 模拟Kafka消息获取（实际场景替换为kafka.FetchMessage(ctx)）
            fmt.Printf("[%d] Received message\n", i)
            
            c.wg.Add(1)
            // 同步处理消息，用闭包捕获当前索引
            func(idx int) {
                defer c.wg.Done() // 处理完成后标记任务结束
                <-time.After(3 * time.Second) // 模拟消息处理耗时
                fmt.Printf("[%d] Stopped processing\n", idx)
            }(i)
            i++
        }
    }
}

修改后的主函数代码 #

package main

import (
    "context"
    "fmt"
    "os"
    "os/signal"
    "syscall"
)

func main() {
    fmt.Println("Main Start")
    sigchan := make(chan os.Signal, 1)
    signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)
    
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    consumer := &Consumer{}
    consumerDone := make(chan struct{}) // 用于等待消费者完全退出
    
    // 启动消费者协程
    go func() {
        consumer.Start(ctx)
        close(consumerDone) // 消费者退出后关闭通道
    }()

    <-sigchan
    fmt.Println("Received signal to shutdown")
    cancel() // 通知消费者停止接收新消息
    <-consumerDone // 等待消费者处理完所有正在进行的任务并退出
    fmt.Println("Main Exit")
}

核心修改说明 #

1. 使用sync.WaitGroup跟踪任务：

   • 每次开始处理消息前调用wg.Add(1)，标记新增一个待完成的任务；
   • 消息处理完成后通过defer wg.Done()标记任务结束；
   • 消费者退出前的defer块中调用wg.Wait()，确保所有正在处理的消息都完成后才退出。

2. 主函数等待消费者完全退出：

   • 通过consumerDone通道监听消费者协程的退出状态；
   • 调用cancel()后，主进程不会立即退出，而是等待消费者处理完最后一条消息、关闭consumerDone通道后才终止，避免协程被强制中断。

3. 优化循环逻辑（可选）：
   实际场景中Kafka的消息获取接口（如FetchMessage）是阻塞式的，此时可以去掉default分支，改为先监听上下文取消，再执行阻塞获取消息的逻辑，避免CPU空转：

   for {
       select {
       case <-ctx.Done():
           fmt.Println("Context is done, waiting for current message to finish")
           return
       default:
       }
   
       // 阻塞获取Kafka消息，上下文取消时会返回错误
       msg, err := kafka.FetchMessage(ctx)
       if err != nil {
           if ctx.Err() != nil {
               break
           }
           fmt.Printf("Failed to fetch message: %v\n", err)
           continue
       }
   
       // 处理消息逻辑同前
       c.wg.Add(1)
       fmt.Printf("[%d] Received message\n", i)
       func(idx int) {
           defer c.wg.Done()
           // 处理消息
       }(i)
       i++
   }
   

内容的提问来源于stack exchange，提问作者Vinicius Barbosa de Medeiros