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..28f302a6 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.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) diff --git a/pkg/awsclienthandler/sqs_client.go b/pkg/awsclienthandler/sqs_client.go index 25548e2a..6773ba38 100644 --- a/pkg/awsclienthandler/sqs_client.go +++ b/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 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), })