Add changing visibility timeout for each call

worldcoin · Dec 2, 2024 · 4f661f0 · 4f661f0
1 parent 3e23538
commit 4f661f0
Show file tree

Hide file tree

Showing 4 changed files with 1 addition and 25 deletions.
diff --git a/consumer.go b/consumer.go
@@ -16,7 +16,6 @@ type Config struct {
 	WorkersNum        int
 	VisibilityTimeout int32
 	BatchSize         int32
-	ExtendEnabled     bool
 }
 
 type Consumer struct {
@@ -56,6 +55,7 @@ loop:
 				MaxNumberOfMessages:   c.cfg.BatchSize,
 				WaitTimeSeconds:       int32(5),
 				MessageAttributeNames: []string{"TraceID", "SpanID"},
+				VisibilityTimeout:     c.cfg.VisibilityTimeout,
 			})
 			if err != nil {
 				zap.S().With(zap.Error(err)).Error("could not receive messages from SQS")
@@ -83,9 +83,6 @@ func (c *Consumer) worker(ctx context.Context, messages <-chan *Message) {
 
 func (c *Consumer) handleMsg(ctx context.Context, m *Message) error {
 	if c.handler != nil {
-		if c.cfg.ExtendEnabled {
-			c.extend(ctx, m)
-		}
 		if err := c.handler.Run(ctx, m); err != nil {
 			return m.ErrorResponse(err)
 		}
@@ -104,15 +101,3 @@ func (c *Consumer) delete(ctx context.Context, m *Message) error {
 	zap.S().Debug("message deleted")
 	return nil
 }
-
-func (c *Consumer) extend(ctx context.Context, m *Message) {
-	_, err := c.sqs.ChangeMessageVisibility(ctx, &sqs.ChangeMessageVisibilityInput{
-		QueueUrl:          &c.cfg.QueueURL,
-		ReceiptHandle:     m.ReceiptHandle,
-		VisibilityTimeout: c.cfg.VisibilityTimeout,
-	})
-	if err != nil {
-		zap.S().With(zap.Error(err)).Error("unable to extend message")
-		return
-	}
-}
diff --git a/consumer_test.go b/consumer_test.go
@@ -67,7 +67,6 @@ func TestConsume(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	consumer := NewConsumer(awsCfg, config, msgHandler)
 	go consumer.Consume(ctx)
@@ -109,7 +108,6 @@ func TestConsume_GracefulShutdown(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	msgHandler := MsgHandler{}
 	consumer := NewConsumer(awsCfg, config, &msgHandler)

diff --git a/consumer_with_idle_trigger.go b/consumer_with_idle_trigger.go
@@ -103,9 +103,6 @@ func (c *ConsumerWithIdleTrigger) handleMsg(ctx context.Context, m *Message) err
 				return m.ErrorResponse(err)
 			}
 		} else {
-			if c.cfg.ExtendEnabled {
-				c.extend(ctx, m)
-			}
 			if err := c.handler.Run(ctx, m); err != nil {
 				return m.ErrorResponse(err)
 			}

diff --git a/consumer_with_idle_trigger_test.go b/consumer_with_idle_trigger_test.go
@@ -57,7 +57,6 @@ func TestConsumeWithIdleTrigger(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	consumer := NewConsumerWithIdleTrigger(awsCfg, config, msgHandler, IdleTimeout, SqsReceiveWaitTimeSeconds)
 	go consumer.Consume(ctx)
@@ -99,7 +98,6 @@ func TestConsumeWithIdleTimeout_GracefulShutdown(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	msgHandler := MsgHandlerWithIdleTrigger{
 		t:                 t,
@@ -150,7 +148,6 @@ func TestConsumeWithIdleTimeout_TimesOut(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	msgHandler := MsgHandlerWithIdleTrigger{
 		t:                 t,
@@ -193,7 +190,6 @@ func TestConsumeWithIdleTimeout_TimesOutAndConsumes(t *testing.T) {
 		WorkersNum:        workersNum,
 		VisibilityTimeout: visibilityTimeout,
 		BatchSize:         batchSize,
-		ExtendEnabled:     true,
 	}
 	msgHandler := handlerWithIdleTrigger(t, expectedMsg, expectedMsgAttributes)
 	consumer := NewConsumerWithIdleTrigger(awsCfg, config, msgHandler, IdleTimeout, SqsReceiveWaitTimeSeconds)