fix: take a consumer handler slot before receiving

internal/consumer/consumer.go

@@ -205,7 +205,9 @@ recvLoop:
 	return subscriptionErr
 }
 
-// runWithSemaphore limits concurrency via a channel-based semaphore.
+// runWithSemaphore limits concurrency via a channel-based semaphore. The slot
+// is taken before receiving: a received message must not wait for a handler
+// while its visibility timeout runs.
 func (c *consumerImpl) runWithSemaphore(ctx context.Context) error {
 	tracer := otel.GetTracerProvider().Tracer("github.com/hookdeck/outpost/internal/consumer")
 
@@ -215,18 +217,19 @@ func (c *consumerImpl) runWithSemaphore(ctx context.Context) error {
 	sem := make(chan struct{}, c.concurrency)
 recvLoop:
 	for {
-		msg, err := c.receiveWithRetry(ctx, &consecutiveErrors)
-		if err != nil {
-			subscriptionErr = err
-			break recvLoop
-		}
-
 		select {
 		case sem <- struct{}{}:
 		case <-ctx.Done():
 			break recvLoop
 		}
 
+		msg, err := c.receiveWithRetry(ctx, &consecutiveErrors)
+		if err != nil {
+			<-sem
+			subscriptionErr = err
+			break recvLoop
+		}
+
 		go func() {
 			defer func() { <-sem }() // Release the semaphore.
 

internal/consumer/consumer_test.go

@@ -3,7 +3,9 @@ package consumer_test
 import (
 	"context"
 	"log"
+	"sync/atomic"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/hookdeck/outpost/internal/consumer"
@@ -278,3 +280,40 @@ var _ consumer.MessageHandler = &handlerImpl{}
 func (h *handlerImpl) Handle(ctx context.Context, msg *mqs.Message) error {
 	return h.handle(ctx, msg)
 }
+
+// The consumer receives only when a handler slot is free.
+func TestConsumer_ReceivesOnlyWithFreeSlot(t *testing.T) {
+	t.Parallel()
+	synctest.Test(t, func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+
+		var receives atomic.Int32
+		sub := &fakeSubscription{
+			receive: func(ctx context.Context) (*mqs.Message, error) {
+				receives.Add(1)
+				return &mqs.Message{Body: []byte("m")}, nil
+			},
+		}
+		release := make(chan struct{})
+		handler := &handlerImpl{
+			handle: func(ctx context.Context, msg *mqs.Message) error {
+				<-release
+				return nil
+			},
+		}
+
+		done := make(chan error)
+		go func() { done <- consumer.New(sub, handler, consumer.WithConcurrency(2)).Run(ctx) }()
+		synctest.Wait()
+		assert.Equal(t, int32(2), receives.Load(), "one receive per free slot")
+
+		release <- struct{}{}
+		synctest.Wait()
+		assert.Equal(t, int32(3), receives.Load(), "a freed slot allows one more receive")
+
+		cancel()
+		close(release)
+		assert.NoError(t, <-done)
+		assert.Equal(t, int32(3), receives.Load(), "nothing received after shutdown")
+	})
+}