feat: move message polling to the client and improve error handling (#24)

Author: dzmitryharchanin-tomtom

cmd/hook/hook.go

@@ -166,7 +166,12 @@ func checkStatus(ghaRunnerName string) bool {
 }
 
 func sendMessageToSqs(message string) {
-	client, err := sqs.NewClient(cfg.AwsRegion, cfg.AwsRunnerSQSUrl)
+	client, err := sqs.NewClient(
+		&sqs.SQSConfig{
+			QueueURL:  cfg.AwsRunnerSQSUrl,
+			AWSRegion: cfg.AwsRegion,
+		},
+	)
 	if err != nil {
 		log.Fatalf("Failed to create sqs client: %v", err)
 	}

cmd/tarter/app/server.go

@@ -23,6 +23,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/caarlos0/env"
 	"github.com/gorilla/mux"
 	"github.com/spf13/cobra"
 	coreApi "github.com/tomtom-international/macos-actions-runner-controller/pkg/core/api"
@@ -48,10 +49,12 @@ func NewTarterCommand() *cobra.Command {
 			if cmd.Name() == "version" {
 				return nil
 			}
-			err := config.LoadTarterConfiguration(&tarterConfig)
-			if err != nil {
-				return errors.New("error: failed to load Tarter configuration. " + err.Error())
+
+			// Load configuration from environment variables
+			if err := env.Parse(&tarterConfig); err != nil {
+				return errors.New("error: failed to load Tarter configuration from environment. " + err.Error())
 			}
+
 			if configFile == "" {
 				return errors.New("error: --config argument is required")
 			}

pkg/clients/sqs/client.go

@@ -18,6 +18,8 @@ package sqs
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"math/rand"
 	"time"
 
@@ -27,29 +29,75 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
 )
 
+const (
+	defaultMinBackoff = 1 * time.Second
+	defaultMaxBackoff = 60 * time.Second
+)
+
 type SQSClient struct {
-	service  *sqs.Client
-	QueueURL string
+	service        *sqs.Client
+	config         *SQSConfig
+	currentBackoff time.Duration
+	errHandle      func(error)
 }
 
-func NewClient(region, queueURL string) (*SQSClient, error) {
-	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(region))
+type SQSConfig struct {
+	QueueURL   string
+	AWSRegion  string
+	MaxBackoff time.Duration
+	MinBackoff time.Duration
+	ErrHandle  func(error)
+}
+
+func NewClient(sqsCfg *SQSConfig) (*SQSClient, error) {
+	err := validateSQSConfig(sqsCfg)
+	if err != nil {
+		return nil, err
+	}
+	cfg, err := config.LoadDefaultConfig(context.TODO(),
+		config.WithRegion(sqsCfg.AWSRegion),
+		config.WithRetryMaxAttempts(3),
+		config.WithRetryMode(aws.RetryModeStandard),
+	)
 	if err != nil {
 		return nil, err
 	}
 
 	svc := sqs.NewFromConfig(cfg)
 
 	return &SQSClient{
-		service:  svc,
-		QueueURL: queueURL,
+		service:        svc,
+		config:         sqsCfg,
+		currentBackoff: sqsCfg.MinBackoff,
+		errHandle:      sqsCfg.ErrHandle,
 	}, nil
 }
 
+func validateSQSConfig(cfg *SQSConfig) error {
+	if cfg.QueueURL == "" {
+		return errors.New("queueURL is required")
+	}
+	if cfg.AWSRegion == "" {
+		return errors.New("region is required")
+	}
+	if cfg.MaxBackoff <= 0 {
+		cfg.MaxBackoff = defaultMaxBackoff
+	}
+	if cfg.MinBackoff <= 0 {
+		cfg.MinBackoff = defaultMinBackoff
+	}
+	if cfg.ErrHandle == nil {
+		cfg.ErrHandle = func(err error) {
+			return
+		}
+	}
+	return nil
+}
+
 func (c *SQSClient) SendMessage(ctx context.Context, messageBody string) (*sqs.SendMessageOutput, error) {
 	result, err := c.service.SendMessage(ctx, &sqs.SendMessageInput{
 		MessageBody: aws.String(messageBody),
-		QueueUrl:    aws.String(c.QueueURL),
+		QueueUrl:    aws.String(c.config.QueueURL),
 	})
 
 	if err != nil {
@@ -61,7 +109,7 @@ func (c *SQSClient) SendMessage(ctx context.Context, messageBody string) (*sqs.S
 
 func (c *SQSClient) DeleteMessage(ctx context.Context, receiptHandle string) (*sqs.DeleteMessageOutput, error) {
 	result, err := c.service.DeleteMessage(ctx, &sqs.DeleteMessageInput{
-		QueueUrl:      aws.String(c.QueueURL),
+		QueueUrl:      aws.String(c.config.QueueURL),
 		ReceiptHandle: aws.String(receiptHandle),
 	})
 
@@ -74,7 +122,7 @@ func (c *SQSClient) DeleteMessage(ctx context.Context, receiptHandle string) (*s
 
 func (c *SQSClient) ReceiveMessages(ctx context.Context, maxMessages int32, waitTimeSeconds int32) ([]types.Message, error) {
 	result, err := c.service.ReceiveMessage(ctx, &sqs.ReceiveMessageInput{
-		QueueUrl:            aws.String(c.QueueURL),
+		QueueUrl:            aws.String(c.config.QueueURL),
 		MaxNumberOfMessages: maxMessages,
 		WaitTimeSeconds:     waitTimeSeconds,
 	})
@@ -89,10 +137,6 @@ func (c *SQSClient) ReceiveMessages(ctx context.Context, maxMessages int32, wait
 // CalculateNextBackoff computes the next backoff duration using an exponential
 // strategy with jitter to prevent synchronized polling.
 //
-// Parameters:
-//   - current: The current backoff duration.
-//   - max: The maximum allowable backoff duration.
-//
 // Returns:
 //
 //	A time.Duration representing the next backoff period.
@@ -101,18 +145,91 @@ func (c *SQSClient) ReceiveMessages(ctx context.Context, maxMessages int32, wait
 // (subtracting up to 25% of the doubled value) to prevent multiple instances
 // from synchronizing their polling cycles. The result is capped at the specified
 // maximum duration.
-func (c *SQSClient) CalculateNextBackoff(current, maxDuration time.Duration) time.Duration {
+func (c *SQSClient) calculateNextBackoff() time.Duration {
 	// Double the current backoff
-	next := current * 2
+	next := c.currentBackoff * 2
 
 	// Apply jitter (randomness) to prevent synchronized polling
 	jitter := time.Duration(rand.Int63n(int64(next / 4)))
 	next -= jitter
 
 	// Ensure we don't exceed the maximum
-	if next > maxDuration {
-		return maxDuration
+	if next > c.config.MaxBackoff {
+		c.currentBackoff = c.config.MaxBackoff
+	} else {
+		c.currentBackoff = next
 	}
 
-	return next
+	return c.currentBackoff
+}
+
+func (c *SQSClient) resetBackoff() {
+	c.currentBackoff = c.config.MinBackoff
+}
+
+func (c *SQSClient) PollForMessages(ctx context.Context, handler func(message *types.Message) error) {
+	errChan := make(chan error, 100)
+	pollCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	go func() {
+		defer close(errChan)
+		for {
+			select {
+			case <-pollCtx.Done():
+				return
+			case err := <-errChan:
+				if err != nil && c.errHandle != nil {
+					c.errHandle(err)
+				}
+			}
+		}
+	}()
+
+	var backoff time.Duration
+
+	for {
+		select {
+		case <-ctx.Done():
+			// Context was canceled - exit gracefully without error
+			return
+		default:
+			messages, err := c.ReceiveMessages(ctx, 10, 20)
+			if err != nil {
+				// Check if the error is due to context cancellation
+				if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+					// Context canceled while receiving messages, stopping polling
+					return
+				}
+				errChan <- fmt.Errorf("failed to receive messages: %w", err)
+
+				// Increase backoff on errors
+				backoff = c.calculateNextBackoff()
+				time.Sleep(backoff)
+				continue
+			}
+
+			if len(messages) == 0 {
+				// No messages found, increase the backoff
+				backoff = c.calculateNextBackoff()
+				time.Sleep(backoff)
+			} else {
+				// Messages found, reset backoff
+				c.resetBackoff()
+
+				for _, msg := range messages {
+					err = handler(&msg)
+					if err != nil {
+						continue
+					}
+
+					_, err = c.DeleteMessage(ctx, *msg.ReceiptHandle)
+					if err != nil {
+						errChan <- fmt.Errorf("failed to delete message %s from queue %s: %w", *msg.MessageId,
+							c.config.QueueURL, err)
+						continue
+					}
+				}
+			}
+		}
+	}
 }

pkg/controller/controller.go

@@ -19,12 +19,12 @@ package controller
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"net"
 	"sync"
 	"time"
 
+	sqsTypes "github.com/aws/aws-sdk-go-v2/service/sqs/types"
 	"github.com/gorilla/websocket"
 	"github.com/tomtom-international/macos-actions-runner-controller/pkg/clients/etcd"
 	"github.com/tomtom-international/macos-actions-runner-controller/pkg/clients/sqs"
@@ -37,11 +37,6 @@ import (
 	"github.com/tomtom-international/macos-actions-runner-controller/pkg/utils"
 )
 
-const (
-	minBackoff = 4 * time.Second
-	maxBackoff = 60 * time.Second
-)
-
 type Controller struct {
 	EtcdClient             *etcd.EtcdClient
 	nodePoolManager        *np.Manager
@@ -58,7 +53,15 @@ func NewController(configuration config.ControllerConfig) (*Controller, error) {
 	if err != nil {
 		return nil, err
 	}
-	sqsClient, err := sqs.NewClient(configuration.AwsRegion, configuration.AwsRunnerRequestSQSUrl)
+	sqsClient, err := sqs.NewClient(
+		&sqs.SQSConfig{
+			QueueURL:  configuration.AwsRunnerRequestSQSUrl,
+			AWSRegion: configuration.AwsRegion,
+			ErrHandle: func(err error) {
+				logger.Errorf("Sqs error: %s", err.Error())
+			},
+		},
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -87,7 +90,7 @@ func (c *Controller) Run(ctx context.Context, wg *sync.WaitGroup) {
 	wg.Add(2)
 	go func() {
 		defer wg.Done()
-		c.listenForNewRunners(ctx, c.createRunner)
+		c.listenForNewRunners(ctx, c.readSQSMessages)
 	}()
 	go func() {
 		defer wg.Done()
@@ -345,9 +348,9 @@ func (c *Controller) GetRunnerList() ([]types.Runner, error) {
 	return runners, nil
 }
 
-func (c *Controller) listenForNewRunners(ctx context.Context, handlerFunc func(runner *types.Runner) error) {
-	logger.Debugf("Listening for new messages from SQS queue: %s", c.sqsRunnerRequestClient.QueueURL)
-	pollForMessages(ctx, c.sqsRunnerRequestClient, handlerFunc)
+func (c *Controller) listenForNewRunners(ctx context.Context, handlerFunc func(message *sqsTypes.Message) error) {
+	logger.Debugf("Listening for new messages from SQS queue")
+	c.sqsRunnerRequestClient.PollForMessages(ctx, handlerFunc)
 }
 
 func (c *Controller) checkRegisteredNodes() {
@@ -373,65 +376,21 @@ func (c *Controller) checkRegisteredNodes() {
 	}
 }
 
-func pollForMessages(ctx context.Context, sqsClient *sqs.SQSClient, handler func(runner *types.Runner) error) {
-	currentBackoff := minBackoff
-
-	for {
-		select {
-		case <-ctx.Done():
-			// Context was canceled - exit gracefully without error
-			logger.Debugf("Context canceled, stopping message polling from %s", sqsClient.QueueURL)
-			return
-		default:
-			messages, err := sqsClient.ReceiveMessages(ctx, 10, 20)
-			if err != nil {
-				// Check if the error is due to context cancellation
-				if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-					logger.Debugf("Context canceled while receiving messages, stopping polling")
-					return
-				}
-
-				logger.Errorf("Failed to receive messages: %s", err.Error())
-
-				// Increase backoff on errors
-				currentBackoff = sqsClient.CalculateNextBackoff(currentBackoff, maxBackoff)
-				time.Sleep(currentBackoff)
-				continue
-			}
+func (c *Controller) readSQSMessages(msg *sqsTypes.Message) error {
+	runner := types.Runner{}
+	logger.Debugf("Received message with id: %s from queue", *msg.MessageId)
+	err := json.Unmarshal([]byte(*msg.Body), &runner)
+	if err != nil {
+		logger.Errorf("Failed to unmarshal message %s from queue: %s", *msg.MessageId, err.Error())
+		return err
+	}
+	runner.Condition.CreateRequestID = *msg.MessageId
 
-			if len(messages) == 0 {
-				// No messages found, increase the backoff
-				currentBackoff = sqsClient.CalculateNextBackoff(currentBackoff, maxBackoff)
-				time.Sleep(currentBackoff)
-			} else {
-				// Messages found, reset backoff
-				currentBackoff = minBackoff
-
-				for _, msg := range messages {
-					runner := types.Runner{}
-					logger.Debugf("Received message with id: %s from queue %s", *msg.MessageId, sqsClient.QueueURL)
-					err := json.Unmarshal([]byte(*msg.Body), &runner)
-					if err != nil {
-						// message will be sent to DLQ after maximum receives
-						logger.Errorf("Failed to unmarshal message %s from queue %s: %s", *msg.MessageId, sqsClient.QueueURL, err.Error())
-						continue
-					}
-					runner.Condition.CreateRequestID = *msg.MessageId
-
-					err = handler(&runner)
-					if err != nil {
-						continue
-					}
-
-					_, err = sqsClient.DeleteMessage(ctx, *msg.ReceiptHandle)
-					if err != nil {
-						logger.Errorf("Failed to delete message %s from queue %s: %s", *msg.MessageId, sqsClient.QueueURL, err.Error())
-						continue
-					}
-				}
-			}
-		}
+	err = c.createRunner(&runner)
+	if err != nil {
+		return err
 	}
+	return nil
 }
 
 func (c *Controller) AddRunnersWatcher(conn *websocket.Conn, notificationChan chan types.WatcherRunnersUpdate) {