From 0561078cbb0c8b50238bb11cf479ed472924ce32 Mon Sep 17 00:00:00 2001
From: Igor Magollo <igraphmagollo@gmail.com>
Date: Fri, 2 Feb 2024 18:23:55 -0300
Subject: [PATCH] feat: add RABBIT_PREFETCH_COUNT to configure channel prefetch
 count option

---
 internal/events/rabbitmq/config.go            | 13 +++++++++++++
 internal/events/rabbitmq/consumer/declares.go |  2 +-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/internal/events/rabbitmq/config.go b/internal/events/rabbitmq/config.go
index 01049b5..2673a61 100644
--- a/internal/events/rabbitmq/config.go
+++ b/internal/events/rabbitmq/config.go
@@ -28,6 +28,7 @@ type Config struct {
 	URL                 string
 	ConsumerName        string
 	ConsumerConcurrency int
+	PrefetchCount       int
 
 	MaxRetries          int
 	InitialInterval     time.Duration
@@ -81,6 +82,18 @@ func LoadConfig(log *zerolog.Logger) Config {
 		log.Info().Msgf("RABBIT_CONSUMER_CONCURRENCY is not set, defaulting to %d", c.ConsumerConcurrency)
 	}
 
+	prefetchCount := os.Getenv("RABBIT_PREFETCH_COUNT")
+	if prefetchCount != "" {
+		parsedPrefetchCount, err := strconv.Atoi(prefetchCount)
+		if err == nil {
+			c.PrefetchCount = parsedPrefetchCount
+		}
+	}
+	if c.PrefetchCount == 0 {
+		c.PrefetchCount = c.ConsumerConcurrency
+		log.Info().Msgf("RABBIT_PREFETCH_COUNT is not set, defaulting to %d", c.PrefetchCount)
+	}
+
 	maxRetries := os.Getenv("RABBIT_MAX_RETRIES")
 	if maxRetries != "" {
 		parsedMaxRetries, err := strconv.Atoi(maxRetries)
diff --git a/internal/events/rabbitmq/consumer/declares.go b/internal/events/rabbitmq/consumer/declares.go
index 7bb7e14..406e81e 100644
--- a/internal/events/rabbitmq/consumer/declares.go
+++ b/internal/events/rabbitmq/consumer/declares.go
@@ -26,7 +26,7 @@ func (r *rabbitmqConsumer) declare(routingKeys []string) error {
 	}
 
 	err = r.chManager.Channel.Qos(
-		r.config.ConsumerConcurrency, 0, false,
+		r.config.PrefetchCount, 0, false,
 	)
 	if err != nil {
 		return eris.Wrap(err, "failed to set QoS")