feat: add publish confirmation function

Author: zhufuyi

pkg/rabbitmq/producer.go

@@ -23,6 +23,9 @@ type producerOptions struct {
 	// If true, the message will be returned to the sender if the queue cannot be
 	// found according to its own exchange type and routeKey rules.
 	mandatory bool
+
+	// only for publish-subscribe mode
+	isPublisherConfirm bool
 }
 
 func (o *producerOptions) apply(opts ...ProducerOption) {
@@ -39,8 +42,9 @@ func defaultProducerOptions() *producerOptions {
 		queueBind:       defaultQueueBindOptions(),
 		deadLetter:      defaultDeadLetterOptions(),
 
-		isPersistent: true,
-		mandatory:    true,
+		isPersistent:       true,
+		mandatory:          true,
+		isPublisherConfirm: false,
 	}
 }
 
@@ -86,6 +90,13 @@ func WithProducerMandatory(enable bool) ProducerOption {
 	}
 }
 
+// WithPublisherConfirm enables publisher confirm.
+func WithPublisherConfirm() ProducerOption {
+	return func(o *producerOptions) {
+		o.isPublisherConfirm = true
+	}
+}
+
 // -------------------------------------------------------------------------------------------
 
 // Producer session
@@ -108,6 +119,9 @@ type Producer struct {
 	exchangeArgs  amqp.Table
 	queueArgs     amqp.Table
 	queueBindArgs amqp.Table
+
+	// only for publish-subscribe mode
+	isPublisherConfirm bool
 }
 
 // NewProducer create a producer

pkg/rabbitmq/publisher.go

@@ -2,6 +2,7 @@ package rabbitmq
 
 import (
 	"context"
+	"fmt"
 
 	amqp "github.com/rabbitmq/amqp091-go"
 	"go.uber.org/zap"
@@ -25,6 +26,15 @@ func NewPublisher(channelName string, connection *Connection, opts ...ProducerOp
 		return nil, err
 	}
 
+	// enable publisher confirm
+	if o.isPublisherConfirm {
+		err = ch.Confirm(false)
+		if err != nil {
+			_ = ch.Close()
+			return nil, err
+		}
+	}
+
 	// declare the exchange type
 	err = ch.ExchangeDeclare(
 		channelName,
@@ -61,7 +71,7 @@ func NewPublisher(channelName string, connection *Connection, opts ...ProducerOp
 }
 
 func (p *Publisher) Publish(ctx context.Context, body []byte) error {
-	return p.ch.PublishWithContext(
+	err := p.ch.PublishWithContext(
 		ctx,
 		p.Exchange.name,
 		p.Exchange.routingKey,
@@ -73,6 +83,24 @@ func (p *Publisher) Publish(ctx context.Context, body []byte) error {
 			Body:         body,
 		},
 	)
+	if err != nil {
+		return err
+	}
+
+	if p.isPublisherConfirm {
+		// wait for publisher confirm
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case confirm := <-p.ch.NotifyPublish(make(chan amqp.Confirmation, 1)):
+			if !confirm.Ack {
+				return fmt.Errorf("publisher confirm failed, exchangeName: %s, routingKey: %s, deliveryTag: %d",
+					p.Exchange.name, p.Exchange.routingKey, confirm.DeliveryTag)
+			}
+		}
+	}
+
+	return nil
 }
 
 // Close publisher