feat(mongodb-to-mongodb]: Read/Write Performance and Reliability Improvements

[michaeltle-goog]

Overview

This PR refactors the mongodb-to-mongodb Dataflow template to improve read-phase scalability, write concurrency, and target cluster protection. It replaces legacy splitting operations with a dynamic, type-aware bounds generator and decouples read/write execution in high-throughput pipelines.

Performance: Migrated 2.5M documents in ~12 minutes.

Key Changes

1. Type-Aware Data-Driven Read Splitting (ReadSplitGenerator)

Replaced splitVector and autoBucket with a deterministic, data-driven splitting algorithm that safely handles single-type and mixed-type _id collections:

• Active Type Discovery: Scans the source collection (find().limit(1)) to identify which BSON _id types are present (NUMBER, STRING, OBJECT_ID, OTHER).
• Per-Type Sampling ($sample): Executes an aggregation pipeline independently for each active BSON type:

[
      { "$match": { "_id": { "$type": "<type>" } } },
      { "$sample": { "size": sampleSize } },
      { "$project": { "_id": 1 } },
      { "$sort": { "_id": 1 } }
]

(Sample size scales dynamically as max(1000, numSplits * 64)).

• Quantile Boundary Extraction: Divides the sorted sample keys into numSplits equal buckets to extract quantile boundaries (25th, 50th, 75th percentiles).
• Type-Scoped Range Filtering: Wraps every generated range query in an explicit BSON $type guard (e.g., _id >= min AND _id < max AND $type: "string"). This prevents MongoDB from silently dropping documents during cross-type range evaluations.
• Multi-Type Merging: Combines the type branches for each split index using an $or wrapper. If only one BSON type exists in the collection, the $or wrapper is automatically omitted.
• Deterministic Fallback: If a collection has too few documents for sampling or duplicate quantile boundaries are detected, it falls back to uniform partition splitting per type ($mod for numbers, hex prefixes for strings, timestamp prefixes for ObjectIDs).

2. Non-Blocking Asynchronous Writes (MongoDbTransforms)

• Decoupled Read and Write Loops: Refactored write execution so that MongoDB batch futures (CompletableFuture.runAsync) no longer block upstream reader threads.
• Bounded Concurrency: Capped parallel write tasks per worker (maxConcurrentAsyncWrites) to pipeline document reads without causing memory exhaustion or GC pressure.

3. Write-Path Rate Limiting

• Linear Ramp-Up Throttling: Added a configurable rate limiter (updateRateLimiterIfNeeded()) that linearly increases write throughput per worker from initialWriteRatePerWorker to maxWriteRatePerWorker over writeRateRampUpMinutes.

4. Logging

• Added more info logging to help customers understand/parse the Dataflow pipeline

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 significantly enhances the scalability and reliability of the MongoDB-to-MongoDB Dataflow template. By replacing brittle legacy splitting mechanisms with a robust, type-aware approach and introducing asynchronous, rate-limited write paths, the template is now better equipped to handle high-volume data migrations while maintaining stability and performance.

Highlights

• Type-Aware Split Reader: Replaced legacy splitVector/autoBucket with a dynamic, $sample-based reader that isolates BSON types to ensure deterministic query bounds.
• Asynchronous Write Paths: Refactored write loops to use asynchronous execution, preventing blocking of upstream readers and improving overall throughput.
• Intelligent Write Rate-Limiting: Introduced a configurable rate limiter with linear ramp-up support to prevent connection bursts and protect target clusters.
• Logging and Maintenance: Reduced log verbosity, added URI credential masking, and improved code style and test coverage.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces parallel index-slice reading and linear write rate limiting to the MongoDB-to-MongoDB Dataflow template. Parallel reading is supported via a new ReadSplitGenerator utility that partitions collections based on key types, while write rate limiting is implemented in MongoDbTransforms using a linear ramp-up strategy. Additionally, a UriSanitizer utility is added to mask sensitive credentials in logs. Feedback was provided on ReadSplitGenerator to use a try-with-resources block when iterating over aggregate results to prevent potential MongoDB cursor leaks.

[gemini-code-assist]

Iterating over col.aggregate(pipeline) directly with a for-each loop can leak the underlying MongoDB cursor if an exception is thrown during iteration. Using a try-with-resources block with MongoCursor ensures that the cursor is always closed properly.

    List<BsonValue> sampledKeys = new ArrayList<>();
    try (com.mongodb.client.MongoCursor<BsonDocument> cursor = col.aggregate(pipeline).iterator()) {
      while (cursor.hasNext()) {
        BsonDocument doc = cursor.next();
        if (doc.containsKey("_id")) {
          sampledKeys.add(doc.get("_id"));
        }
      }
    }

[codecov]

Codecov Report

❌ Patch coverage is 57.62004% with 203 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.19%. Comparing base (66dfb9c) to head (1e3d277).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
.../cloud/teleport/v2/templates/MongoDbToMongoDb.java	0.00%	82 Missing ⚠️
...loud/teleport/v2/transforms/MongoDbTransforms.java	63.74%	40 Missing and 22 partials ⚠️
...oud/teleport/v2/transforms/ReadSplitGenerator.java	72.93%	31 Missing and 28 partials ⚠️

❌ Your patch check has failed because the patch coverage (57.62%) 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    #4079    +/-   ##
==========================================
  Coverage     56.18%   56.19%            
+ Complexity     7342     6932   -410     
==========================================
  Files          1126     1128     +2     
  Lines         68766    69197   +431     
  Branches       7785     7871    +86     
==========================================
+ Hits          38637    38885   +248     
- Misses        27635    27767   +132     
- Partials       2494     2545    +51     

Components	Coverage Δ
spanner-templates	87.53% <ø> (-0.01%)	⬇️
spanner-import-export	68.91% <ø> (ø)
spanner-live-forward-migration	89.29% <ø> (-0.02%)	⬇️
spanner-live-reverse-replication	83.48% <ø> (-0.02%)	⬇️
spanner-bulk-migration	92.14% <ø> (-0.01%)	⬇️
gcs-spanner-dv	88.58% <ø> (-0.02%)	⬇️

Files with missing lines	Coverage Δ
...gle/cloud/teleport/v2/transforms/UriSanitizer.java	100.00% <100.00%> (ø)
...oud/teleport/v2/transforms/ReadSplitGenerator.java	72.93% <72.93%> (ø)
...loud/teleport/v2/transforms/MongoDbTransforms.java	64.79% <63.74%> (-1.78%)	⬇️
.../cloud/teleport/v2/templates/MongoDbToMongoDb.java	0.00% <0.00%> (ø)

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.