fix(mongodb-to-mongodb): optimize type discovery using canonical BSON order and O(1) covered seeks

[michaeltle-goog]

Why This Change Was Needed

When running Dataflow backfill migrations from a sharded MongoDB cluster to Ignite, ReadSplitGenerator timed out during _id BSON type discovery (probeActiveTypeBounds).

1. Type Discovery (probeActiveTypeBounds)

Before (legacyProbeBuckets)

• How it worked: Spawned a 6-thread pool per collection (Executors.newFixedThreadPool(6)) that fired 6 concurrent $type queries across all candidate type buckets (number, string, objectId, binData, object, date), each with .sort({_id: 1}).limit(1) and a 3-second maxTimeMS (followed by a second .sort({_id: -1}).limit(1) query for matching buckets).
• Why it failed on sharded clusters:
  1. Connection Pool Cold-Start Burst: The short-lived metadata MongoClient is initialized with minSize(0) (0 warm connections). Firing 6 threads concurrently forced 6 simultaneous TCP + TLS + auth handshakes to mongos.
  2. Scatter-Gather Fanout: Because {_id: {$type: ...}} is not targeted to a single shard key value, mongos fanned out 6 × N_shards concurrent cursor requests across every shard.
  3. Cascading Fallback Penalty: When the 3s/5s timeouts tripped, probeActiveTypeBounds returned empty [], causing generateTypeIsolatedSplits to abort before running $sample or ObjectId interpolation and instead fall back to generateIndexSliceFilters (which wraps 4 types in a $or query with $mod and a hardcoded 5-year ObjectId window).

After (Canonical BSON Order + O(1) Covered Index Seeks)

• How it works:
  1. Aligns KNOWN_TYPE_BUCKETS with MongoDB's canonical BSON comparison order in the {_id: 1} B-tree index (number < string < object < binData < objectId < bool < date).
  2. Executes 2 sequential, unfiltered O(1) index endpoint seeks on a single connection:
     • Global Min: col.find().projection({_id: 1}).sort({_id: 1}).limit(1) -> minVal
     • Global Max: col.find().projection({_id: 1}).sort({_id: -1}).limit(1) -> maxVal
  3. Homogeneous Short-Circuit (99%+ of collections): If getBucketForValue(minVal) and getBucketForValue(maxVal) belong to the same BSON type bucket, by B-tree ordering no other BSON type can exist between them. The method immediately returns ProbedTypeBounds(minBucket, minVal, maxVal) in 2 queries (~2ms) with zero $type queries sent to the cluster.
  4. Mixed-Type Collections: Only probes candidate buckets strictly between minIdx and maxIdx in canonical BSON order.

2. Split Generation Hierarchy in the New System

The new system uses a multi-tiered fallback hierarchy across both Type Discovery (Phase 1) and Split Generation (Phase 2) so that even if a downstream step fails or times out, it degrades gracefully using information already gathered:

Phase 0: Small Collection Short-Circuit

• Before probing any types, generateTypeIsolatedSplits checks col.estimatedDocumentCount().
• If estimatedDocs <= 5,000 (or effectiveSplits <= 1), it immediately returns a single unfiltered split [{}] with zero index probes and zero sampling.

Phase 1 Fallbacks: Type & Boundary Discovery (probeActiveTypeBounds)

1. Tier 1 (Primary — Homogeneous O(1) Short-Circuit):
   • Runs 2 unfiltered covered seeks (sort({_id: 1}).limit(1) and sort({_id: -1}).limit(1)).
   • If minBucket == maxBucket, returns [ProbedTypeBounds(minBucket, minVal, maxVal)] immediately.
2. Tier 2 (Fallback 1A — Narrowed Mixed-Type Window):
   • Triggered when minVal and maxVal belong to different known BSON buckets (minIdx < maxIdx).
   • Only probes candidate buckets strictly between minIdx and maxIdx in canonical BSON order (for (i = minIdx + 1; i < maxIdx; i++)), using an unsorted limit(1) existence check first and only running sort({_id: ±1}).limit(1) if that intermediate type actually exists.
3. Tier 3 (Fallback 1B — legacyProbeBuckets):
   • Triggered only if minVal or maxVal is an unrecognized/exotic BSON type (e.g., RegEx, Timestamp) or if an unexpected exception occurs during the global min/max seek.
   • Falls back to parallel per-bucket $type probes.

Phase 2 Fallbacks: Split Generation (generateSplitsForTypeBounds)

Once Phase 1 yields ProbedTypeBounds(bucket, minKey, maxKey), each active type is partitioned using the following cascade:

1. Tier 1 (Primary — $sample Quantile Boundaries):
   • Runs [$sample, $project: {_id: 1}, $sort: {_id: 1}] (30s timeout) to compute data-skew-aware quantile boundaries ($lte: b_0, $gt: b_{i-1}, $lte: b_i, ..., $gt: b_{last}).
