Repository: google/go-cloud (Go)

Issue:
pubsub/all: Using long polling when consuming from a FIFO SQS queue ends up blocking requests
### Describe the bug

The `pubsub` package has trouble consuming from a **FIFO enabled SQS** queue when using **long polling** on messages with the same `MessageGroupID` and blocks requests by the amount of wait time configured for **long polling**.

The reason this is happening is due to the way the `pubsub` library is structured and also due to the basic operation mode of **FIFO** SQS queues.


#### Details about SQS FIFO queue operation
When an SQS queue has **FIFO** enabled, then as the name suggests, it guarantees that the messages belonging to the same **group** (using the same `MessageGroupID`), will be delivered in-order and with an exactly-once delivery guarantee. This means, that unless a message is acknowledged, the next in-line will not be made available when reading from the queue. Sending a **poll** request when the last message has not been acknowledged, will return 0 messages from the queue. More info [here](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-fifo-queues.html)

**Long Polling** is a mode during which you configure a **Wait Time**, which shows how much time a poll request will wait before returning until it gets a message or the configured wait time expires. This is used to reduce the number of repeated requests that might return 0 messages. For example, if you configure a wait time of 10 seconds and then send a **poll** request, it will wait up to 10 seconds, either returning as soon as a message is available or when the 10 seconds have passed by (returning 0 messages again). If **short polling** is used and a poll request is sent when 0 messages, are available the request will return immediately, sending an other request right after, repeating until a message is available. It is generally advised to use long polling. More info can be found [here](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-short-and-long-polling.html)

#### `pubsub` operation
The package supports an internal auto-scaling mechanism that tries to increase the message consumption throughput based on the current speed of message processing (time from the message received, until it is acknowledged).

When the message rate is increased, the so-called **batch size** of messages requested is increased and requests are split into **batches** depending on the allowed limits of each driver/cloud-provider.

