Client gracefully stops

This ensures that when the client wants to stop it will ensure that all
goroutines have stopped.

Author: akarl

client.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -361,7 +362,11 @@ func (c *Client) runOnce() error {
 	notifyInputChClose := inputCh.NotifyClose(make(chan *amqp.Error, 1))
 	notifyOutputChClose := outputCh.NotifyClose(make(chan *amqp.Error, 1))
 
-	err = c.runRepliesConsumer(inputCh)
+	// We will wait on this to ensure that all go routines are done before we
+	// exit this function.
+	wg := sync.WaitGroup{}
+
+	repliesConsumerTag, err := c.runRepliesConsumer(inputCh, &wg)
 	if err != nil {
 		return err
 	}
@@ -376,13 +381,18 @@ func (c *Client) runOnce() error {
 			return err
 		}
 
+		wg.Add(1) // Confirms consumer.
+
 		go c.runConfirmsConsumer(
 			outputCh.NotifyPublish(make(chan amqp.Confirmation)),
 			outputCh.NotifyReturn(make(chan amqp.Return)),
+			&wg,
 		)
 	}
 
-	go c.runPublisher(outputCh)
+	wg.Add(1) // Publisher.
+
+	go c.runPublisher(outputCh, &wg)
 
 	_, err = monitorAndWait(
 		make(chan struct{}),
@@ -393,28 +403,39 @@ func (c *Client) runOnce() error {
 		notifyOutputChClose,
 	)
 	if err != nil {
+		// We don't have a graceful exit, just return the error.
 		return err
 	}
 
+	// 1. Stop the publisher by closing the output channel. This also closes
+	// the confirms consumer if it's running.
+	outputCh.Close()
+
+	// 2. Stop the replies consumer by canceling the consumer.
+	err = inputCh.Cancel(repliesConsumerTag, false)
+	if err != nil {
+		return err
+	}
+
+	// 3. The consumer is stopped, we can now close the input channel.
+	inputCh.Close()
+
+	// 3. Wait for all the go routines to finish.
+	wg.Wait()
+
 	return nil
 }
 
 // runPublisher consumes messages from chan requests and publishes them on the
 // amqp exchange. The method will stop consuming if the underlying amqp channel
 // is closed for any reason, and when this happens the messages will be put back
 // in chan requests unless we have retried to many times.
