From 213cbefdcfd3ab472a2cf8ac5227d61f03473394 Mon Sep 17 00:00:00 2001 From: Shikhar Gupta Date: Wed, 8 Jul 2026 18:07:47 +0530 Subject: [PATCH 1/2] Replace SQS singleton with per-pipeline client cache with TTL eviction --- cmd/main.go | 9 +- .../controller/sourcecrawler_controller.go | 18 +-- pkg/awsclienthandler/sqs_client.go | 104 +++++++++++++++--- 3 files changed, 109 insertions(+), 22 deletions(-) diff --git a/cmd/main.go b/cmd/main.go index 5c9dcbcf..93cbb1ea 100644 --- a/cmd/main.go +++ b/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) } diff --git a/internal/controller/sourcecrawler_controller.go b/internal/controller/sourcecrawler_controller.go index 2feeb387..75c9d8e0 100644 --- a/internal/controller/sourcecrawler_controller.go +++ b/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.Name) + return ctrl.Result{}, nil + } logger.Error(err, "failed to get SourceCrawler CR") return ctrl.Result{}, err } @@ -161,24 +166,23 @@ 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 { + if _, err := awsclienthandler.NewSQSClientFromConfig(ctx, sourceAWSConfig, sourceCrawlerCR.Name); 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, sourceCrawlerCR.Name), 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) diff --git a/pkg/awsclienthandler/sqs_client.go b/pkg/awsclienthandler/sqs_client.go index 25548e2a..5f81c179 100644 --- a/pkg/awsclienthandler/sqs_client.go +++ b/pkg/awsclienthandler/sqs_client.go @@ -19,25 +19,89 @@ 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 + 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) { + 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 { + logger.Info("Using existing SQS client for the pipeline", "pipelineName", pipelineName) + return entry.client, nil } + logger.Info("No client is present right now for the pipeline, creating the new one") cfg, err := getAWSConfig(ctx, awsConfig) if err != nil { return nil, err @@ -48,17 +112,29 @@ 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) + logger.Info("New SQS client successfully created for the pipeline", "pipelineName", pipelineName) + + newEntry := &sqsEntry{client: sqsClient, 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 From 6f74e1e6e2fb509e4621bd24a9f06f8453069448 Mon Sep 17 00:00:00 2001 From: Shikhar Gupta Date: Tue, 21 Jul 2026 19:02:13 +0530 Subject: [PATCH 2/2] Address comments --- .../controller/sourcecrawler_controller.go | 7 +++-- pkg/awsclienthandler/sqs_client.go | 29 ++++++++++++------- test/e2e/unstructured_test.go | 11 +++++++ 3 files changed, 33 insertions(+), 14 deletions(-) diff --git a/internal/controller/sourcecrawler_controller.go b/internal/controller/sourcecrawler_controller.go index 75c9d8e0..28f302a6 100644 --- a/internal/controller/sourcecrawler_controller.go +++ b/internal/controller/sourcecrawler_controller.go @@ -79,7 +79,7 @@ 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.Name) + awsclienthandler.DeleteSQSClient(req.String()) return ctrl.Result{}, nil } logger.Error(err, "failed to get SourceCrawler CR") @@ -166,10 +166,11 @@ 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, sourceCrawlerCR.Name); 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, sourceCrawlerCR.Name), nil + return handleSQSWakeUp(ctx, sqsQueueURL, sourceCrawlerConfig.S3Config.Bucket, sourceCrawlerConfig.S3Config.Prefix, namespacedName), nil } } return ctrl.Result{RequeueAfter: defaultCrawlerResyncInterval}, nil diff --git a/pkg/awsclienthandler/sqs_client.go b/pkg/awsclienthandler/sqs_client.go index 5f81c179..6773ba38 100644 --- a/pkg/awsclienthandler/sqs_client.go +++ b/pkg/awsclienthandler/sqs_client.go @@ -34,8 +34,9 @@ type SQSClientCache struct { } type sqsEntry struct { - client *sqs.Client - lastUsed time.Time + client *sqs.Client + awsConfig AWSConfig + lastUsed time.Time } var sqsClientCache = &SQSClientCache{ @@ -46,7 +47,7 @@ func InitSQSCache(ctx context.Context, ttl time.Duration) { sqsClientCache.startCleanup(ctx, ttl) } -func (c *SQSClientCache) GetClient(pipelineName string) (*sqsEntry, bool) { +func (c *SQSClientCache) getClient(pipelineName string) (*sqsEntry, bool) { c.mu.Lock() defer c.mu.Unlock() entry, ok := c.clients[pipelineName] @@ -56,7 +57,7 @@ func (c *SQSClientCache) GetClient(pipelineName string) (*sqsEntry, bool) { return entry, ok } -func (c *SQSClientCache) SetClient(pipelineName string, entry *sqsEntry) { +func (c *SQSClientCache) setClient(pipelineName string, entry *sqsEntry) { c.mu.Lock() defer c.mu.Unlock() if c.clients == nil { @@ -66,6 +67,9 @@ func (c *SQSClientCache) SetClient(pipelineName string, entry *sqsEntry) { } func (c *SQSClientCache) startCleanup(ctx context.Context, ttl time.Duration) { + if ttl <= 0 { + return + } ticker := time.NewTicker(ttl / 2) go func() { defer ticker.Stop() @@ -95,13 +99,13 @@ func (c *SQSClientCache) evictStale(ttl time.Duration) { func NewSQSClientFromConfig(ctx context.Context, awsConfig *AWSConfig, pipelineName string) (*sqs.Client, error) { logger := log.FromContext(ctx) - entry, ok := sqsClientCache.GetClient(pipelineName) - if ok { + 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("No client is present right now for the pipeline, creating the new one") + logger.Info("Creating new SQS client for the pipeline", "pipelineName", pipelineName) cfg, err := getAWSConfig(ctx, awsConfig) if err != nil { return nil, err @@ -114,10 +118,13 @@ func NewSQSClientFromConfig(ctx context.Context, awsConfig *AWSConfig, pipelineN } sqsClient := sqs.NewFromConfig(cfg, sqsOptions) - logger.Info("New SQS client successfully created for the pipeline", "pipelineName", pipelineName) - newEntry := &sqsEntry{client: sqsClient, lastUsed: time.Now()} - sqsClientCache.SetClient(pipelineName, newEntry) + var cfgVal AWSConfig + if awsConfig != nil { + cfgVal = *awsConfig + } + newEntry := &sqsEntry{client: sqsClient, awsConfig: cfgVal, lastUsed: time.Now()} + sqsClientCache.setClient(pipelineName, newEntry) return sqsClient, nil } @@ -130,7 +137,7 @@ func DeleteSQSClient(pipelineName string) { // GetSQSClient returns the initialized Amazon SQS client instance. func GetSQSClient(pipelineName string) (*sqs.Client, bool) { - entry, ok := sqsClientCache.GetClient(pipelineName) + entry, ok := sqsClientCache.getClient(pipelineName) if !ok { return nil, false } diff --git a/test/e2e/unstructured_test.go b/test/e2e/unstructured_test.go index 92c3f274..5e0efa3b 100644 --- a/test/e2e/unstructured_test.go +++ b/test/e2e/unstructured_test.go @@ -87,6 +87,13 @@ func TestUnstructuredDataLoad(t *testing.T) { } 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) } @@ -131,6 +138,10 @@ func TestUnstructuredDataLoad(t *testing.T) { } // 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), })