For example, SQS supports up to 10 messages returned in a single request. So, if the driver determines that we can process more than 10 messages it will increase the batch size to eg: 12. This will have to be split into two batches, one requesting 10 (the max SQS limit) and one requesting the remaining 2. These requests will be sent at the same time from the function [getNextBatch](https://github.com/google/go-cloud/blob/60313ac53f50306c322ab91d8bc48b07ac3c29cc/pubsub/pubsub.go#L641) and once resolved they will be passed over to the execution flow.

Now, this here is the problematic part:
- Execution flow won't continue until both requests have returned their results
- Messages are processed/acked only after the above flow is completed
- SQS Queue will not provide the next group item unless the previous one is acknowledged

This leads to the following results:
1. 1st request for 10 messages will be served immediately returning up to 10 messages from the queue
2. the 2nd request of the remaining 2 will be sent at the same time
3. However, since the 10 messages from the 1st request have not been acknowledged (message acknowledgment, is after we exit this flow), the second request will not receive any messages
4. With long polling enabled, the second request will not return immediately, but wait up to the configured **WaitTime** until it eventually returns 0, because the first batch was not acked, blocking the whole flow.
5. After that the flow will continue normally (only with the messages from the 1st request) and repeat the cycle continuously

### To Reproduce
We have to create an **SQS** queue with the following:
- SQS Queue has **FIFO** enabled
- We use **Long Polling** in our connection (`awssnssqs.SubscriptionOptions.WaitTime` option)
- There are a lot of messages in the **same group (using the same `MessageGroupID`)**
- We consume fast enough so that auto-scaling will increase the batch size enough to split requests into more than one batch. (We can also emulate this by setting a lower `MaxBatchSize` limit on our connection).

Steps to reproduce the behavior
1. Publish quite some messages to the **FIFO SQS** queue (to have room for consuming enough) **using the same `MessageGroupId`**
2. Open a connection to the queue using long polling:
```go
subURL := <OUR_QUEUE's URL>

urlMux := &pubsub.URLMux{}

u := &awssnssqs.URLOpener{
  UseV2: true,
  SubscriptionOptions: awssnssqs.SubscriptionOptions{
    WaitTime: 10 * time.Second, // enable long-polling
    ReceiveBatcherOptions: batcher.Options{
      MaxBatchSize: 2, // we reduce batch size to speed up bug appearance
    },
  },
}
urlMux.RegisterSubscription(awssnssqs.SQSScheme, u)
// Open a *pubsub.Subscription using the URL.
sub, err := urlMux.OpenSubscription(ctx, subURL)
```
3. Start consuming the messages using some logging to be able to see the added delay:
```go
for {
  now := time.Now()
  fmt.Println("Start message reception", now)
  m, err := sub.Receive(ctx)
  fmt.Println("End of message reception", time.Since(now))
  if err != nil {
    log.Print(err)
    return err
  }
  fmt.Printf("%s\n", m.Body)
  m.Ack()
}
```
Almost immediately, you will start noticing that you will be getting 10 seconds delay (the configured `WaitTime`) between start and end of the Message Receival (due to the 2nd batch request being blocked)

## Expected behavior

I would expect that the batches are processed independently and the flow not being blocked. Since, `Subscription.Receive` returns a single message, it could work without getting blocked by independent requests

### Version

**v0.38.0**

### Additional context

**The problematic code lies in the `pubsub` basic package (that's why I used `all` in the issue name) and not in the `awssnssqs` one, but because the bug is based on AWS operation, I am not sure if it is entirely transferrable to other cloud providers.**



Discussion:

Patch:
diff --git a/pubsub/pubsub.go b/pubsub/pubsub.go
index f68b17e03e..077637fb4e 100644
--- a/pubsub/pubsub.go
+++ b/pubsub/pubsub.go
@@ -565,17 +565,28 @@ func (s *Subscription) Receive(ctx context.Context) (_ *Message, err error) {
 				if s.preReceiveBatchHook != nil {
 					s.preReceiveBatchHook(batchSize)
 				}
-				msgs, err := s.getNextBatch(batchSize)
-				s.mu.Lock()
-				defer s.mu.Unlock()
-				if err != nil {
-					// Non-retryable error from ReceiveBatch -> permanent error.
-					s.err = err
-				} else if len(msgs) > 0 {
-					s.q = append(s.q, msgs...)
+				resultChannel := s.getNextBatch(batchSize)
+				for msgsOrError := range resultChannel {
+					if msgsOrError.msgs != nil && len(msgsOrError.msgs) > 0 {
+						// messages received from channel
+						s.mu.Lock()
+						s.q = append(s.q, msgsOrError.msgs...)
+						s.mu.Unlock()
+						// notify that queue should now have messages
+						s.waitc <- struct{}{}
+					} else if msgsOrError.err != nil {
+						// err can receive message only after batch group completes
+						// Non-retryable error from ReceiveBatch -> permanent error
+						s.mu.Lock()
+						s.err = msgsOrError.err
+						s.mu.Unlock()
+					}
 				}
+				// batch reception finished
+				s.mu.Lock()
 				close(s.waitc)
 				s.waitc = nil
+				s.mu.Unlock()
 			}()
 		}
 		if len(s.q) > 0 {
@@ -625,11 +636,11 @@ func (s *Subscription) Receive(ctx context.Context) (_ *Message, err error) {
 		}
 		// A call to ReceiveBatch must be in flight. Wait for it.
 		waitc := s.waitc
-		s.mu.Unlock()
+		s.mu.Unlock() // unlock to allow message or error processing from background goroutine
 		select {
 		case <-waitc:
-			s.mu.Lock()
 			// Continue to top of loop.
+			s.mu.Lock()
 		case <-ctx.Done():
 			s.mu.Lock()
 			return nil, ctx.Err()
@@ -637,16 +648,19 @@ func (s *Subscription) Receive(ctx context.Context) (_ *Message, err error) {
 	}
 }
 
-// getNextBatch gets the next batch of messages from the server and returns it.
-func (s *Subscription) getNextBatch(nMessages int) ([]*driver.Message, error) {
-	var mu sync.Mutex
-	var q []*driver.Message
+type msgsOrError struct {
+	msgs []*driver.Message
+	err  error
+}
 
+// getNextBatch gets the next batch of messages from the server. It will return a channel that will itself return the
+// messages as they come from each independent batch, or an operation error
+func (s *Subscription) getNextBatch(nMessages int) chan msgsOrError {
 	// Split nMessages into batches based on recvBatchOpts; we'll make a
 	// separate ReceiveBatch call for each batch, and aggregate the results in
 	// msgs.
 	batches := batcher.Split(nMessages, s.recvBatchOpts)
-
+	result := make(chan msgsOrError, len(batches))
 	g, ctx := errgroup.WithContext(s.backgroundCtx)
 	for _, maxMessagesInBatch := range batches {
 		// Make a copy of the loop variable since it will be used by a goroutine.
@@ -663,16 +677,18 @@ func (s *Subscription) getNextBatch(nMessages int) ([]*driver.Message, error) {
 			if err != nil {
 				return wrapError(s.driver, err)
 			}
-			mu.Lock()
-			defer mu.Unlock()
-			q = append(q, msgs...)
+			result <- msgsOrError{msgs: msgs}
 			return nil
 		})
 	}
-	if err := g.Wait(); err != nil {
-		return nil, err
-	}
-	return q, nil
+	go func() {
+		// wait on group completion on the background and proper channel closing
+		if err := g.Wait(); err != nil {
+			result <- msgsOrError{err: err}
+		}
+		close(result)
+	}()
+	return result
 }
 
 var errSubscriptionShutdown = gcerr.Newf(gcerr.FailedPrecondition, nil, "pubsub: Subscription has been Shutdown")


Test patch:
diff --git a/pubsub/pubsub_test.go b/pubsub/pubsub_test.go
index 3c42442c68..ddce3147b9 100644
--- a/pubsub/pubsub_test.go
+++ b/pubsub/pubsub_test.go
@@ -20,6 +20,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -281,6 +282,43 @@ func TestCancelTwoReceives(t *testing.T) {
 	}
 }
 
+type secondReceiveBlockedDriverSub struct {
+	driver.Subscription
+	receiveCounter atomic.Uint64
+}
+
+func (s *secondReceiveBlockedDriverSub) ReceiveBatch(ctx context.Context, _ int) ([]*driver.Message, error) {
+	s.receiveCounter.Add(1)
+	if s.receiveCounter.Load() > 1 {
+		// wait after 1st request for the context to finish before returning the batch result
+		<-ctx.Done()
+	}
+	msg := &driver.Message{Body: []byte(fmt.Sprintf("message #%d", s.receiveCounter.Load()))}
+	return []*driver.Message{msg}, nil
+}
+func (*secondReceiveBlockedDriverSub) CanNack() bool          { return false }
+func (*secondReceiveBlockedDriverSub) IsRetryable(error) bool { return false }
+func (*secondReceiveBlockedDriverSub) Close() error           { return nil }
+
+func TestIndependentBatchReturn(t *testing.T) {
+	// We want to test the scenario when multiple batch requests are sent, as long as one of them succeeds, it should
+	// not block the Subscription.Receive result
+	s := NewSubscription(
+		&secondReceiveBlockedDriverSub{},
+		&batcher.Options{MaxBatchSize: 1, MaxHandlers: 2}, // force 2 batches, by allowing 2 handlers and 1 msg per batch
+		nil,
+	)
+	// set the false calculated subscription batch size to force 2 batches to be called
+	s.runningBatchSize = 2
+	ctx := context.Background()
+	defer s.Shutdown(ctx)
+	_, err := s.Receive(ctx)
+	if err != nil {
+		t.Fatal("Receive should not fail", err)
+		return
+	}
+}
+
 func TestRetryTopic(t *testing.T) {
 	// Test that Send is retried if the driver returns a retryable error.
 	ctx := context.Background()