2. Tier 2 (Fallback 2A — Probed Boundary Interpolation / Type-Specific Splits):
   • Triggered if $sample times out on a large sharded cluster or returns fewer samples than splits:
     • objectId: Calls generateProbedObjectIdSplits(minHex, maxHex, splits) using the exact minKey and maxKey captured in Phase 1 to linearly interpolate splits contiguous 12-byte hex ranges (step = (maxBig - minBig) / splits), leaving the final slice open ($gte: lowHex).
     • number: Calls generateNumberFilters(splits) ($mod: [splits, r]).
     • string: Calls generateStringFilters(splits) (alphanumeric prefix ranges 0-9A-Za-z).
     • object / binData / bool / date: Returns a single type-isolated filter [{"_id": {"$type": "<bucket>"}}].
3. Tier 3 (Fallback 2B — Last-Resort Static $or Splits):
   • Triggered only if both Phase 1 Tier 1 and Tier 3 (legacyProbeBuckets) return zero active types (e.g., cluster unreachable). Calls generateIndexSliceFilters(effectiveSplits).

---

3. Final Split Upper Bound (maxKey)

Before

• The final partition (i == boundaryCount or i == numSplits - 1 in generateProbedObjectIdSplits, and single-split filters) appended "$lte": maxKey to _id.
• Any new documents inserted at the tail of the _id index after probeActiveTypeBounds ran were excluded from the backfill.

After

• maxKey is still used to compute evenly spaced interpolation step sizes (step = (maxBig - minBig) / numSplits), but "$lte": maxKey is omitted from the final split and single-split filters ({"_id": {"$type": "objectId", "$gte": {"$oid": lowHex}}}).
• This leaves the top partition open-ended within its BSON type so full backfills capture newly inserted tail documents.

[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 optimizes the MongoDB-to-MongoDB migration process by replacing inefficient, parallelized type discovery with a high-performance, O(1) covered index seek strategy. By leveraging the canonical BSON order of the _id index, the system can now determine collection boundaries with minimal overhead, particularly for homogeneous collections. Additionally, the change improves data completeness by removing restrictive upper bounds on splits, ensuring that newly inserted documents are included in the migration.

Highlights

• Optimized Type Discovery: Replaced parallel multi-threaded $type queries with a sequential, O(1) covered index seek strategy that leverages canonical BSON order to identify collection boundaries significantly faster.
• Homogeneous Collection Short-Circuit: Implemented a fast-path for homogeneous collections that resolves bounds in ~2ms with only two queries, bypassing expensive type-specific probing.
• Open-Ended Split Generation: Removed the hard $lte upper bound constraint from split generation, allowing the migration process to capture new documents inserted at the tail of the collection.
• Multi-Tiered Fallback Hierarchy: Introduced a robust fallback mechanism that degrades from O(1) seeks to mixed-type window probing and finally to legacy parallel probes if necessary.

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 refactors the BSON type boundary probing in ReadSplitGenerator to optimize performance. It replaces the parallel executor-based probing with a sequential global min/max check, allowing immediate resolution for homogeneous collections and targeted probing for mixed collections. Additionally, it adds support for boolean types, removes upper-bound constraints on tail slices, and updates unit tests. The feedback suggests further optimizing the intermediate candidate buckets loop by merging the existence check and the minimum bound query into a single sorted query to eliminate a redundant roundtrip.

[codecov]

Codecov Report

❌ Patch coverage is 47.76119% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.34%. Comparing base (6e68e20) to head (5078af2).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...oud/teleport/v2/transforms/ReadSplitGenerator.java	47.76%	29 Missing and 6 partials ⚠️

❌ Your patch check has failed because the patch coverage (47.76%) 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    #4270      +/-   ##
============================================
- Coverage     56.35%   56.34%   -0.02%     
+ Complexity     7858     7408     -450     
============================================
  Files          1154     1154              
  Lines         73194    73209      +15     
  Branches       8580     8585       +5     
============================================
- Hits          41252    41251       -1     
- Misses        29158    29160       +2     
- Partials       2784     2798      +14     

Components	Coverage Δ
spanner-templates	84.47% <ø> (-0.02%)	⬇️
spanner-import-export	69.00% <ø> (-0.06%)	⬇️
spanner-live-forward-migration	88.94% <ø> (-0.04%)	⬇️
spanner-live-reverse-replication	80.84% <ø> (-0.03%)	⬇️
spanner-bulk-migration	88.98% <ø> (-0.02%)	⬇️
gcs-spanner-dv	88.06% <ø> (-0.04%)	⬇️

Files with missing lines	Coverage Δ
...oud/teleport/v2/transforms/ReadSplitGenerator.java	49.22% <47.76%> (+1.24%)	⬆️

... and 6 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.