Replace SQS singleton with per-pipeline client cache with TTL eviction - #291
Replace SQS singleton with per-pipeline client cache with TTL eviction#291gshikhar2021 wants to merge 2 commits into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refactors the SQS client management strategy to improve scalability and resource management. By moving from a single global client to a per-pipeline cache, the system can now effectively handle multiple concurrent pipelines. The addition of TTL-based eviction ensures that inactive clients are cleaned up, preventing memory leaks and stale connections. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request replaces the single global SQS client with a multi-tenant, TTL-evicted client cache (SQSClientCache) to support multiple pipelines. The review feedback highlights critical issues that need to be addressed: using only the resource name as a cache key can lead to cross-namespace cache collisions, so namespaced names should be used instead. Additionally, the cache currently lacks support for credential rotation because it does not compare incoming AWS configurations with cached ones, and the cleanup routine is susceptible to a panic if initialized with an invalid TTL.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if apierrors.IsNotFound(err) { | ||
| awsclienthandler.DeleteSQSClient(req.Name) | ||
| return ctrl.Result{}, nil | ||
| } |
There was a problem hiding this comment.
Using only req.Name as the cache key can lead to cache collisions if multiple SourceCrawler resources with the same name exist in different namespaces. Since this operator can watch multiple namespaces, we should use the namespaced name (e.g., req.NamespacedName.String()) as the cache key to prevent cross-namespace client leakage.
| if apierrors.IsNotFound(err) { | |
| awsclienthandler.DeleteSQSClient(req.Name) | |
| return ctrl.Result{}, nil | |
| } | |
| if apierrors.IsNotFound(err) { | |
| awsclienthandler.DeleteSQSClient(req.NamespacedName.String()) | |
| return ctrl.Result{}, 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 |
There was a problem hiding this comment.
Using sourceCrawlerCR.Name as the cache key can cause cache collisions across different namespaces. We should use the namespaced name string of the resource to ensure uniqueness.
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, namespacedName), nil| type sqsEntry struct { | ||
| client *sqs.Client | ||
| lastUsed time.Time | ||
| } |
There was a problem hiding this comment.
To support credential rotation and updates, we should store the AWSConfig in the sqsEntry so that we can detect when the configuration has changed and recreate the client accordingly.
| type sqsEntry struct { | |
| client *sqs.Client | |
| lastUsed time.Time | |
| } | |
| type sqsEntry struct { | |
| client *sqs.Client | |
| awsConfig AWSConfig | |
| lastUsed time.Time | |
| } |
| entry, ok := sqsClientCache.GetClient(pipelineName) | ||
| if ok { | ||
| logger.Info("Using existing SQS client for the pipeline", "pipelineName", pipelineName) | ||
| return entry.client, nil | ||
| } |
There was a problem hiding this comment.
If the AWS credentials are rotated or updated in the Secret, the controller will continue using the cached client with the old/stale credentials because it ignores the incoming awsConfig when a cached client exists. We should compare the incoming configuration with the cached one and recreate the client if they differ.
| entry, ok := sqsClientCache.GetClient(pipelineName) | |
| if ok { | |
| logger.Info("Using existing SQS client for the pipeline", "pipelineName", pipelineName) | |
| return entry.client, 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 | |
| } |
| newEntry := &sqsEntry{client: sqsClient, lastUsed: time.Now()} | ||
| sqsClientCache.SetClient(pipelineName, newEntry) |
There was a problem hiding this comment.
Store the AWSConfig in the new cache entry to enable configuration comparison on subsequent calls.
| 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) |
| func (c *SQSClientCache) startCleanup(ctx context.Context, ttl time.Duration) { | ||
| ticker := time.NewTicker(ttl / 2) |
There was a problem hiding this comment.
If ttl is less than or equal to 0, time.NewTicker(ttl / 2) will panic. We should add a defensive check to prevent panics if an invalid TTL is passed.
func (c *SQSClientCache) startCleanup(ctx context.Context, ttl time.Duration) {
if ttl <= 0 {
return
}
ticker := time.NewTicker(ttl / 2)
05bc52d to
fd51748
Compare
fd51748 to
6f74e1e
Compare
📝 WalkthroughWalkthroughThe application now initializes a TTL-based, pipeline-scoped SQS client cache. SourceCrawler reconciliation creates, retrieves, and deletes clients by pipeline identifier, and e2e setup uses named client lookup. ChangesSQS client cache lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SourceCrawlerReconciler
participant SQSClientCache
participant SQSClient
SourceCrawlerReconciler->>SQSClientCache: Create or retrieve client by namespace/name
SourceCrawlerReconciler->>SQSClientCache: Look up pipeline client
SQSClientCache-->>SourceCrawlerReconciler: Return client and presence status
SourceCrawlerReconciler->>SQSClient: Handle SQS wake-up
SourceCrawlerReconciler->>SQSClientCache: Delete client when resource is not found
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
test/e2e/unstructured_test.go (1)
65-65: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winStale call will not compile:
NewSQSClientFromConfignow requires apipelineNameargument.
sqs_client.gochanged the signature toNewSQSClientFromConfig(ctx, awsConfig, pipelineName), but line 101 still calls it with only 2 args — this is a build-breaking compile error. This call (and the outersqsClientdeclared at line 65) also appear to be dead leftovers: the test's actual SQS admin operations exclusively use the named"e2e-test"cached client fetched viaGetSQSClientat line 141, which currently shadows this outer variable. The simplest fix is to remove the stale block entirely rather than just adding the missing argument.🐛 Proposed fix
- var sqsClient *sqs.Client- sqsClient, err = awsclienthandler.NewSQSClientFromConfig(ctx, e2eAWS) - if err != nil { - t.Fatal(err) - } -Removing these also eliminates the local-vs-outer
sqsClientshadowing at line 141.Also applies to: 101-104, 141-144
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/e2e/unstructured_test.go` at line 65, Remove the unused outer sqsClient declaration and the stale NewSQSClientFromConfig initialization block in the test; retain the named “e2e-test” client obtained via GetSQSClient for all SQS admin operations, eliminating the local shadowing.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@test/e2e/unstructured_test.go`:
- Around line 89-99: Preserve and validate the error returned by
NewDestinationS3ClientFromConfig before initializing the SQS client, so it
cannot be overwritten by the subsequent NewSQSClientFromConfig assignment. Keep
the existing fatal behavior for either client initialization failure.
---
Outside diff comments:
In `@test/e2e/unstructured_test.go`:
- Line 65: Remove the unused outer sqsClient declaration and the stale
NewSQSClientFromConfig initialization block in the test; retain the named
“e2e-test” client obtained via GetSQSClient for all SQS admin operations,
eliminating the local shadowing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Enterprise
Run ID: 6d2898ae-c045-49c6-bc75-cb266c1a7849
📒 Files selected for processing (4)
cmd/main.gointernal/controller/sourcecrawler_controller.gopkg/awsclienthandler/sqs_client.gotest/e2e/unstructured_test.go
| 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) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Newly inserted code swallows the error from NewDestinationS3ClientFromConfig.
Line 89 assigns err from NewDestinationS3ClientFromConfig, but the inserted SQS-client-init block (lines 90-96) immediately reassigns err before it's checked — the original error is discarded and only the new call's error is validated at line 97.
🐛 Proposed fix
err = awsclienthandler.NewDestinationS3ClientFromConfig(ctx, e2eAWS)
+ if err != nil {
+ t.Fatal(err)
+ }
+
// 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)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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) | |
| } | |
| err = awsclienthandler.NewDestinationS3ClientFromConfig(ctx, e2eAWS) | |
| if err != nil { | |
| t.Fatal(err) | |
| } | |
| // 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) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/e2e/unstructured_test.go` around lines 89 - 99, Preserve and validate
the error returned by NewDestinationS3ClientFromConfig before initializing the
SQS client, so it cannot be overwritten by the subsequent NewSQSClientFromConfig
assignment. Keep the existing fatal behavior for either client initialization
failure.
1 participant
Summary by CodeRabbit
Bug Fixes
Performance
Tests