Repository: dapr/components-contrib (go)

Issue:
 snssqs: fix consumer starvation
## Actual Behavior
The current implementation uses separate goroutines to process each message retrieved by sqsClient.ReceiveMessageWithContext, which is executed when concurrencyMode is set to parallel. The consumeSubscription function waits for all spawned goroutines to complete before receiving the next batch. However, this design has a critical flaw: it starves the consumer if new messages are published during the processing of an existing batch. These newly published messages won't be processed until all goroutines from the original batch have finished their work, effectively blocking the consumption of the new messages.

## Expected Behavior
Message handling should not be block until the batch is completely processed.

## Steps to Reproduce the Problem

- Set a handler that blocks for a long time.
- Publish one sns message.
- Message will be handled.
- Publish another message
- Message will not be handled until the first finishes.

Pull request:
snssqs: fix consumer starvation
Closes #3479

Patch:
diff --git a/pubsub/aws/snssqs/metadata.go b/pubsub/aws/snssqs/metadata.go
index 4b469106..5ffd575b 100644
--- a/pubsub/aws/snssqs/metadata.go
+++ b/pubsub/aws/snssqs/metadata.go
@@ -57,6 +57,8 @@ type snsSqsMetadata struct {
 	AccountID string `mapstructure:"accountID"`
 	// processing concurrency mode
 	ConcurrencyMode pubsub.ConcurrencyMode `mapstructure:"concurrencyMode"`
+	// limits the number of concurrent goroutines
+	ConcurrencyLimit int `mapstructure:"concurrencyLimit"`
 }
 
 func maskLeft(s string) string {
@@ -130,6 +132,10 @@ func (s *snsSqs) getSnsSqsMetadata(meta pubsub.Metadata) (*snsSqsMetadata, error
 		return nil, err
 	}
 
+	if md.ConcurrencyLimit < 0 {
+		return nil, errors.New("concurrencyLimit must be greater than or equal to 0")
+	}
+
 	s.logger.Debug(md.hideDebugPrintedCredentials())
 
 	return md, nil
diff --git a/pubsub/aws/snssqs/metadata.yaml b/pubsub/aws/snssqs/metadata.yaml
index 641ab4ea..e121eaf0 100644
--- a/pubsub/aws/snssqs/metadata.yaml
+++ b/pubsub/aws/snssqs/metadata.yaml
@@ -128,6 +128,15 @@ metadata:
     default: '"parallel"'
     example: '"single", "parallel"'
     type: string
+  - name: concurrencyLimit
+    required: false
+    description: |
+      Defines the maximum number of concurrent workers handling messages.
+      This value is ignored when "concurrencyMode" is set to “single“.
+      To avoid limiting the number of concurrent workers set this to “0“.
+    type: number
+    default: '0'
+    example: '100'
   - name: accountId
     required: false
     description: |
diff --git a/pubsub/aws/snssqs/snssqs.go b/pubsub/aws/snssqs/snssqs.go
index 4e237176..e84dbeb8 100644
--- a/pubsub/aws/snssqs/snssqs.go
+++ b/pubsub/aws/snssqs/snssqs.go
@@ -595,6 +595,13 @@ func (s *snsSqs) consumeSubscription(ctx context.Context, queueInfo, deadLetters
 		WaitTimeSeconds:     aws.Int64(s.metadata.MessageWaitTimeSeconds),
 	}
 
+	// sem is a semaphore used to control the concurrencyLimit.
+	// It is set only when we are in parallel mode and limit is > 0.
+	var sem chan (struct{}) = nil
+	if (s.metadata.ConcurrencyMode == pubsub.Parallel) && s.metadata.ConcurrencyLimit > 0 {
+		sem = make(chan struct{}, s.metadata.ConcurrencyLimit)
+	}
+
 	for {
 		// If the context is canceled, stop requesting messages
 		if ctx.Err() != nil {
@@ -629,7 +636,6 @@ func (s *snsSqs) consumeSubscription(ctx context.Context, queueInfo, deadLetters
 		}
 		s.logger.Debugf("%v message(s) received on queue %s", len(messageResponse.Messages), queueInfo.arn)
 
-		var wg sync.WaitGroup
 		for _, message := range messageResponse.Messages {
 			if err := s.validateMessage(ctx, message, queueInfo, deadLettersQueueInfo); err != nil {
 				s.logger.Errorf("message is not valid for further processing by the handler. error is: %v", err)
@@ -637,25 +643,30 @@ func (s *snsSqs) consumeSubscription(ctx context.Context, queueInfo, deadLetters
 			}
 
 			f := func(message *sqs.Message) {
-				defer wg.Done()
 				if err := s.callHandler(ctx, message, queueInfo); err != nil {
 					s.logger.Errorf("error while handling received message. error is: %v", err)
 				}
 			}
 
-			wg.Add(1)
 			switch s.metadata.ConcurrencyMode {
 			case pubsub.Single:
 				f(message)
 			case pubsub.Parallel:
-				wg.Add(1)
+				// This is the back pressure mechanism.
+				// It will block until another goroutine frees a slot.
+				if sem != nil {
+					sem <- struct{}{}
+				}
+
 				go func(message *sqs.Message) {
-					defer wg.Done()
+					if sem != nil {
+						defer func() { <-sem }()
+					}
+
 					f(message)
 				}(message)
 			}
 		}
-		wg.Wait()
 	}
 }
 


Test patch:
diff --git a/pubsub/aws/snssqs/snssqs_test.go b/pubsub/aws/snssqs/snssqs_test.go
index f396ddef..82f58bcb 100644
--- a/pubsub/aws/snssqs/snssqs_test.go
+++ b/pubsub/aws/snssqs/snssqs_test.go
@@ -51,6 +51,7 @@ func Test_getSnsSqsMetadata_AllConfiguration(t *testing.T) {
 		"consumerID":               "consumer",
 		"Endpoint":                 "endpoint",
 		"concurrencyMode":          string(pubsub.Single),
+		"concurrencyLimit":         "42",
 		"accessKey":                "a",
 		"secretKey":                "s",
 		"sessionToken":             "t",
@@ -68,6 +69,7 @@ func Test_getSnsSqsMetadata_AllConfiguration(t *testing.T) {
 	r.Equal("consumer", md.SqsQueueName)
 	r.Equal("endpoint", md.Endpoint)
 	r.Equal(pubsub.Single, md.ConcurrencyMode)
+	r.Equal(42, md.ConcurrencyLimit)
 	r.Equal("a", md.AccessKey)
 	r.Equal("s", md.SecretKey)
 	r.Equal("t", md.SessionToken)
@@ -105,6 +107,7 @@ func Test_getSnsSqsMetadata_defaults(t *testing.T) {
 	r.Equal("", md.SessionToken)
 	r.Equal("r", md.Region)
 	r.Equal(pubsub.Parallel, md.ConcurrencyMode)
+	r.Equal(0, md.ConcurrencyLimit)
 	r.Equal(int64(10), md.MessageVisibilityTimeout)
 	r.Equal(int64(10), md.MessageRetryLimit)
 	r.Equal(int64(2), md.MessageWaitTimeSeconds)
@@ -273,6 +276,20 @@ func Test_getSnsSqsMetadata_invalidMetadataSetup(t *testing.T) {
 			}}},
 			name: "invalid message concurrencyMode",
 		},
+		// invalid concurrencyLimit
+		{
+			metadata: pubsub.Metadata{Base: metadata.Base{Properties: map[string]string{
+				"consumerID":        "consumer",
+				"Endpoint":          "endpoint",
+				"AccessKey":         "acctId",
+				"SecretKey":         "secret",
+				"awsToken":          "token",
+				"Region":            "region",
+				"messageRetryLimit": "10",
+				"concurrencyLimit":  "-1",
+			}}},
+			name: "invalid message concurrencyLimit",
+		},
 	}
 
 	l := logger.NewLogger("SnsSqs unit test")