-func (c *Client) runPublisher(ouputChan *amqp.Channel) {
-	c.logger.Debug("running publisher...")
+func (c *Client) runPublisher(ouputChan *amqp.Channel, wg *sync.WaitGroup) {
+	defer wg.Done()
 
-	// Monitor the closing of this channel. We need to do this in a separate,
-	// goroutine to ensure we won't get a deadlock inside the select below
-	// which can itself close this channel.
-	onClose := make(chan struct{})
+	onClose := ouputChan.NotifyClose(make(chan *amqp.Error, 1))
 
-	go func() {
-		<-ouputChan.NotifyClose(make(chan *amqp.Error))
-		close(onClose)
-	}()
+	c.logger.Debug("running publisher...")
 
 	// Delivery tags always starts at 1 but we increase it before we do any
 	// .Publish() on the channel.
@@ -456,6 +477,9 @@ func (c *Client) runPublisher(ouputChan *amqp.Channel) {
 				request.Publishing,
 			)
 			if err != nil {
+				// Normally a Publish that results in an error will
+				// automatically close the channel and connection. But if the
+				// error occurs during a flush, that doesn't happen.
 				ouputChan.Close()
 
 				c.retryRequest(request, err)
@@ -518,7 +542,9 @@ func (c *Client) retryRequest(request *Request, err error) {
 // runConfirmsConsumer will consume both confirmations and returns and since
 // returns always arrives before confirmations we want to finish handling any
 // return before we handle any confirmations.
-func (c *Client) runConfirmsConsumer(confirms chan amqp.Confirmation, returns chan amqp.Return) {
+func (c *Client) runConfirmsConsumer(confirms chan amqp.Confirmation, returns chan amqp.Return, wg *sync.WaitGroup) {
+	defer wg.Done()
+
 	for {
 		select {
 		case ret, ok := <-returns:
@@ -623,7 +649,7 @@ func (c *Client) respondToRequest(request *Request, d *amqp.Delivery, err error)
 // runRepliesConsumer will declare and start consuming from the queue where we
 // expect replies to come back. The method will stop consuming if the
 // underlying amqp channel is closed for any reason.
-func (c *Client) runRepliesConsumer(inChan *amqp.Channel) error {
+func (c *Client) runRepliesConsumer(inChan *amqp.Channel, wg *sync.WaitGroup) (consumerTag string, err error) {
 	// RabbitMQ will soon no longer support what they call "non-exclusive
 	// transient queues". We want to support reconnects and so we cannot set
 	// the exclusive flag since that would delete the queue on automatically on disconnect.
@@ -637,23 +663,29 @@ func (c *Client) runRepliesConsumer(inChan *amqp.Channel) error {
 		c.replyToQueueDeclareArgs,
 	)
 	if err != nil {
-		return err
+		return "", err
 	}
 
+	tag := uuid.NewString()
+
 	messages, err := inChan.Consume(
 		queue.Name,
-		"",    // consumer tag. Auto-generated by the server.
+		tag,   // consumer tag.
 		true,  // auto-ack. We don't support manual ack for the reply-to queue.
 		true,  // exclusive. We must be the only consumer.
 		false, // no-local.
 		false, // no-wait.
 		c.replyToConsumerArgs,
 	)
 	if err != nil {
-		return err
+		return "", err
 	}
 
+	wg.Add(1)
+
 	go func() {
+		defer wg.Done()
+
 		c.logger.Debug("running replies consumer...")
 
 		for response := range messages {
@@ -680,7 +712,7 @@ func (c *Client) runRepliesConsumer(inChan *amqp.Channel) error {
 		c.logger.Debug("replies consumer is done")
 	}()
 
-	return nil
+	return tag, nil
 }
 
 // Send will send a Request by using a amqp.Publishing.

client_test.go

@@ -118,7 +118,10 @@ func TestClient_ConfirmsConsumer(t *testing.T) {
 	returns := make(chan amqp.Return)
 	confirms := make(chan amqp.Confirmation)
 
-	go client.runConfirmsConsumer(confirms, returns)
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
+	go client.runConfirmsConsumer(confirms, returns, &wg)
 
 	t.Cleanup(func() {
 		close(confirms)
@@ -246,12 +249,11 @@ func TestClient_ConfirmsConsumer(t *testing.T) {
 	t.Run("closing returns will not stop select", func(t *testing.T) {
 		confirms := make(chan amqp.Confirmation)
 		returns := make(chan amqp.Return)
-		finished := make(chan struct{})
 
-		go func() {
-			client.runConfirmsConsumer(confirms, returns)
-			close(finished)
-		}()
+		wg := sync.WaitGroup{}
+		wg.Add(1)
+
+		go client.runConfirmsConsumer(confirms, returns, &wg)
 
 		close(returns)
 
@@ -272,6 +274,13 @@ func TestClient_ConfirmsConsumer(t *testing.T) {
 
 		close(confirms)
 
+		finished := make(chan struct{})
+
+		go func() {
+			wg.Wait()
+			close(finished)
+		}()
+
 		select {
 		case <-finished:
 		case <-time.After(5 * time.Second):
@@ -280,7 +289,35 @@ func TestClient_ConfirmsConsumer(t *testing.T) {
 	})
 }
 
+func TestClientStop(t *testing.T) {
+	t.Parallel()
+
+	client := NewClient(testURL)
+
+	request := NewRequest().
+		WithResponse(false)
+
+	// Ensure that the client has started.
+	_, err := client.Send(request)
+	require.ErrorIs(t, err, ErrRequestReturned)
+
+	stopped := make(chan struct{})
+
+	go func() {
+		client.Stop()
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(5 * time.Second):
+		t.Fatal("did not exit")
+	}
+}
+
 func TestClientStopWhenCannotStart(t *testing.T) {
+	t.Parallel()
+
 	client := NewClient(testURL)
 
 	request := NewRequest().
@@ -290,43 +327,37 @@ func TestClientStopWhenCannotStart(t *testing.T) {
 	_, err := client.Send(request)
 	require.Error(t, err)
 
-	var stopped sync.WaitGroup
-
-	stopped.Add(1)
+	stopped := make(chan struct{})
 
 	go func() {
 		client.Stop()
-		stopped.Done()
+		close(stopped)
 	}()
 
-	assert.Eventually(t, func() bool {
-		stopped.Wait()
-		return true
-	},
-		1*time.Second,
-		500*time.Millisecond,
-	)
+	select {
+	case <-stopped:
+	case <-time.After(5 * time.Second):
+		t.Fatal("did not exit")
+	}
 }
 
 func TestClientStopWhenNeverStarted(t *testing.T) {
-	client := NewClient(testURL)
+	t.Parallel()
 
-	var stopped sync.WaitGroup
+	client := NewClient(testURL)
 
-	stopped.Add(1)
+	stopped := make(chan struct{})
 
 	go func() {
 		client.Stop()
-		stopped.Done()
+		close(stopped)
 	}()
 
-	assert.Eventually(t, func() bool {
-		stopped.Wait()
-		return true
-	},
-		1*time.Second,
-		500*time.Millisecond,
-	)
+	select {
+	case <-stopped:
+	case <-time.After(5 * time.Second):
+		t.Fatal("did not exit")
+	}
 }
 
 func TestClientConfig(t *testing.T) {