Repository navigation
perf(mongodb-to-mongodb): single-stream CDC optimization, stateful deduplication, and timestamp-aware coalescing - #4298
Conversation
… order and O(1) covered seeks
…nd background CDC cursor prefetcher
…and timestamp-aware coalescing - Enforce stateful deduplication whenever CDC mutations are enabled (standalone or combined with backfill) to drop stale out-of-order mutations across change streams and retries. - Add timestamp sort key comparison in WriteBatchesFn batch coalescing so older events within an unordered batch do not overwrite newer replacements or deletes. - Add unit tests for stale insert after delete and newer insert after delete coalescing.
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 delivers comprehensive performance and correctness enhancements for the mongodb-to-mongodb Dataflow Flex Template. By introducing robust stateful deduplication, timestamp-aware batch coalescing, and optimized binary serialization, the changes significantly improve the reliability and efficiency of streaming CDC pipelines. Additionally, the inclusion of a background prefetcher for change stream cursors helps mitigate latency bottlenecks. 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 introduces several performance optimizations and robustness improvements to the MongoDB-to-MongoDB pipeline. Key changes include caching deduplication keys, migrating the DocumentWithMetadataCoder to deterministic binary BSON serialization, implementing a background prefetcher thread to pipeline change stream events across SDF boundaries, and optimizing type-isolated split generation with O(1) covered index scans. Additionally, bulk write coalescing now respects event timestamps to prevent stale updates. The reviewer feedback suggests closing the MongoDB cursor before joining the prefetcher thread to prevent blocking on network calls, uniquely naming prefetcher threads using an atomic counter for better debugging, and catching ArithmeticException when converting Decimal128 values to avoid potential crashes on NaN or infinite values.
Codecov Report❌ Patch coverage is ❌ Your patch check has failed because the patch coverage (68.40%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #4298 +/- ##
============================================
+ Coverage 56.35% 57.13% +0.77%
+ Complexity 7858 7722 -136
============================================
Files 1154 1164 +10
Lines 73194 74592 +1398
Branches 8580 8848 +268
============================================
+ Hits 41252 42621 +1369
+ Misses 29158 29100 -58
- Partials 2784 2871 +87
🚀 New features to boost your workflow:
|
…dress review feedback - Initialize PartitionCursorHolder.latestPostBatchToken to null and clear it whenever a non-null event is prefetched so a coalescing poll timeout during an in-flight tryNext() does not rewind lastSeenResumeToken to an older token. - Add unique thread ID suffix to mongodb-cdc-prefetcher threads for clearer thread dumps. - Return null from ReadSplitGenerator.asBigDecimal for non-finite Double/Decimal128 bounds (NaN, +/-Infinity) so generateProbedNumberSplits falls back to the unsplit $type filter without excluding NaN or mis-interpolating Infinity, while mapping Decimal128 negative zero (-0) to BigDecimal.ZERO. - Update dedupStateRetentionHours helpText and apply Spotless formatting in MongoDbToMongoDb. - Add unit tests covering prefetcher coalescing timeout, token mismatch cursor rebuilds, out-of-order stale delete coalescing, connection string pool/socket timeouts, and non-finite/negative-zero numeric split bounds.
…SON split edge cases
Ensure consistent BSON encoding, exact change stream checkpoint tokens, and safe split generation across edge-case BSON types:
- Configure a shared Document codec with UuidRepresentation.STANDARD in DocumentWithMetadata, DocumentWithMetadataCoder, and ApplyUdfFn so java.util.UUID values round-trip through binary BSON coding and Extended JSON deduplication keys without CodecConfigurationException.
- Store BsonDocument directly in ChangeStreamRestriction with lazy JSON serialization, eliminating RESUME_TOKEN_SERIALIZATION_STRIDE so every event claims its exact resume token without per-event JSON serialization overhead.
- Synchronize PartitionCursorHolder.getLatestPostBatchResumeToken() and guard against non-empty eventQueue, clear lastTryNextWasNull on fresh poll requests, and mark running=false on prefetcher exceptions.
- Guard {$unset: "updateDescription"} in buildChangeStreamPipeline when fullDocumentStrategy is "default".
- Handle non-finite numeric bounds before $sample quantile splitting, support shared BSON sort rank between symbol and string in probeActiveTypeBounds, deduplicate rounded numeric split boundaries, and omit empty filters in generateIndexSliceFilters when numSplits exceeds 62 for string-only keys.
- Return null from DocumentWithMetadata.getId() for empty or malformed documentKey strings, validate 32-bit epoch-second bounds in parseStartAtOperationTime, and avoid duplicate pre-UDF JSON serialization in ApplyUdfFn.
Change the default migrationMode parameter and fallback from BACKFILL_AND_STREAMING to BACKFILL so one-time historical backfill runs by default unless streaming CDC is explicitly requested.
…support UUID in backfill checkpointing Prevent DoFn @teardown in MongoDbBackfillReader and MongoDbChangeStreamReader from closing JVM-wide shared MongoClient instances cached in MongoDbTransforms.CLIENT_CACHE, and serialize backfill checkpoint _id values using DocumentWithMetadata.DOCUMENT_CODEC so documents with UUID _ids or nested UUID fields do not throw CodecConfigurationException.
…oDbToMongoDbIT) and fix DLQ/composite _id bugs
…bToMongoDbIT and apply Java readability simplifications
Apache Beam's PTransform.getName() returns the default class name unless a name is passed to the PTransform(String) constructor, causing UDF, validation, and write WriteToDlq stages in bounded pipelines to share the same output prefix and overwrite earlier stage DLQ files with empty shards.
Pass explicit stage names ('udf', 'validate', 'write') to WriteToDlq so each pipeline stage writes to distinct error-<stage> output files.
…Helper. Embedding heterogeneous BSON fixtures, endpoint provisioning, launch configuration, DLQ artifact inspection, and deep BSON parity checks directly inside MongoDbToMongoDbIT obscured the end-to-end test scenarios beneath setup boilerplate. Move all heterogeneous BSON, CDC, UDF, and DLQ document fixtures into immutable static RawBsonDocument constants in MongoDbTestHelper, and encapsulate MongoDB/Firestore Enterprise provisioning, LaunchConfig construction, UDF/DLQ staging, and BSON parity verification inside TestEnvironment.
There was a problem hiding this comment.
Code Review
This pull request introduces significant performance optimizations, robustness improvements, and feature enhancements to the MongoDB-to-MongoDB template. Key updates include rewriting DocumentWithMetadataCoder to use binary BSON serialization, implementing a background prefetcher thread for change streams to pipeline network RPCs, optimizing split generation via O(1) covered index seeks, and ensuring strict BSON fidelity (such as standard UUID representation and NaN/Infinity handling). Additionally, the default migration mode is changed to BACKFILL, connection pool/socket timeouts are made configurable, and extensive integration tests are added. The reviewer feedback suggests wrapping database operations within the ConditionCheck in MongoDbTestHelper in a try-catch block to gracefully handle transient errors and prevent immediate test failures.
| @Override | ||
| protected CheckResult check() { | ||
| long count = tgtStreamCol.countDocuments(); | ||
| Document d1 = tgtStreamCol.find(Filters.eq("_id", "post_t0_valid_1")).first(); | ||
| Document d2 = tgtStreamCol.find(Filters.eq("_id", "post_t0_valid_2")).first(); | ||
| boolean ready = count == 2 && d1 != null && d2 != null; | ||
| return new CheckResult(ready, "count=" + count); | ||
| } |
There was a problem hiding this comment.
In integration tests, database operations within a ConditionCheck should be wrapped in a try-catch block that returns a failed CheckResult (false) when an exception is caught. This allows transient errors to be handled gracefully by retrying the check rather than failing the test immediately.
@Override
protected CheckResult check() {
try {
long count = tgtStreamCol.countDocuments();
Document d1 = tgtStreamCol.find(Filters.eq("_id", "post_t0_valid_1")).first();
Document d2 = tgtStreamCol.find(Filters.eq("_id", "post_t0_valid_2")).first();
boolean ready = count == 2 && d1 != null && d2 != null;
return new CheckResult(ready, "count=" + count);
} catch (Exception e) {
return new CheckResult(false, "Exception during check: " + e.getMessage());
}
}References
- In integration tests, database operations within a ConditionCheck should be wrapped in a try-catch block that returns a failed CheckResult (false) when an exception is caught. This allows transient errors to be handled gracefully by retrying the check rather than failing the test immediately.
…c writer pool. Running bounded backfill pipelines on DataflowRunner previously defaulted to Legacy Batch execution, which replaced GroupIntoBatches with a full-dataset disk-shuffle barrier and skewed shard keys onto a subset of workers. Additionally, WriteBatchesFn created a separate MongoClient connection pool per DoFn instance and executed bulkWrite batches synchronously on the calling thread, causing initial TLS handshake storms and multi-minute bundle processing stalls. Enable Streaming Engine execution when running on DataflowRunner so bounded backfill reads and sharded batched writes pipeline concurrently across all workers while still terminating automatically upon completion. Share a single JVM-wide MongoClient and bounded async write executor pool across WriteBatchesFn instances on each worker, preserving per-document ordering across overlapping batches within a bundle and reducing initial retry backoff from 2s to 200ms.
Large ObjectId collections experience severe partition size skew under WiredTiger's random cursor $sample stage, which under-samples shallower B-tree branches near the end of high-rate inserts and leaves multi-million-document tail partitions on a single worker. In addition, MongoDbBackfillReader previously processed all assigned partitions within a concurrency slot sequentially without supporting non-zero trySplit requests, preventing idle workers from stealing unstarted partitions from slower slots, and checkpointed every 2,000 documents with 2,000-document getMore batches. Weight adjacent sampled ObjectId intervals by their 24-bit monotonic counter delta, refine wide multi-second gaps via bounded O(1) covered index seeks at 1-second boundaries, and interpolate partition boundaries while preserving the 5-byte processUnique identifier (raising the $sample cap from 5,000 to 50,000). Track an exclusive upper partition bound on BackfillSlotRestriction so BackfillSlotRestrictionTracker.trySplit can hand off remaining unstarted whole partitions to idle workers, raise default backfill cursor batch size to 5,000 and slice size to 25,000, and expose --backfillCursorBatchSize and --backfillMaxDocsPerSlice for document-size-specific tuning.
Per-DoFn Guava RateLimiter instances previously throttled each deserialized WriteBatchesFn thread independently rather than bounding aggregate worker throughput, continued enforcing a fixed ceiling after target database tablets finished splitting, and started the ramp-up clock during @setup before the first buffered batch arrived. Share a single WorkerRateLimiter instance per worker JVM per WriteWithDlq transform, lazily anchor the warm-up clock on the first acquire() call, and support uncapped exponential warm-up (when initialWriteRatePerWorker > 0 and maxWriteRatePerWorker <= 0) that doubles per minute over writeRateRampUpMinutes before disabling itself for lock-free steady-state writes.
Description
This PR provides performance, scalability, and correctness improvements for the
mongodb-to-mongodbDataflow Flex Template across batchBACKFILL, combinedBACKFILL_AND_STREAMING, and high-throughputSTREAMING_CDCworkloads, along with automated End-to-End (E2E) correctness integration tests (MongoDbToMongoDbIT+MongoDbTestHelper) migrating from MongoDB to Cloud Firestore Enterprise (MongoDB compatibility mode):1. Default
migrationModeSet toBACKFILLmigrationModeinMongoDbToMongoDb.OptionsfromBACKFILL_AND_STREAMINGtoBACKFILLso bounded batch migrations are the default unless continuous change stream replication (BACKFILL_AND_STREAMINGorSTREAMING_CDC) is explicitly configured.2. Fast Canonical BSON Type Discovery & Split Generation (
ReadSplitGenerator)O(1)covered_idindex probes across MongoDB's canonical BSON comparison order (Number,String,Symbol,Document,BinData,ObjectId,Boolean,Date,Timestamp,Regex,MinKey/MaxKey).Dates (BsonDateTime(Long.MIN_VALUE)),Double.NaN/±Infinity/-0.0numeric bounds,Decimal128full precision,BsonTimestamp(0, 1), and$typeprobes for unboundedstringandbinDatabuckets.3. Background CDC Cursor Prefetching & Per-Event SDF Token Safety (
MongoDbChangeStreamReader)PartitionCursorHolderwith a dedicated background prefetcher thread per change stream partition that drainscursor.tryNext()asynchronously into a boundedArrayBlockingQueue(PREFETCH_QUEUE_CAPACITY = 4096), pipelininggetMoreRPCs across Beam Splittable DoFn (SDF) slice boundaries.cursor.getResumeToken()access on the prefetcher thread viagetLatestPostBatchResumeToken()and ensured exact per-event resume token serialization beforetracker.tryClaim(...)so SDF checkpoints never rewind or skip events mid-slice.$unset: "updateDescription"to the change stream aggregation pipeline whenfullDocumentpost-images are used (updateLookup,whenAvailable,required), avoiding redundant network transfer.4. Stateful Deduplication & Sharded
documentKeyNormalization (StatefulDeduplication,DocumentWithMetadata)StatefulDeduplicationwheneverincludeCdcis true (bothBACKFILL_AND_STREAMINGandSTREAMING_CDC) to enforce Last-Write-Wins (LWW) monotonic ordering per document across partitioned change streams, worker retries, andT0backfill/CDC overlap.documentKeys (which include shard key fields alongside_id) to canonical{"_id": ...}keys viaDocumentWithMetadata.serializeIdDocument(...)incomputeDedupKey()andextractIdFromDocumentKey().5. Timestamp-Aware Batch Write Coalescing, Composite
_idKey Safety, & SharedMongoClientLifecycle (MongoDbTransforms)WriteBatchesFn.flushBatch()to compareTimestampSortKeywhen coalescing multiple mutations for the same_idwithin a batch, preventing stale inserts/updates from overwriting newer updates or deletes.WriteBatchesFn.putCoalescedOp()to keycolMapbyitem.getDedupKey()(Canonical Extended JSON) when_idis aMap/Documentorbyte[], preserving MongoDB subdocument_idfield order so composite keys with different field orders (e.g.{"tenant": "acme", "seq": 1}vs{"seq": 1, "tenant": "acme"}) are not coalesced together.MongoDbTransforms.closeIfNotShared(MongoClient)and updated@TeardowninMongoDbBackfillReaderandMongoDbChangeStreamReaderso worker teardown of individualDoFninstances never closes the JVM-wide sharedMongoClientinCLIENT_CACHE.6. Binary BSON Coder, Lazy
originalDocument,UUIDCodec Hardening, & Stage-Isolated Bounded DLQ (DocumentWithMetadataCoder,MongoDbBackfillReader,MongoDbTransforms)DocumentWithMetadataCoderwith raw binary BSON wire-format encoding (RawBsonDocumentCodec) and lazyoriginalDocumentcomputation (preserving pre-UDF payloads acrosswithDocument(...)transformations for DLQ reporting).DocumentWithMetadata.DOCUMENT_CODEC(UuidRepresentation.STANDARD) acrossDocumentWithMetadata,DocumentWithMetadataCoder, andMongoDbBackfillReadercheckpoint_idextraction, preventingCodecConfigurationExceptionon documents withUUID_ids or nestedUUIDfields.@DoFn.BoundedPerElementtoMongoDbBackfillReader.ProcessBackfillSlotFn, passed explicit stage prefixes (udf,validate,write) toWriteToDlq(/error-<stage>), and clamped near-GlobalWindow.maxTimestamp()elements so batch DLQ writes fromGroupIntoBatchesare never dropped or overwritten across stages.7. Automated End-to-End Integration Tests: MongoDB -> Cloud Firestore Enterprise (
MongoDbToMongoDbIT,MongoDbTestHelper)MongoDbToMongoDbITandMongoDbTestHelpermigrating from an isolatedmongo:7.0replica-set source Testcontainer (MongoDBResourceManager) to an ephemeral Cloud Firestore Enterprise database with MongoDB compatibility enabled (FirestoreAdminClient) per test (with optional-DsourceMongoUri,-DtargetMongoUri, and-DuseTestcontainersTarget=trueoverrides):backfill_heterogeneousBsonTypesAndMultiCollection_preservesAllDocuments: DefaultBACKFILLmode, database-wide multi-collection auto-discovery, fine-grained_idsplit probing (targetBackfillChunkSize=3), and byte/type-exactRawBsonDocumentverification acrossObjectId,String,Int32,Int64(> 2^53),Double(±Infinity),Binary(0x00,0x04UUID,0x80,0xFF), and compositeDocument_ids plusDecimal128,Date,Timestamp,Boolean, IEEE 754-0.0/NaN,Regex,Null,MinKey/MaxKey,Array, and nestedDocumentpayloads.backfillAndStreaming_concurrentOverlapMutations_deduplicatesAndConverges:BACKFILL_AND_STREAMINGwithchangeStreamFullDocument=whenAvailable, testing concurrent updates, deletes (includingUUID_id), delete-then-reinsert resurrections, and net-new inserts during backfill/CDC overlap.streamingCdc_rapidSameKeyMutationsAndMultiPartition_coalescesAndConverges:STREAMING_CDCwithnumChangeStreamSplits=4andchangeStreamFullDocument=updateLookup, exercising rapid same-key mutation bursts (INSERT -> UPDATE -> REPLACE -> UPDATE,INSERT -> UPDATE -> DELETE,INSERT -> DELETE -> INSERTresurrection,UUID_iddelete), a >64-event burst, and post-idle prefetcher wake-up.streamingCdcAndBackfill_udfStartAtOperationTimeAndDlq_routesValidAndFailedDocs:startAtOperationTimeoplog replay filtering, JavaScript UDF transformation andnullfiltering, UDF exceptions routed toretryable/DLQ with lazyoriginalDocumentpreserved, and target permanent write rejections (__fs_invalid__reserved field /$jsonSchema) routed topermanent/DLQ.Testing
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 mvn spotless:check checkstyle:check test -pl v2/mongodb-to-mongodb: 179/179 tests passing (0 failures).MongoDbToMongoDbIT):Dataflow Templates Integration Smoke Tests(DirectRunner-> Cloud Firestore Enterprise: 2/2 passing in 52.3s) andDataflow Templates Integration Tests(Dataflow Flex Template -> Cloud Firestore Enterprise: 2/2 passing in 569.7s): 4/4 suites passing (0 failures).