Conversation
WalkthroughAdds a new S3 connector and concurrent batch writer: parses s3:// URIs, configures an AWS S3 client, implements planning/list/read/write flows with per-namespace metadata and optimistic-lock handling, provides a buffered concurrent BatchProcessor, CLI flags, and updates AWS SDK dependencies. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant Conn as S3 Connector
participant BP as BatchProcessor
participant S3 as AWS S3
App->>Conn: WriteData(namespace, documents)
activate Conn
Conn->>Conn: Validate & convert docs (JSON/BSON)
Conn->>BP: Add(namespace, [][]byte)
activate BP
BP->>BP: Check memory & append to buffer
alt Memory pressure
BP->>BP: Evict largest buffer / wait for space
end
alt Buffer >= MaxFileSize or Close triggered
BP->>S3: PutObject (async upload)
activate S3
S3-->>BP: Upload result (ETag / error)
deactivate S3
BP->>S3: Get/Update `.metadata.json` (optimistic-lock / retry)
end
BP-->>Conn: Add result / error
deactivate BP
Conn-->>App: Return success / error
deactivate Conn
sequenceDiagram
participant App as Application
participant Conn as S3 Connector
participant S3 as AWS S3
App->>Conn: GeneratePlan(requested namespaces)
activate Conn
Conn->>S3: ListObjectsV2 (prefix) [paged]
activate S3
S3-->>Conn: Object pages
deactivate S3
Conn->>Conn: Filter objects by namespaces, build metadata keys
loop per-namespace
Conn->>S3: GetObject(.metadata.json)
alt metadata found
S3-->>Conn: metadata JSON
Conn->>Conn: Sum counts, cache
else not found
S3-->>Conn: NotFound
Conn->>Conn: Use zero counts
end
end
Conn->>App: Emit partitions (object key cursor + estimated counts)
deactivate Conn
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
go.mod (1)
10-58: Consider minor AWS SDK version updates for latest stability patches.The AWS SDK v2 dependency structure is correct with proper indirect dependencies declared. However, newer patch versions are available: S3 service can upgrade from v1.92.1 to v1.93.0, and the base SDK from v1.40.0 to v1.40.1. No security advisories exist for either current version, but these patch updates may include bug fixes and stability improvements worth adopting.
🧹 Nitpick comments (4)
connectors/s3/s3writer.go (2)
70-94: Memory accounting may undercount actual usage.The
dataLencalculation (lines 70-73) accounts for the raw input size, but the actual buffer grows larger due to:
- JSON array delimiters (
[,],,)- Pretty-print formatting (newlines, indentation) when
PrettyJSONis enabledThis could cause actual memory usage to exceed
MaxTotalMemoryby a non-trivial margin, especially with many small documents and pretty-printing enabled.Consider accounting for formatting overhead, or document that
MaxTotalMemoryis approximate:dataLen := int64(0) for _, d := range data { - dataLen += int64(len(d)) + overhead := int64(2) // comma + potential newline + if bp.config.PrettyJSON { + overhead += int64(len(d) / 10) // rough estimate for indentation + } + dataLen += int64(len(d)) + overhead }
242-254: Potential key collision with nanosecond timestamps.Under high concurrency, two goroutines could generate the same
UnixNano()timestamp, causing one upload to overwrite another. While unlikely, this could cause data loss.Consider adding a unique identifier:
+import "github.com/google/uuid" + func (bp *BatchProcessor) objectKey(namespace string) string { nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") if nsPath == "" { nsPath = "default" } - fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) + fileName := fmt.Sprintf("%d-%s.json", time.Now().UnixNano(), uuid.New().String()[:8])internal/app/options/connectorflags.go (1)
608-622: Credentials in CLI flags may be exposed in process listings.The
access-key-id,secret-access-key, andsession-tokenflags will be visible in process listings (ps aux) and shell history. This is a security concern for production use.Consider:
- Documenting that environment variables (
AWS_ACCESS_KEY_ID, etc.) or IAM roles are preferred- Adding a note in the flag usage about security implications
- Supporting file-based credential input
altsrc.NewStringFlag(&cli.StringFlag{ Name: "access-key-id", - Usage: "Static AWS access key ID", + Usage: "Static AWS access key ID (prefer AWS_ACCESS_KEY_ID env var or IAM roles for security)", Destination: &settings.AccessKeyID, }), altsrc.NewStringFlag(&cli.StringFlag{ Name: "secret-access-key", - Usage: "Static AWS secret access key", + Usage: "Static AWS secret access key (prefer AWS_SECRET_ACCESS_KEY env var or IAM roles for security)", Destination: &settings.SecretAccessKey, }),connectors/s3/connector.go (1)
465-504: Code duplication with s3writer.go for metadata handling.The
metadataKeyandreadMetadatafunctions duplicate logic froms3writer.go(lines 256-291). This creates maintenance burden and risk of inconsistency.Consider refactoring to share the metadata logic. Options:
- Export
MetadataKeyfrom theBatchProcessoror make it a standalone function- Create a shared
metadata.gofile with common helpers+// In s3writer.go, export the function: +func MetadataKey(prefix, namespace string) string { + nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") + if nsPath == "" { + nsPath = "default" + } + prefix = strings.Trim(prefix, "/") + if prefix == "" { + return path.Join(nsPath, ".metadata.json") + } + return path.Join(prefix, nsPath, ".metadata.json") +} // In connector.go, use the shared function: -func (c *connector) metadataKey(namespace string) string { - // ... duplicated code ... -} +func (c *connector) metadataKey(namespace string) string { + return MetadataKey(c.settings.Prefix, namespace) +}
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (4)
connectors/s3/connector.go(1 hunks)connectors/s3/s3writer.go(1 hunks)go.mod(2 hunks)internal/app/options/connectorflags.go(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
internal/app/options/connectorflags.go (1)
connectors/s3/connector.go (2)
ConnectorSettings(37-52)NewConn(84-157)
connectors/s3/connector.go (4)
connectors/s3/s3writer.go (3)
BatchProcessor(36-52)NewBatchProcessor(54-63)Config(21-27)gen/adiom/v1/adiomv1connect/adiom.connect.go (1)
ConnectorServiceHandler(214-225)gen/adiom/v1/messages.pb.go (2)
DataType_DATA_TYPE_JSON_ID(29-29)DataType_DATA_TYPE_MONGO_BSON(28-28)protocol/iface/connector.go (1)
Namespace(104-107)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (5)
connectors/s3/s3writer.go (1)
223-240: LGTM - Clean shutdown logic.The shutdown sequence correctly:
- Sets the shutdown flag under lock
- Collects buffer names before iteration to avoid map modification during iteration
- Releases the lock before
wg.Wait()to allow upload goroutines to completeConsider adding a timeout to
wg.Wait()for production resilience, but this is acceptable for initial implementation.internal/app/options/connectorflags.go (1)
300-316: LGTM - S3 connector registration follows established patterns.The registration correctly:
- Uses
strings.HasPrefix(strings.ToLower(s), "s3://")for case-insensitive matching- Initializes settings with the URI before flag parsing
- Delegates to
CreateHelperconsistent with other connectors like CosmosDB and MongoDBconnectors/s3/connector.go (3)
83-157: LGTM - Well-structured connector initialization.The constructor properly:
- Validates required fields before proceeding
- Applies sensible defaults
- Supports multiple credential sources (default chain, profile, static)
- Enables S3-compatible service usage via endpoint and path-style options
392-418: LGTM - WriteData handles errors appropriately.The implementation correctly:
- Checks for existing errors to fail-fast
- Returns early for empty batches
- Propagates errors to block subsequent writes via
setError()- Converts data types appropriately before batching
425-427:Teardown()is properly part of the connector lifecycle.The
Teardown()method is required by theConnectorinterface (protocol/iface/connector.go:129), not theConnectorServiceHandlerinterface. The S3 connector correctly implements this interface, and the method is invoked during shutdown through the runner orchestration (runners/local/runner.go:312-319 callsr.src.Teardown()andr.dst.Teardown()as part of the application lifecycle). ThebatchProcessor.Close()call will execute as expected, and there is no risk of data loss.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
connectors/s3/connector.go (2)
361-377: JSON to BSON conversion loses Extended JSON type fidelity.This concern was raised in a previous review. Using
json.Unmarshaltomap[string]interface{}thenbson.Marshalloses MongoDB Extended JSON types ($oid,$date,$numberLongbecome nested maps instead of proper BSON types).Consider using
bson.UnmarshalExtJSONfor proper Extended JSON handling:case adiomv1.DataType_DATA_TYPE_MONGO_BSON: var rawDocs []json.RawMessage if err := json.NewDecoder(getOut.Body).Decode(&rawDocs); err != nil { return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("decode json array from %s: %w", key, err)) } for _, d := range rawDocs { - // Convert JSON to BSON - var doc map[string]interface{} - if err := json.Unmarshal(d, &doc); err != nil { - return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("unmarshal json to map in %s: %w", key, err)) + // Convert Extended JSON to BSON + var doc bson.D + if err := bson.UnmarshalExtJSON(d, true, &doc); err != nil { + return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("unmarshal extended json in %s: %w", key, err)) } - bsonDoc, err := bson.Marshal(doc) + bsonDoc, err := bson.Marshal(doc)
456-465: BSON to JSON conversion loses Extended JSON type information.This concern was raised in a previous review. Using
bson.Unmarshaltomap[string]anythenjson.Marshalloses BSON type information (ObjectIds, Dates, Decimals serialize as plain strings/numbers).Use
bson.MarshalExtJSONfor type-preserving conversion:case adiomv1.DataType_DATA_TYPE_MONGO_BSON: - var doc map[string]any - if err := bson.Unmarshal(data, &doc); err != nil { + var doc bson.D + if err := bson.Unmarshal(data, &doc); err != nil { return nil, fmt.Errorf("bson to json: %w", err) } - converted, err := json.Marshal(doc) + converted, err := bson.MarshalExtJSON(doc, true, false) if err != nil { return nil, fmt.Errorf("marshal json: %w", err) }
🧹 Nitpick comments (3)
connectors/s3/util.go (1)
5-7: Consider using a switch or set for better readability.The long OR-chain is hard to scan. A switch or map lookup would be clearer and easier to extend.
-func IsS3OptimisticLockFailedError(apiError smithy.APIError) bool { - return apiError.ErrorCode() == "PreconditionFailed" || apiError.ErrorCode() == "412" || apiError.ErrorCode() == "409" || apiError.ErrorCode() == "ConditionalRequestConflict" -} +func IsS3OptimisticLockFailedError(apiError smithy.APIError) bool { + switch apiError.ErrorCode() { + case "PreconditionFailed", "412", "409", "ConditionalRequestConflict": + return true + default: + return false + } +}connectors/s3/s3writer.go (1)
158-228: LGTM with operational note.Good improvements: timeout context (line 189) and proper variable capture in the goroutine. The memory tracking strategy is well-documented.
Lines 202-208: Upload failures are logged but silently dropped. Consider whether failed uploads should be tracked for observability (e.g., a counter metric) or reported back to the caller. The comment acknowledges this as a future improvement.
connectors/s3/connector.go (1)
92-95: Original parse error is discarded.The error message loses the detail from
parseS3ConnectionString. Consider wrapping instead:bucket, prefix, err := parseS3ConnectionString(settings.Uri) if err != nil { - return nil, fmt.Errorf("bad uri format %v", settings.Uri) + return nil, fmt.Errorf("bad uri format: %w", err) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
connectors/s3/connector.go(1 hunks)connectors/s3/s3writer.go(1 hunks)connectors/s3/util.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/s3/s3writer.go (3)
connectors/s3/connector.go (2)
MAX_S3_WRITE_TIMEOUT_SEC(39-39)S3_METADATA_UPDATES_RETRY(40-40)connectors/s3/metadata.go (1)
MetadataKey(9-20)connectors/s3/util.go (1)
IsS3OptimisticLockFailedError(5-7)
🔇 Additional comments (19)
connectors/s3/util.go (1)
9-20: LGTM!The custom error type correctly implements both
Error()andUnwrap()for proper error chain support. The privateerrfield is appropriate.connectors/s3/s3writer.go (7)
20-33: LGTM!The
Configstruct has clear documentation for the limit fields, andbufferInfoappropriately encapsulates buffer state.
35-63: LGTM!The
BatchProcessorstruct properly separates concerns (state, lifecycle, metadata), and the constructor correctly initializes the condition variable with the mutex.
65-134: LGTM with minor note.The backpressure mechanism, buffer management, and JSON formatting are well-implemented. The fix for the infinite loop case (lines 79-81) properly addresses the edge case.
Note: Memory accounting at line 99 uses the input
dataLenrather than the actual formatted buffer size. WithPrettyJSONenabled, the real memory usage will be higher due to indentation. This is acceptable as an approximation but worth documenting.
136-156: LGTM!Simple and effective eviction strategy. The O(n) scan is acceptable given typical namespace counts.
230-248: LGTM!The shutdown sequence correctly flushes remaining buffers and releases the lock before waiting on the WaitGroup to avoid deadlock.
268-314: LGTM!Proper use of ETag-based optimistic locking with conditional PutObject. The
NoSuchKeyhandling correctly returns nil for non-existent metadata files.
346-366: Concurrency pattern is sound.The combination of
metadataMutex(serializes within-process) and ETag-based optimistic locking with retries (handles S3-level races) correctly addresses the race condition concern from previous reviews. This is a valid read-modify-write pattern for S3.connectors/s3/connector.go (11)
30-41: LGTM!Sentinel errors and constants are well-defined with clear names and appropriate default values.
43-70: LGTM!
ConnectorSettingsprovides comprehensive configuration. The embeddedUnimplementedConnectorServiceHandlerensures forward compatibility with interface changes.
72-88: LGTM!The parsing logic correctly handles both
s3://bucketands3://bucket/prefixformats with case-insensitive prefix matching.
166-188: LGTM!Capabilities are correctly advertised for S3 as both source and sink with appropriate data type support.
190-295: LGTM!The plan generation correctly paginates S3 objects, extracts namespaces from paths (fixing the previous asymmetry issue), and efficiently caches metadata per namespace.
297-329: LGTM!Graceful handling of missing or unreadable metadata files by returning a count of 0.
388-396: LGTM!Streaming methods correctly return unimplemented errors as S3 doesn't support change streams.
398-424: LGTM!Proper error state checking before processing, document conversion, and error propagation to prevent further writes after a failure.
426-433: LGTM!
WriteUpdatescorrectly returns unsupported, andTeardownproperly closes the batch processor.
435-447: LGTM!Thread-safe error state with RWMutex and first-error-wins semantics is appropriate for async error propagation.
471-502: LGTM!Correct implementation with
NoSuchKeyhandling consistent with the pattern ins3writer.go.
| func (bp *BatchProcessor) objectKey(namespace string) string { | ||
| nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") | ||
| if nsPath == "" { | ||
| nsPath = "default" | ||
| } | ||
| fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) | ||
|
|
||
| prefix := strings.Trim(bp.config.Prefix, "/") | ||
| if prefix == "" { | ||
| return path.Join(nsPath, fileName) | ||
| } | ||
| return path.Join(prefix, nsPath, fileName) | ||
| } |
There was a problem hiding this comment.
Potential filename collision under high concurrency.
time.Now().UnixNano() can return the same value for multiple goroutines executing simultaneously, leading to S3 key collisions and data loss.
Consider adding a random suffix or atomic counter:
+import "math/rand"
+
func (bp *BatchProcessor) objectKey(namespace string) string {
nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/")
if nsPath == "" {
nsPath = "default"
}
- fileName := fmt.Sprintf("%d.json", time.Now().UnixNano())
+ fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), rand.Int63())
prefix := strings.Trim(bp.config.Prefix, "/")
if prefix == "" {
return path.Join(nsPath, fileName)
}
return path.Join(prefix, nsPath, fileName)
}Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In connectors/s3/s3writer.go around lines 250 to 262, the objectKey() uses
time.Now().UnixNano() alone which can collide under high concurrency; change the
filename generation to append a short unique suffix (for example a UUID or a
process-wide atomic counter combined with a few random bytes) to ensure
uniqueness across goroutines, make the counter atomic or use a thread-safe
UUID/random generator, and update imports accordingly so the returned key is
time-based plus the unique suffix (e.g., "<timestamp>-<uniq>.json") to eliminate
S3 key collisions.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
connectors/s3/s3writer.go (1)
251-263: Filename collision risk remains unaddressed.The
objectKeymethod still uses onlytime.Now().UnixNano()(line 256) without additional uniqueness guarantees. Under high concurrency, multiple goroutines can execute within the same nanosecond, leading to S3 key collisions and potential data loss.As suggested in the previous review, add a random suffix or process-wide atomic counter:
+import ( + "sync/atomic" +) + +var fileCounter uint64 + func (bp *BatchProcessor) objectKey(namespace string) string { nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") if nsPath == "" { nsPath = "default" } - fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) + seq := atomic.AddUint64(&fileCounter, 1) + fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), seq) prefix := strings.Trim(bp.config.Prefix, "/") if prefix == "" { return path.Join(nsPath, fileName) } return path.Join(prefix, nsPath, fileName) }Alternatively, use a random suffix:
fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), rand.Int63())
🧹 Nitpick comments (1)
connectors/s3/s3writer.go (1)
36-53: Well-structured processor with proper synchronization primitives.The
BatchProcessoruses appropriate synchronization mechanisms: mutex + condition variable for backpressure, WaitGroup for lifecycle management, and a dedicated metadataMutex for metadata updates.Optional: Consider per-namespace locks instead of a global
metadataMutexto allow concurrent metadata updates for different namespaces. However, since metadata updates involve S3 I/O, the current serialization overhead is likely negligible.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
connectors/s3/s3writer.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/s3/s3writer.go (3)
connectors/s3/connector.go (2)
MAX_S3_WRITE_TIMEOUT_SEC(39-39)S3_METADATA_UPDATES_RETRY(40-40)connectors/s3/metadata.go (1)
MetadataKey(9-20)connectors/s3/util.go (2)
IsS3OptimisticLockFailedError(5-7)ErrS3OptimisticLockFailed(10-12)
🔇 Additional comments (11)
connectors/s3/s3writer.go (11)
21-34: LGTM! Clear configuration and buffer tracking structures.The
Configstruct provides all necessary parameters for S3 batch processing, andbufferInfoappropriately tracks per-namespace buffer state.
55-64: LGTM! Constructor properly initializes all components.The condition variable is correctly bound to the mutex, and all fields are appropriately initialized.
66-135: Excellent implementation of backpressure and buffering logic.The
Addmethod correctly:
- Prevents infinite loops by rejecting batches exceeding total memory (lines 80-82) ✓
- Implements backpressure with eviction and condition variable signaling ✓
- Handles buffer initialization and JSON array formatting ✓
- Provides fallback for pretty JSON failures ✓
- Triggers flush when file size limit is reached ✓
137-157: LGTM! Effective eviction strategy with helpful logging.The method correctly identifies and flushes the largest buffer under memory pressure, with appropriate debug information.
159-229: Excellent async upload implementation with proper resource management.The method correctly:
- Snapshots buffer data before launching goroutine ✓
- Uses timeout context to prevent indefinite blocking (line 190) ✓
- Removes buffer from map immediately to unblock new writes ✓
- Properly manages WaitGroup for lifecycle tracking ✓
- Reclaims memory and signals waiting threads upon completion ✓
- Handles metadata updates with retry logic ✓
231-249: LGTM! Proper shutdown sequence with graceful flush.The
Closemethod correctly flushes all remaining buffers, unlocks before waiting (avoiding deadlock), and ensures all uploads complete via WaitGroup.
265-267: LGTM! Clean delegation to shared metadata key generation.The wrapper appropriately delegates to the
MetadataKeyutility function.
269-291: LGTM! Proper metadata retrieval with ETag capture.The method correctly:
- Handles non-existent metadata by returning nil (line 279) ✓
- Captures ETag for optimistic locking (line 290) ✓
- Properly closes the response body ✓
293-315: LGTM! Proper optimistic locking implementation.The method correctly:
- Uses
IfMatchwith ETag for conditional writes (line 305) ✓- Detects and wraps optimistic lock failures for retry handling (lines 309-310) ✓
- Sets appropriate content type ✓
317-346: Excellent retry logic with exponential backoff and jitter.The method correctly:
- Implements proper exponential backoff (line 341) ✓
- Adds jitter to reduce thundering herd (line 342) ✓
- Only retries optimistic lock failures (lines 328-331) ✓
- Returns non-retryable errors immediately ✓
This addresses previous review concerns about backoff calculation.
348-368: LGTM! Proper metadata update with local and distributed synchronization.The method correctly combines:
- Local mutex to prevent concurrent updates in the same process (line 349) ✓
- Optimistic locking via ETag to handle cross-process/cross-instance races ✓
- Clean read-modify-write pattern with proper error propagation ✓
This addresses the previous review concern about metadata race conditions.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
connectors/s3/s3writer.go (2)
10-10: Seedmath/randfor non-deterministic jitter.The
math/randpackage is used for jitter (line 343) but isn't seeded, producing the same sequence on every run. Add aninit()function or switch tomath/rand/v2(Go 1.22+) which auto-seeds.+func init() { + rand.Seed(time.Now().UnixNano()) +} + // Config holds our limits type Config struct {
251-263: Potential filename collision under high concurrency.
time.Now().UnixNano()can return identical values for concurrent goroutines, risking S3 key collisions and data overwrites.Add a random suffix or atomic counter:
+var objectCounter uint64 + func (bp *BatchProcessor) objectKey(namespace string) string { nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") if nsPath == "" { nsPath = "default" } - fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) + counter := atomic.AddUint64(&objectCounter, 1) + fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), counter) prefix := strings.Trim(bp.config.Prefix, "/")Add
"sync/atomic"to imports.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
connectors/s3/s3writer.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/s3/s3writer.go (3)
connectors/s3/connector.go (2)
MAX_S3_WRITE_TIMEOUT_SEC(39-39)S3_METADATA_UPDATES_RETRY(40-40)connectors/s3/metadata.go (1)
MetadataKey(9-20)connectors/s3/util.go (2)
IsS3OptimisticLockFailedError(5-7)ErrS3OptimisticLockFailed(10-12)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (7)
connectors/s3/s3writer.go (7)
36-64: Well-structured concurrent processor design.Good use of
sync.Condfor backpressure signaling and separatemetadataMutexfor serializing metadata operations without blocking data ingestion.
76-97: Backpressure mechanism correctly prevents infinite loop.The pre-check at lines 80-82 properly guards against the case where a single batch exceeds
MaxTotalMemory, preventing the wait loop from becoming infinite.
231-249: Clean shutdown implementation.Properly collects buffer keys before iteration to avoid concurrent modification, and uses
WaitGroupto ensure all in-flight uploads complete.
269-291: Correct handling of missing metadata and ETag extraction.Gracefully handles
NoSuchKeyfor first-time writes and properly returns the ETag for subsequent optimistic locking.
293-315: Proper optimistic locking with ETag.Uses
IfMatchfor conditional writes and correctly wraps precondition failures for retry handling.
317-347: Solid retry logic with exponential backoff.Correctly identifies retryable precondition failures and implements proper exponential backoff (50ms, 100ms, 200ms) with jitter.
349-369: Serialized metadata updates with optimistic locking.The
metadataMutexprevents local race conditions, while ETag-based conditional writes with retry logic handle cross-process conflicts.
| if err != nil { | ||
| slog.Error("Upload failed", | ||
| "namespace", ns, | ||
| "key", key, | ||
| "error", err, | ||
| ) | ||
| // In a real system, you might implement a retry mechanism or Dead Letter Queue here. | ||
| } else { | ||
| // Update metadata on successful upload | ||
| if err := bp.updateMetadataWithRetries(ctx, ns, key, uint64(numDocs)); err != nil { | ||
| slog.Error("Failed to update metadata after retries", | ||
| "namespace", ns, | ||
| "key", key, | ||
| "error", err, | ||
| ) | ||
| } |
There was a problem hiding this comment.
Upload failures result in silent data loss.
When PutObject fails, the data is logged and discarded. The comment mentions retry/DLQ but nothing is implemented. Consider adding a retry mechanism or at minimum persisting failed uploads to a local file for recovery.
Would you like me to generate a retry mechanism with exponential backoff for failed uploads, or a simple file-based dead letter queue?
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
connectors/s3/s3writer.go (1)
251-263: Potential filename collision under high concurrency.
time.Now().UnixNano()can return the same value for multiple goroutines executing simultaneously, leading to S3 key collisions and data loss.This was flagged in a previous review but appears unaddressed. Consider adding a random suffix or atomic counter:
+var keyCounter uint64 + func (bp *BatchProcessor) objectKey(namespace string) string { nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") if nsPath == "" { nsPath = "default" } - fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) + counter := atomic.AddUint64(&keyCounter, 1) + fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), counter) prefix := strings.Trim(bp.config.Prefix, "/")
🧹 Nitpick comments (2)
connectors/s3/s3writer.go (2)
203-219: Upload failures result in silent data loss.When
PutObjectfails, the data is logged and discarded. The XXX comment acknowledges this, but for production reliability, consider tracking failed uploads for later retry or alerting.One approach is to track failed uploads in a channel for later processing or persistence:
type failedUpload struct { namespace string key string data []byte err error } // Add to BatchProcessor: failedUploads chan failedUploadAlternatively, ensure upstream callers are aware that
Close()returning does not guarantee all data was successfully persisted.
231-249:Close()may block indefinitely if S3 uploads are slow or unresponsive.While individual uploads now have a 300-second timeout,
Close()callsbp.wg.Wait()without any overall timeout. If there are many in-flight uploads or the network is degraded,Close()could block for an extended period.Consider adding an optional deadline to
Close():func (bp *BatchProcessor) Close() error { // ... flush logic ... done := make(chan struct{}) go func() { bp.wg.Wait() close(done) }() select { case <-done: slog.Debug("Finished final S3 flush") return nil case <-time.After(10 * time.Minute): return fmt.Errorf("timeout waiting for uploads to complete") } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
connectors/s3/s3writer.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/s3/s3writer.go (3)
connectors/s3/connector.go (2)
MAX_S3_WRITE_TIMEOUT_SEC(39-39)S3_METADATA_UPDATES_RETRY(40-40)connectors/s3/metadata.go (1)
MetadataKey(9-20)connectors/s3/util.go (2)
IsS3OptimisticLockFailedError(5-7)ErrS3OptimisticLockFailed(10-12)
🔇 Additional comments (9)
connectors/s3/s3writer.go (9)
1-19: LGTM!Good choice using
math/rand/v2which auto-seeds, addressing the previous concern about deterministic jitter.
21-28: LGTM!Config struct is well-defined. Consider adding validation in the constructor to ensure
MaxFileSize <= MaxTotalMemoryand both are positive, but this is optional.
30-53: LGTM!Good use of
sync.Condfor backpressure signaling andmetadataMutexto serialize metadata updates within the process. The optimistic locking with retries handles cross-process concurrency.
55-64: LGTM!Constructor correctly initializes all state and the condition variable.
137-157: LGTM!Eviction strategy correctly identifies and flushes the largest buffer to maximize freed memory.
265-267: LGTM!Clean delegation to the shared
MetadataKeyutility function.
269-291: LGTM!Proper handling of
NoSuchKeyfor new namespaces and correct ETag extraction for optimistic locking.
317-346: LGTM!Well-implemented retry logic with exponential backoff and jitter. Only retries on optimistic lock failures (precondition errors), correctly propagating other errors immediately.
348-368: LGTM!Proper serialization with
metadataMutexfor intra-process concurrency combined with optimistic locking for inter-process safety. The read-modify-write pattern is correctly protected.
| for i, doc := range data { | ||
| if bp.buffers[namespace].docCount > 0 || i > 0 { | ||
| bp.buffers[namespace].buffer.WriteByte(',') | ||
| } | ||
| if bp.config.PrettyJSON { | ||
| var prettyBuf bytes.Buffer | ||
| prettyBuf.Grow(len(doc) + len(doc)/10) | ||
| if err := json.Indent(&prettyBuf, doc, "", " "); err == nil { | ||
| bp.buffers[namespace].buffer.WriteByte('\n') | ||
| doc = prettyBuf.Bytes() | ||
| } else { | ||
| slog.Warn("Failed JSON indentation. Falling back to no-indent") | ||
| } | ||
| } | ||
| bp.buffers[namespace].buffer.Write(doc) | ||
| bp.buffers[namespace].docCount++ | ||
| } |
There was a problem hiding this comment.
Memory tracking may undercount when PrettyJSON is enabled.
When PrettyJSON is true, json.Indent can significantly expand the document size (typically 20-50% larger). However, dataLen was calculated from the original compact data, and currentUsage is incremented by this smaller value. The actual buffer may grow larger than tracked, potentially exceeding MaxTotalMemory in practice.
Consider recalculating the size after indentation or using a conservative multiplier:
for i, doc := range data {
if bp.buffers[namespace].docCount > 0 || i > 0 {
bp.buffers[namespace].buffer.WriteByte(',')
}
if bp.config.PrettyJSON {
var prettyBuf bytes.Buffer
prettyBuf.Grow(len(doc) + len(doc)/10)
if err := json.Indent(&prettyBuf, doc, "", " "); err == nil {
bp.buffers[namespace].buffer.WriteByte('\n')
+ // Update memory tracking for expanded size
+ bp.currentUsage += int64(prettyBuf.Len() - len(doc))
doc = prettyBuf.Bytes()
} else {
slog.Warn("Failed JSON indentation. Falling back to no-indent")
}
}📝 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.
| for i, doc := range data { | |
| if bp.buffers[namespace].docCount > 0 || i > 0 { | |
| bp.buffers[namespace].buffer.WriteByte(',') | |
| } | |
| if bp.config.PrettyJSON { | |
| var prettyBuf bytes.Buffer | |
| prettyBuf.Grow(len(doc) + len(doc)/10) | |
| if err := json.Indent(&prettyBuf, doc, "", " "); err == nil { | |
| bp.buffers[namespace].buffer.WriteByte('\n') | |
| doc = prettyBuf.Bytes() | |
| } else { | |
| slog.Warn("Failed JSON indentation. Falling back to no-indent") | |
| } | |
| } | |
| bp.buffers[namespace].buffer.Write(doc) | |
| bp.buffers[namespace].docCount++ | |
| } | |
| for i, doc := range data { | |
| if bp.buffers[namespace].docCount > 0 || i > 0 { | |
| bp.buffers[namespace].buffer.WriteByte(',') | |
| } | |
| if bp.config.PrettyJSON { | |
| var prettyBuf bytes.Buffer | |
| prettyBuf.Grow(len(doc) + len(doc)/10) | |
| if err := json.Indent(&prettyBuf, doc, "", " "); err == nil { | |
| bp.buffers[namespace].buffer.WriteByte('\n') | |
| // Update memory tracking for expanded size | |
| bp.currentUsage += int64(prettyBuf.Len() - len(doc)) | |
| doc = prettyBuf.Bytes() | |
| } else { | |
| slog.Warn("Failed JSON indentation. Falling back to no-indent") | |
| } | |
| } | |
| bp.buffers[namespace].buffer.Write(doc) | |
| bp.buffers[namespace].docCount++ | |
| } |
🤖 Prompt for AI Agents
In connectors/s3/s3writer.go around lines 111 to 127, memory accounting uses the
original compact document length even when PrettyJSON=true and json.Indent
expands the bytes, causing currentUsage to undercount and potentially exceed
MaxTotalMemory; recalculate the actual byte size after indentation (e.g., use
len(prettyBytes) when indent succeeds) or apply a conservative multiplier before
reserving memory, then increment currentUsage by that actual/reserved size and
validate against MaxTotalMemory before writing to the buffer; ensure the
fallback path still uses the correct length and update any buffer.Grow logic to
match the recalculated size.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
connectors/s3/s3writer.go (2)
112-128: Memory tracking undercount when PrettyJSON expands documents.When
PrettyJSONis enabled,json.Indenttypically expands documents by 20-50%, butcurrentUsage(incremented at line 101) reflects only the original compact size. The actual buffer grows larger than tracked, potentially exceedingMaxTotalMemoryand causing OOM despite backpressure.Consider recalculating after indentation:
for i, doc := range data { if bp.buffers[namespace].docCount > 0 || i > 0 { bp.buffers[namespace].buffer.WriteByte(',') } if bp.config.PrettyJSON { var prettyBuf bytes.Buffer - prettyBuf.Grow(len(doc) + len(doc)/10) + originalLen := len(doc) + prettyBuf.Grow(originalLen + originalLen/10) if err := json.Indent(&prettyBuf, doc, "", " "); err == nil { bp.buffers[namespace].buffer.WriteByte('\n') + // Account for expansion + expansion := prettyBuf.Len() - originalLen + bp.currentUsage += int64(expansion) doc = prettyBuf.Bytes() } else { slog.Warn("Failed JSON indentation. Falling back to no-indent") } }
252-264: Potential filename collision under high concurrency.
time.Now().UnixNano()can return identical values for concurrent goroutines, leading to S3 key collisions and data loss.Add a random suffix using
math/rand/v2:func (bp *BatchProcessor) objectKey(namespace string) string { nsPath := strings.ReplaceAll(strings.Trim(namespace, "/"), ".", "/") if nsPath == "" { nsPath = "default" } - fileName := fmt.Sprintf("%d.json", time.Now().UnixNano()) + fileName := fmt.Sprintf("%d_%d.json", time.Now().UnixNano(), rand.Uint64()) prefix := strings.Trim(bp.config.Prefix, "/") if prefix == "" { return path.Join(nsPath, fileName) } return path.Join(prefix, nsPath, fileName) }
🧹 Nitpick comments (1)
connectors/s3/s3writer.go (1)
354-374: Consider per-namespace locking for better concurrency.The current
metadataMutexserializes all metadata updates across all namespaces. While this prevents races, updates to different namespaces could safely proceed in parallel.Consider using a per-namespace lock map:
// Add to BatchProcessor struct: metadataLocks sync.Map // map[string]*sync.Mutex per namespace // In updateMetadata: func (bp *BatchProcessor) updateMetadata(ctx context.Context, namespace string, fileKey string, recordCount uint64) error { // Get or create per-namespace mutex lockVal, _ := bp.metadataLocks.LoadOrStore(namespace, &sync.Mutex{}) nsLock := lockVal.(*sync.Mutex) nsLock.Lock() defer nsLock.Unlock() // ... rest of function unchanged }This allows concurrent metadata updates for different namespaces while still preventing races within the same namespace.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
connectors/s3/s3writer.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
connectors/s3/s3writer.go (3)
connectors/s3/connector.go (2)
MAX_S3_WRITE_TIMEOUT_SEC(39-39)S3_METADATA_UPDATES_RETRY(40-40)connectors/s3/metadata.go (1)
MetadataKey(9-20)connectors/s3/util.go (2)
IsS3OptimisticLockFailedError(5-7)ErrS3OptimisticLockFailed(10-12)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (4)
connectors/s3/s3writer.go (4)
1-65: LGTM! Clean structure and proper initialization.The use of
math/rand/v2(which auto-seeds) addresses previous concerns about deterministic jitter. The BatchProcessor struct correctly initializes the condition variable with its mutex.
77-83: LGTM! Prevents infinite loop scenario.The pre-check ensures that oversized batches are rejected immediately rather than entering an infinite wait loop.
294-321: LGTM! Optimistic locking correctly implemented.The conditional write logic properly handles both new file creation (using
IfNoneMatch="*"for put-if-not-exists) and updates (usingIfMatchwith ETag for compare-and-swap). This prevents concurrent metadata overwrites.
323-352: LGTM! Robust retry logic with proper backoff.The exponential backoff (
50*(1<<i)ms) ensures non-zero first retry, andrand.Nfrommath/rand/v2provides non-deterministic jitter. Correctly retries only on optimistic lock failures.
| if err != nil { | ||
| slog.Error("Upload failed", | ||
| "namespace", ns, | ||
| "key", key, | ||
| "error", err, | ||
| ) | ||
| //XXX: might want to implement a retry mechanism or Dead Letter Queue here. | ||
| } else { | ||
| // Update metadata on successful upload | ||
| if err := bp.updateMetadataWithRetries(ctx, ns, key, uint64(numDocs)); err != nil { | ||
| slog.Error("Failed to update metadata after retries", | ||
| "namespace", ns, | ||
| "key", key, | ||
| "error", err, | ||
| ) | ||
| } | ||
| } |
There was a problem hiding this comment.
Upload failures result in data loss.
Failed uploads are logged but the data is discarded. While the XXX comment (line 210) acknowledges this, transient S3 errors will cause permanent data loss.
Would you like me to generate a retry mechanism with exponential backoff for failed uploads, or a simple file-based dead letter queue for recovery?
1 participant
Summary by CodeRabbit
New Features
Chores
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.