fix(tracing): add amqp prefetch

sensorbucket · Nov 25, 2024 · dcbc8eb · dcbc8eb
1 parent e51ccb9
commit dcbc8eb
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/services/tracing/ingress-archiver/service/transport.go b/services/tracing/ingress-archiver/service/transport.go
@@ -14,13 +14,15 @@ import (
 	"sensorbucket.nl/sensorbucket/pkg/mq"
 )
 
-func StartIngressDTOConsumer(conn *mq.AMQPConnection, svc *Application, queue, xchg, topic string) {
+func StartIngressDTOConsumer(conn *mq.AMQPConnection, svc *Application, queue, xchg, topic string, prefetch int) {
 	consume := conn.Consume(queue, func(c *amqp091.Channel) error {
+		if err := c.Qos(prefetch, 0, false); err != nil {
+			return fmt.Errorf("error setting Qos with prefetch on amqp: %w", err)
+		}
 		_, err := c.QueueDeclare(queue, true, false, false, false, nil)
 		if err != nil {
 			return err
 		}
-
 		// Create exchange and bind if both arguments are provided, this is optional
 		if xchg != "" && topic != "" {
 			if err := c.ExchangeDeclare(xchg, "topic", true, false, false, false, nil); err != nil {

diff --git a/services/tracing/main.go b/services/tracing/main.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"strconv"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -40,6 +41,7 @@ var (
 	AMQP_QUEUE_INGRESS               = env.Could("AMQP_QUEUE_INGRESS", "archive-ingress")
 	AMQP_XCHG_INGRESS                = env.Could("AMQP_XCHG_INGRESS", "ingress")
 	AMQP_XCHG_INGRESS_TOPIC          = env.Could("AMQP_XCHG_INGRESS_TOPIC", "ingress.*")
+	AMQP_PREFETCH                    = env.Could("AMQP_PREFETCH", "5")
 	AUTH_JWKS_URL                    = env.Could("AUTH_JWKS_URL", "http://oathkeeper:4456/.well-known/jwks.json")
 )
 
@@ -61,6 +63,11 @@ func Run(cleanup cleanupper.Cleanupper) error {
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer cancel()
 
+	prefetch, err := strconv.Atoi(AMQP_PREFETCH)
+	if err != nil {
+		return err
+	}
+
 	stopProfiler, err := web.RunProfiler()
 	if err != nil {
 		log.Printf("could not setup profiler server: %s\n", err)
@@ -93,6 +100,7 @@ func Run(cleanup cleanupper.Cleanupper) error {
 		go ingressarchiver.StartIngressDTOConsumer(
 			mqConn, svc,
 			AMQP_QUEUE_INGRESS, AMQP_XCHG_INGRESS, AMQP_XCHG_INGRESS_TOPIC,
+			prefetch,
 		)
 		ingressarchiver.CreateHTTPTransport(r, svc)
 	}
@@ -107,6 +115,7 @@ func Run(cleanup cleanupper.Cleanupper) error {
 			AMQP_QUEUE_PIPELINEMESSAGES,
 			AMQP_XCHG_PIPELINEMESSAGES,
 			AMQP_XCHG_PIPELINEMESSAGES_TOPIC,
+			prefetch,
 		)
 		tracinghttp := tracingtransport.NewHTTP(tracingService, HTTP_BASE)
 		tracinghttp.SetupRoutes(r)

diff --git a/services/tracing/tracing/transport/mq.go b/services/tracing/tracing/transport/mq.go
@@ -13,8 +13,8 @@ import (
 	"sensorbucket.nl/sensorbucket/services/tracing/tracing"
 )
 
-func StartMQ(svc *tracing.Service, conn *mq.AMQPConnection, queue, xchg, topic string) {
-	pipelineMessages := mq.Consume(conn, queue, setupFunc(queue, xchg, topic))
+func StartMQ(svc *tracing.Service, conn *mq.AMQPConnection, queue, xchg, topic string, prefetch int) {
+	pipelineMessages := mq.Consume(conn, queue, setupFunc(prefetch, queue, xchg, topic))
 
 	log.Println("Measurement MQ Transport running")
 	go processMessage(pipelineMessages, svc)
@@ -89,8 +89,11 @@ func processMessage(deliveries <-chan amqp091.Delivery, svc *tracing.Service) {
 	}
 }
 
-func setupFunc(queue, xchg, topic string) mq.AMQPSetupFunc {
+func setupFunc(prefetch int, queue, xchg, topic string) mq.AMQPSetupFunc {
 	return func(c *amqp091.Channel) error {
+		if err := c.Qos(prefetch, 0, false); err != nil {
+			return fmt.Errorf("error setting Qos with prefetch on amqp: %w", err)
+		}
 		_, err := c.QueueDeclare(queue, true, false, false, false, nil)
 		if err != nil {
 			return fmt.Errorf("error declaring amqp queue: %w", err)