goroutine

rabbitmq/consumer.go

@@ -19,7 +19,8 @@ type Consumer struct {
 	prefetchCount int  // QoS: 预取数量（并发控制）
 	autoAck       bool // 是否自动确认
 	handler       MessageHandler
-	workerCount   int // worker 数量
+	workerCount   int                // worker 数量
+	cancel        context.CancelFunc // 用于停止 worker
 }
 
 // ConsumerOption 消费者配置选项
@@ -74,6 +75,9 @@ func NewConsumer(queue string, handler MessageHandler, opts ...ConsumerOption) *
 
 // Start 启动消费者
 func (c *Consumer) Start(ctx context.Context) error {
+	// 创建可取消的 context
+	workerCtx, cancel := context.WithCancel(ctx)
+	c.cancel = cancel
 	ch, err := GetChannel()
 	if err != nil {
 		return err
@@ -108,7 +112,7 @@ func (c *Consumer) Start(ctx context.Context) error {
 
 	// 启动多个 worker
 	for i := 0; i < c.workerCount; i++ {
-		go c.worker(ctx, i, msgs)
+		go c.worker(workerCtx, i, msgs)
 	}
 
 	return nil
@@ -118,29 +122,40 @@ func (c *Consumer) Start(ctx context.Context) error {
 func (c *Consumer) worker(ctx context.Context, workerID int, msgs <-chan amqp.Delivery) {
 	g.Log().Debugf(ctx, "Worker %d 已启动", workerID)
 
-	for msg := range msgs {
-		// 处理消息
-		err := c.handler(ctx, msg.Body)
-
-		if err != nil {
-			g.Log().Errorf(ctx, "Worker %d 处理消息失败: %v", workerID, err)
-
-			// 如果不是自动确认，需要手动 Nack
-			if !c.autoAck {
-				// requeue=false: 不重新入队，进入死信队列
-				msg.Nack(false, false)
-			}
-		} else {
-			// 处理成功，手动确认
-			if !c.autoAck {
-				msg.Ack(false)
+	for {
+		select {
+		case <-ctx.Done():
+			// Context 取消，退出
+			g.Log().Infof(ctx, "Worker %d 收到停止信号，正在退出", workerID)
+			return
+		case msg, ok := <-msgs:
+			if !ok {
+				// Channel 关闭，退出
+				g.Log().Infof(ctx, "Worker %d 消息通道已关闭，退出", workerID)
+				return
 			}
 
-			g.Log().Debugf(ctx, "Worker %d 处理消息成功", workerID)
+			// 处理消息
+			err := c.handler(ctx, msg.Body)
+
+			if err != nil {
+				g.Log().Errorf(ctx, "Worker %d 处理消息失败: %v", workerID, err)
+
+				// 如果不是自动确认，需要手动 Nack
+				if !c.autoAck {
+					// requeue=false: 不重新入队，进入死信队列
+					msg.Nack(false, false)
+				}
+			} else {
+				// 处理成功，手动确认
+				if !c.autoAck {
+					msg.Ack(false)
+				}
+
+				g.Log().Debugf(ctx, "Worker %d 处理消息成功", workerID)
+			}
 		}
 	}
-
-	g.Log().Debugf(ctx, "Worker %d 已停止", workerID)
 }
 
 // StartTypedConsumer 启动类型化消费者（自动反序列化）
@@ -163,3 +178,12 @@ func StartTypedConsumer[T any](
 	consumer := NewConsumer(queue, wrappedHandler, opts...)
 	return consumer.Start(ctx)
 }
+
+// Stop 停止消费者
+func (c *Consumer) Stop(ctx context.Context) {
+	if c.cancel != nil {
+		g.Log().Infof(ctx, "正在停止消费者: queue=%s", c.queue)
+		c.cancel()
+		c.cancel = nil
+	}
+}