Replace SQS singleton with per-pipeline client cache with TTL eviction

cmd/main.go

@@ -21,6 +21,7 @@ import (
 	"flag"
 	"os"
 	"path/filepath"
+	"time"
 
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
 	// to ensure that exec-entrypoint and run can make use of them.
@@ -40,6 +41,7 @@ import (
 
 	operatorv1alpha1 "github.com/redhat-data-and-ai/unstructured-data-controller/api/v1alpha1"
 	"github.com/redhat-data-and-ai/unstructured-data-controller/internal/controller"
+	"github.com/redhat-data-and-ai/unstructured-data-controller/pkg/awsclienthandler"
 	// +kubebuilder:scaffold:imports
 )
 
@@ -289,8 +291,13 @@ func main() {
 		os.Exit(1)
 	}
 
+	ctx := ctrl.SetupSignalHandler()
+
+	setupLog.Info("initializing SQS client cache")
+	awsclienthandler.InitSQSCache(ctx, 30*time.Minute)
+
 	setupLog.Info("starting manager")
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "problem running manager")
 		os.Exit(1)
 	}

internal/controller/sourcecrawler_controller.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -77,6 +78,10 @@ func (r *SourceCrawlerReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 
 	sourceCrawlerCR := &operatorv1alpha1.SourceCrawler{}
 	if err := r.Get(ctx, req.NamespacedName, sourceCrawlerCR); err != nil {
+		if apierrors.IsNotFound(err) {
+			awsclienthandler.DeleteSQSClient(req.String())
+			return ctrl.Result{}, nil
+		}
 		logger.Error(err, "failed to get SourceCrawler CR")
 		return ctrl.Result{}, err
 	}
@@ -161,24 +166,24 @@ func (r *SourceCrawlerReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 			if err != nil {
 				return ctrl.Result{}, r.handleError(ctx, sourceCrawlerCR, fmt.Errorf("failed to get source credentials for SQS: %w", err))
 			}
-			if _, err := awsclienthandler.NewSQSClientFromConfig(ctx, sourceAWSConfig); err != nil {
+			namespacedName := types.NamespacedName{Namespace: sourceCrawlerCR.Namespace, Name: sourceCrawlerCR.Name}.String()
+			if _, err := awsclienthandler.NewSQSClientFromConfig(ctx, sourceAWSConfig, namespacedName); err != nil {
 				return ctrl.Result{}, r.handleError(ctx, sourceCrawlerCR, fmt.Errorf("failed to create SQS client: %w", err))
 			}
-			return handleSQSWakeUp(ctx, sqsQueueURL, sourceCrawlerConfig.S3Config.Bucket, sourceCrawlerConfig.S3Config.Prefix), nil
+			return handleSQSWakeUp(ctx, sqsQueueURL, sourceCrawlerConfig.S3Config.Bucket, sourceCrawlerConfig.S3Config.Prefix, namespacedName), nil
 		}
 	}
 	return ctrl.Result{RequeueAfter: defaultCrawlerResyncInterval}, nil
 }
 
-func handleSQSWakeUp(ctx context.Context, queueURL, bucket, prefix string) ctrl.Result {
+func handleSQSWakeUp(ctx context.Context, queueURL, bucket, prefix string, pipelineName string) ctrl.Result {
 	logger := log.FromContext(ctx)
 
-	sqsClient, err := awsclienthandler.GetSQSClient()
-	if err != nil {
-		logger.Error(err, "failed to initialize SQS client")
+	sqsClient, ok := awsclienthandler.GetSQSClient(pipelineName)
+	if !ok {
+		logger.Info("SQS client not found for the pipeline, will try again in a bit ...")
 		return ctrl.Result{RequeueAfter: 10 * time.Second}
 	}
-
 	// DrainSQSQueue long-polls (up to 20s), so it blocks until messages
 	// arrive or the timeout expires — no separate poll interval needed.
 	hasMessages, err := awsclienthandler.DrainSQSQueue(ctx, sqsClient, queueURL, bucket, prefix)

pkg/awsclienthandler/sqs_client.go

@@ -19,25 +19,93 @@ package awsclienthandler
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 )
 
-var (
-	SQSClient *sqs.Client
-)
+type SQSClientCache struct {
+	mu      sync.Mutex
+	clients map[string]*sqsEntry
+}
+
+type sqsEntry struct {
+	client    *sqs.Client
+	awsConfig AWSConfig
+	lastUsed  time.Time
+}
+
+var sqsClientCache = &SQSClientCache{
+	clients: make(map[string]*sqsEntry),
+}
+
+func InitSQSCache(ctx context.Context, ttl time.Duration) {
+	sqsClientCache.startCleanup(ctx, ttl)
+}
+
+func (c *SQSClientCache) getClient(pipelineName string) (*sqsEntry, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	entry, ok := c.clients[pipelineName]
+	if ok {
+		entry.lastUsed = time.Now()
+	}
+	return entry, ok
+}
+
+func (c *SQSClientCache) setClient(pipelineName string, entry *sqsEntry) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.clients == nil {
+		c.clients = make(map[string]*sqsEntry)
+	}
+	c.clients[pipelineName] = entry
+}
+
+func (c *SQSClientCache) startCleanup(ctx context.Context, ttl time.Duration) {
+	if ttl <= 0 {
+		return
+	}
+	ticker := time.NewTicker(ttl / 2)
+	go func() {
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				c.evictStale(ttl)
+			}
+		}
+	}()
+}
+
+func (c *SQSClientCache) evictStale(ttl time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	now := time.Now()
+	for key, entry := range c.clients {
+		if now.Sub(entry.lastUsed) > ttl {
+			delete(c.clients, key)
+		}
+	}
+}
 
 // NewSQSClientFromConfig creates and returns an Amazon SQS client using the provided context and AWS configuration.
