sync: 同步 v1.0.4 开源代码并脱敏

internal/library/rabbitmq/rabbitmq.go

@@ -241,24 +241,29 @@ func publishWithConfirm(ctx context.Context, ch *amqp.Channel, queue string, mes
 
 // Consume 启动消费者，使用手动 Ack，handler 成功才确认消息。
 func Consume(ctx context.Context, queue, consumer string, handler Handler) error {
+	_, err := ConsumeWithDone(ctx, queue, consumer, handler)
+	return err
+}
+
+func ConsumeWithDone(ctx context.Context, queue, consumer string, handler Handler) (<-chan struct{}, error) {
 	ch, err := Channel(ctx)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	if err = declareQueue(ch, queue); err != nil {
 		_ = ch.Close()
-		return err
+		return nil, err
 	}
 	if err = ch.Qos(1, 0, false); err != nil {
 		_ = ch.Close()
-		return gerror.Wrap(err, "rabbitmq set qos failed")
+		return nil, gerror.Wrap(err, "rabbitmq set qos failed")
 	}
 
 	// autoAck=false，业务处理成功后手动 Ack，处理失败时 Nack 并重新入队。
 	deliveries, err := ch.Consume(queue, consumer, false, false, false, false, nil)
 	if err != nil {
 		_ = ch.Close()
-		return gerror.Wrap(err, "rabbitmq consume failed")
+		return nil, gerror.Wrap(err, "rabbitmq consume failed")
 	}
 
 	consumerState := &managedConsumer{
@@ -269,6 +274,8 @@ func Consume(ctx context.Context, queue, consumer string, handler Handler) error
 	registerConsumer(consumerState)
 	go func() {
 		defer close(consumerState.done)
+		defer ch.Close()
+		defer unregisterConsumer(consumerState)
 		for {
 			if ctx.Err() != nil {
 				return
@@ -302,7 +309,7 @@ func Consume(ctx context.Context, queue, consumer string, handler Handler) error
 			}
 		}
 	}()
-	return nil
+	return consumerState.done, nil
 }
 
 // ConsumeJSON 启动 JSON 消费者，统一完成反序列化和 Ack/Nack 处理。
@@ -423,3 +430,16 @@ func cloneHeaders(headers amqp.Table) map[string]any {
 func isClosedError(err error) bool {
 	return err == nil || errors.Is(err, amqp.ErrClosed)
 }
+
+func unregisterConsumer(target *managedConsumer) {
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	for index, consumer := range consumers {
+		if consumer == target {
+			remaining := make([]*managedConsumer, 0, len(consumers)-1)
+			remaining = append(remaining, consumers[:index]...)
+			consumers = append(remaining, consumers[index+1:]...)
+			return
+		}
+	}
+}

internal/library/rabbitmq/rabbitmq_test.go

@@ -149,3 +149,22 @@ func resetRabbitMQState(t *testing.T) {
 		connMu.Unlock()
 	})
 }
+
+func TestUnregisterConsumerPreservesShutdownSnapshot(t *testing.T) {
+	consumerMu.Lock()
+	original := consumers
+	first, second, third := &managedConsumer{}, &managedConsumer{}, &managedConsumer{}
+	consumers = []*managedConsumer{first, second, third}
+	snapshot := consumers
+	consumerMu.Unlock()
+	defer func() { consumerMu.Lock(); consumers = original; consumerMu.Unlock() }()
+	unregisterConsumer(second)
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	if len(consumers) != 2 || consumers[0] != first || consumers[1] != third {
+		t.Fatal("consumer not removed")
+	}
+	if snapshot[1] != second || snapshot[2] != third {
+		t.Fatal("shutdown snapshot mutated")
+	}
+}