-func NewSQSClientFromConfig(ctx context.Context, awsConfig *AWSConfig) (*sqs.Client, error) {
+func NewSQSClientFromConfig(ctx context.Context, awsConfig *AWSConfig, pipelineName string) (*sqs.Client, error) {
 	logger := log.FromContext(ctx)
-	if SQSClient != nil {
-		return SQSClient, nil
+
+	entry, ok := sqsClientCache.getClient(pipelineName)
+	if ok && awsConfig != nil && entry.awsConfig == *awsConfig {
+		logger.Info("Using existing SQS client for the pipeline", "pipelineName", pipelineName)
+		return entry.client, nil
 	}
 
+	logger.Info("Creating new SQS client for the pipeline", "pipelineName", pipelineName)
 	cfg, err := getAWSConfig(ctx, awsConfig)
 	if err != nil {
 		return nil, err
@@ -48,17 +116,32 @@ func NewSQSClientFromConfig(ctx context.Context, awsConfig *AWSConfig) (*sqs.Cli
 			o.BaseEndpoint = aws.String(awsConfig.Endpoint)
 		}
 	}
-	SQSClient = sqs.NewFromConfig(cfg, sqsOptions)
-	logger.Info("SQS client initialized ...")
-	return SQSClient, nil
+
+	sqsClient := sqs.NewFromConfig(cfg, sqsOptions)
+
+	var cfgVal AWSConfig
+	if awsConfig != nil {
+		cfgVal = *awsConfig
+	}
+	newEntry := &sqsEntry{client: sqsClient, awsConfig: cfgVal, lastUsed: time.Now()}
+	sqsClientCache.setClient(pipelineName, newEntry)
+
+	return sqsClient, nil
+}
+
+func DeleteSQSClient(pipelineName string) {
+	sqsClientCache.mu.Lock()
+	defer sqsClientCache.mu.Unlock()
+	delete(sqsClientCache.clients, pipelineName)
 }
 
 // GetSQSClient returns the initialized Amazon SQS client instance.
-func GetSQSClient() (*sqs.Client, error) {
-	if SQSClient == nil {
-		return nil, errors.New("SQS client not initialized yet")
+func GetSQSClient(pipelineName string) (*sqs.Client, bool) {
+	entry, ok := sqsClientCache.getClient(pipelineName)
+	if !ok {
+		return nil, false
 	}
-	return SQSClient, nil
+	return entry.client, true
 }
 
 const sqsLongPollSeconds = 20

test/e2e/unstructured_test.go

@@ -62,7 +62,7 @@
 	var kubeClient klient.Client
 	var sourceS3Client *s3.Client
 	var destS3Client *s3.Client
 	var sqsClient *sqs.Client
 
 	feature.Setup(
 		func(ctx context.Context, t *testing.T, cfg *envconf.Config) context.Context {
@@ -87,11 +87,18 @@
 			}
 
 			err = awsclienthandler.NewDestinationS3ClientFromConfig(ctx, e2eAWS)
+			// create SQS client
+			_, err = awsclienthandler.NewSQSClientFromConfig(ctx, &awsclienthandler.AWSConfig{
+				Region:          "us-east-1",
+				AccessKeyID:     "test",
+				SecretAccessKey: "test",
+				Endpoint:        localstackURL,
+			}, "e2e-test")
 			if err != nil {
 				t.Fatal(err)
 			}
 
 			sqsClient, err = awsclienthandler.NewSQSClientFromConfig(ctx, e2eAWS)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -131,6 +138,10 @@
 			}
 
 			// create SQS queue
+			sqsClient, ok := awsclienthandler.GetSQSClient("e2e-test")
+			if !ok {
+				t.Fatal("SQS client not found for e2e-test")
+			}
 			_, err = sqsClient.CreateQueue(ctx, &sqs.CreateQueueInput{
 				QueueName: aws.String(unstructuredQueueName),
 			})