storage: budget MySQL snapshot partition probing by table size

Bound the prefix partitioner to a per-table probe query budget scaled
to the estimated row count, 2500 requests per billion rows by default
via the mysql_source_snapshot_partition_requests_per_billion_rows
dyncfg, with a floor of 256 so small tables can afford their handful
of splits. Probes run sequentially, so this caps the wall-clock that
partitioning can add to snapshot start in proportion to the snapshot
work it optimizes. An exhausted budget stops splitting early and
leaves coarser buckets, never incorrect ones.

Adds a mock test asserting the partitioner never exceeds its budget
and still returns a valid ordered boundary list when truncated.

Author: peterdukelarsen

src/mysql-util/src/partition.rs

@@ -18,13 +18,18 @@ const TARGET_RANGES_PER_WORKER: f64 = 8.0;
 
 /// Computes up to `num_workers - 1` partition boundaries that divide the primary key space
 /// into `num_workers` roughly even partitions. This should be run in a transaction.
+///
+/// At most `max_requests` probes are issued against the server. When the
+/// budget runs out, remaining ranges stay unsplit, which skews partition
+/// sizes but never correctness.
 pub async fn partition_table(
     conn: &mut mysql_async::Conn,
     table: QualifiedTableRef<'_>,
     pk_col: &str,
     num_workers: usize,
     estimated_row_count: u64,
     min_rows_per_worker: u64,
+    max_requests: u64,
 ) -> Result<Vec<String>, MySqlError> {
     let (schema_name, table_name) = (table.schema_name, table.table_name);
     let mut db = KeyProber::new(conn, table, pk_col);
@@ -33,6 +38,7 @@ pub async fn partition_table(
         num_workers,
         estimated_row_count,
         min_rows_per_worker,
+        max_requests,
     )
     .await?;
     tracing::trace!(
@@ -61,6 +67,7 @@ async fn partition<D: PrimaryKeyProber>(
     workers: usize,
     estimated_row_count: u64,
     min_rows_per_worker: u64,
+    max_requests: u64,
 ) -> Result<Vec<String>, MySqlError> {
     if workers <= 1 {
         return Ok(Vec::new());
@@ -70,7 +77,13 @@ async fn partition<D: PrimaryKeyProber>(
     let target_max_rows_per_range =
         get_target_max_rows_per_range(workers, estimated_row_count, min_rows_per_worker);
     // Should be many more ranges than workers unless the overall row count of the table is quite small.
-    let ranges = split_into_ranges(db, estimated_row_count, target_max_rows_per_range).await?;
+    let ranges = split_into_ranges(
+        db,
+        estimated_row_count,
+        target_max_rows_per_range,
+        max_requests,
+    )
+    .await?;
     Ok(assign_boundaries(&ranges, workers))
 }
 
@@ -101,7 +114,9 @@ async fn split_into_ranges<D: PrimaryKeyProber>(
     db: &mut D,
     estimated_row_count: f64,
     target_rows_per_bucket: f64,
+    max_requests: u64,
 ) -> Result<Vec<Range>, MySqlError> {
+    let mut budget = max_requests;
     let mut ranges = vec![Range {
         prefix: None,
         end: None,
@@ -112,9 +127,12 @@ async fn split_into_ranges<D: PrimaryKeyProber>(
         let mut split_any = false;
         let mut next: Vec<Range> = Vec::with_capacity(ranges.len());
         for range in ranges {
-            if range.estimated_rows > target_rows_per_bucket {
+            // Entering a split costs one probe for the first prefix plus
+            // up to four for its first walk step, so a budget below that keeps
+            // the range as a leaf.
+            if range.estimated_rows > target_rows_per_bucket && budget >= 4 {
                 split_any = true;
-                next.extend(split_range(db, &range, target_rows_per_bucket).await?);
+                next.extend(split_range(db, &range, target_rows_per_bucket, &mut budget).await?);
             } else {
                 next.push(range);
             }
@@ -124,6 +142,11 @@ async fn split_into_ranges<D: PrimaryKeyProber>(
             break;
         }
     }
+    tracing::debug!(
+        ranges = ranges.len(),
+        requests_spent = max_requests - budget,
+        "split key space into ranges"
+    );
     Ok(ranges)
 }
 
@@ -160,14 +183,20 @@ fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec<String> {
 /// Splits `parent` at every distinct key prefix one character longer than the
 /// prefix `parent`, i.e. "a" depth: 1 for a table with keys "a", "aa", "aaa", "ab" would be split to
 /// "a" depth: 2 estimate 1, "aa" depth: 2 estimate 2, "ab" depth 2, estimate 1.
+///
+/// `budget` is decremented once per probe. The caller must provide at least
+/// 4, one for the first prefix and three for a walk step, and the walk closes
+/// out with a tail child once it cannot afford another step.
 async fn split_range<D: PrimaryKeyProber>(
     db: &mut D,
     parent: &Range,
     target_rows: f64,
+    budget: &mut u64,
 ) -> Result<Vec<Range>, MySqlError> {
     let depth = parent.depth + 1;
     let mut children = Vec::new();
 
+    *budget -= 1;
     let Some(mut cur) = db
         .prefix_of_first_key_in_range(parent.prefix.as_deref(), parent.end.as_deref(), depth)
         .await?
@@ -179,12 +208,15 @@ async fn split_range<D: PrimaryKeyProber>(
 
     loop {
         // Small optimization to return early if the remaining rows past the current start prefix
-        // fits within the threshold we're looking for.
+        // fits within the threshold we're looking for. An exhausted budget likewise closes out
+        // the split, continuing costs up to three more probes plus the next iteration's estimate before
+        // another chance to stop.
+        *budget -= 1;
         let remaining = db
             .estimate_range_rows(start.as_deref(), parent.end.as_deref())
             .await?;
         let remaining = f64::cast_lossy(remaining).max(1.0);
-        if remaining <= target_rows {
+        if remaining <= target_rows || *budget < 4 {
             children.push(Range {
                 prefix: start,
                 end: parent.end.clone(),
@@ -198,12 +230,14 @@ async fn split_range<D: PrimaryKeyProber>(
         // extensions still need visiting. The first key past `cur` exposes
         // them, so the walk can keep splitting instead of retrying the whole
         // range at greater depths forever.
+        *budget -= 1;
         let next = match db
             .prefix_of_first_row_not_matching_prefix(&cur, parent.end.as_deref(), depth)
             .await?
         {
             Some(next) => Some(next),
             None => {
+                *budget -= 1;
                 db.prefix_of_first_key_in_range(Some(&cur), parent.end.as_deref(), depth)
                     .await?
             }
@@ -212,6 +246,7 @@ async fn split_range<D: PrimaryKeyProber>(
         // server produces one), stop splitting here.
         match next.filter(|next| next != &cur) {
             Some(next) => {
+                *budget -= 1;
                 let estimated_rows = db
                     .estimate_range_rows(start.as_deref(), Some(&next))
                     .await?;
@@ -304,9 +339,15 @@ mod tests {
     /// "estimates". Byte order stands in for the collation.
     struct MockDb {
         keys: Vec<String>,
+        /// Probes served, for asserting on the request budget.
+        requests: u64,
     }
 
     impl MockDb {
+        fn new(keys: Vec<String>) -> Self {
+            MockDb { keys, requests: 0 }
+        }
+
         fn bounds(&self, start: Option<&str>, end: Option<&str>) -> (usize, usize) {
             // The lower bound is exclusive, a key equal to `start` is skipped.
             let lo = match start {
@@ -327,6 +368,7 @@ mod tests {
             start: Option<&str>,
             end: Option<&str>,
         ) -> Result<u64, MySqlError> {
+            self.requests += 1;
             let (lo, hi) = self.bounds(start, end);
             Ok(u64::cast_from(hi - lo))
         }
@@ -337,6 +379,7 @@ mod tests {
             end: Option<&str>,
             len: usize,
         ) -> Result<Option<String>, MySqlError> {
+            self.requests += 1;
             let (lo, hi) = self.bounds(start, end);
             if lo >= hi {
                 return Ok(None);
@@ -350,6 +393,7 @@ mod tests {
             end: Option<&str>,
             len: usize,
         ) -> Result<Option<String>, MySqlError> {
+            self.requests += 1;
             let (_, hi) = self.bounds(None, end);
             // Find the last key matching `cur`, byte prefixes stand in for
             // the collation's LIKE matching.
@@ -370,38 +414,36 @@ mod tests {
 
     #[mz_ore::test(tokio::test)]
     async fn single_worker_gets_no_boundaries() -> Result<(), MySqlError> {
-        let mut db = MockDb { keys: keys(1000) };
+        let mut db = MockDb::new(keys(1000));
         let count = u64::cast_from(db.keys.len());
-        let boundaries = partition(&mut db, 1, count, MIN_BUCKET_ROWS).await?;
+        let boundaries = partition(&mut db, 1, count, MIN_BUCKET_ROWS, u64::MAX).await?;
         assert!(boundaries.is_empty());
         Ok(())
     }
 
     #[mz_ore::test(tokio::test)]
     async fn small_table_gets_no_boundaries() -> Result<(), MySqlError> {
         // Under `min_bucket_rows` everything stays in one bucket.
-        let mut db = MockDb { keys: keys(10_000) };
+        let mut db = MockDb::new(keys(10_000));
         let count = u64::cast_from(db.keys.len());
-        let boundaries = partition(&mut db, 4, count, MIN_BUCKET_ROWS).await?;
+        let boundaries = partition(&mut db, 4, count, MIN_BUCKET_ROWS, u64::MAX).await?;
         assert!(boundaries.is_empty());
         Ok(())
     }
 
     #[mz_ore::test(tokio::test)]
     async fn empty_table_gets_no_boundaries() -> Result<(), MySqlError> {
-        let mut db = MockDb { keys: vec![] };
-        let boundaries = partition(&mut db, 4, 0, MIN_BUCKET_ROWS).await?;
+        let mut db = MockDb::new(vec![]);
+        let boundaries = partition(&mut db, 4, 0, MIN_BUCKET_ROWS, u64::MAX).await?;
         assert!(boundaries.is_empty());
         Ok(())
     }
 
     #[mz_ore::test(tokio::test)]
     async fn splits_evenly_across_workers() -> Result<(), MySqlError> {
-        let mut db = MockDb {
-            keys: keys(200_000),
-        };
+        let mut db = MockDb::new(keys(200_000));
         let count = u64::cast_from(db.keys.len());
-        let boundaries = partition(&mut db, 4, count, MIN_BUCKET_ROWS).await?;
+        let boundaries = partition(&mut db, 4, count, MIN_BUCKET_ROWS, u64::MAX).await?;
         assert_eq!(boundaries.len(), 3);
         // Boundaries must be sorted and split the keys into ~50k chunks.
         let mut prev = 0;
@@ -420,9 +462,9 @@ mod tests {
 
     #[mz_ore::test(tokio::test)]
     async fn low_min_bucket_rows_splits_small_tables() -> Result<(), MySqlError> {
-        let mut db = MockDb { keys: keys(1000) };
+        let mut db = MockDb::new(keys(1000));
         let count = u64::cast_from(db.keys.len());
-        let boundaries = partition(&mut db, 4, count, 10).await?;
+        let boundaries = partition(&mut db, 4, count, 10, u64::MAX).await?;
         assert_eq!(boundaries.len(), 3);
         let mut prev = 0;
         for b in &boundaries {
@@ -445,9 +487,9 @@ mod tests {
         // range unsplittable at any depth.
         let mut all_keys = vec!["U".to_string()];
         all_keys.extend((0..1000).map(|i| format!("U{i:06}")));
-        let mut db = MockDb { keys: all_keys };
+        let mut db = MockDb::new(all_keys);
         let count = u64::cast_from(db.keys.len());
-        let boundaries = partition(&mut db, 4, count, 10).await?;
+        let boundaries = partition(&mut db, 4, count, 10, u64::MAX).await?;
         assert_eq!(boundaries.len(), 3);
         for b in &boundaries {
             assert!(
@@ -458,6 +500,73 @@ mod tests {
         Ok(())
     }
 
+    #[mz_ore::test(tokio::test)]
+    async fn request_budget_bounds_probes() -> Result<(), MySqlError> {
+        // Unlimited budget as a baseline: splitting this table to the target
+        // costs far more probes than the budget below.
+        let mut db = MockDb::new(keys(200_000));
+        let count = u64::cast_from(db.keys.len());
+        partition(&mut db, 4, count, 10, u64::MAX).await?;
+        assert!(db.requests > 50, "baseline requests={}", db.requests);
+
+        // A small budget stops splitting early but still yields a valid,
+        // ordered boundary list.
+        let mut db = MockDb::new(keys(200_000));
+        let budget = 24;
+        let boundaries = partition(&mut db, 4, count, 10, budget).await?;
+        assert!(db.requests <= budget, "requests={}", db.requests);
+        assert!(
+            !boundaries.is_empty() && boundaries.len() <= 3,
+            "{boundaries:?}"
+        );
+        for pair in boundaries.windows(2) {
+            assert!(pair[0] < pair[1], "{boundaries:?}");
+        }
+        Ok(())
+    }
+
+    /// A database whose next-prefix wraps around instead of advancing, like a
+    /// numeric key column compared as a string. The request budget must still
+    /// bound the walk, and duplicate boundaries must never be emitted.
+    struct WrappingDb;
+
+    impl PrimaryKeyProber for WrappingDb {
+        async fn estimate_range_rows(
+            &mut self,
+            _: Option<&str>,
+            _: Option<&str>,
+        ) -> Result<u64, MySqlError> {
+            Ok(1_000_000)
+        }
+        async fn prefix_of_first_key_in_range(
+            &mut self,
+            _: Option<&str>,
+            _: Option<&str>,
+            _: usize,
+        ) -> Result<Option<String>, MySqlError> {
+            Ok(Some("9".to_string()))
+        }
+        async fn prefix_of_first_row_not_matching_prefix(
+            &mut self,
+            _: &str,
+            _: Option<&str>,
+            _: usize,
+        ) -> Result<Option<String>, MySqlError> {
+            // Never advances past "9".
+            Ok(Some("1".to_string()))
+        }
+    }
+
+    #[mz_ore::test(tokio::test)]
+    async fn non_advancing_prefixes_terminate() -> Result<(), MySqlError> {
+        // A finite budget is the only bound here: the wrapping server would
+        // otherwise feed the walk identical children forever. Boundaries may
+        // repeat, downstream monotonicity validation rejects them before use.
+        let boundaries = partition(&mut WrappingDb, 4, 1_000_000, MIN_BUCKET_ROWS, 100).await?;
+        assert!(boundaries.len() <= 3);
+        Ok(())
+    }
+
     /// Exercises the partitioner against a live MySQL server, covering what
     /// the mock cannot see: `EXPLAIN` estimates over prepared statements,
     /// `LIKE` pattern semantics, and the nested next-prefix query.
@@ -514,12 +623,14 @@ mod tests {
         };
 
         // A minimum above the table size yields no boundaries.
-        let boundaries = partition_table(&mut conn, table.clone(), "id", 4, 1004, 50_000).await?;
+        let boundaries =
+            partition_table(&mut conn, table.clone(), "id", 4, 1004, 50_000, u64::MAX).await?;
         assert!(boundaries.is_empty(), "{boundaries:?}");
 
         // A low minimum splits the table, and MySQL agrees the boundaries are
         // strictly increasing under the column collation.
-        let boundaries = partition_table(&mut conn, table.clone(), "id", 4, 1004, 10).await?;
+        let boundaries =
+            partition_table(&mut conn, table.clone(), "id", 4, 1004, 10, u64::MAX).await?;
         assert!(
             !boundaries.is_empty() && boundaries.len() <= 3,
             "{boundaries:?}"

src/storage-types/src/dyncfgs.rs

@@ -225,6 +225,18 @@ pub static MYSQL_SOURCE_SNAPSHOT_PARTITION_MIN_ROWS: Config<usize> = Config::new
      stops splitting; also the smallest table considered worth splitting.",
 );
 
+/// Probe query budget for the MySQL snapshot prefix partitioner, scaled to
+/// the table's estimated size so probing effort stays proportional to the
+/// snapshot work it optimizes. A small floor applies so modest tables can
+/// still afford their handful of splits. When a table's budget runs out,
+/// splitting stops early and buckets come out coarser, never incorrect.
+pub static MYSQL_SOURCE_SNAPSHOT_PARTITION_REQUESTS_PER_BILLION_ROWS: Config<usize> = Config::new(
+    "mysql_source_snapshot_partition_requests_per_billion_rows",
+    2_500,
+    "Cap on MySQL snapshot PK-prefix partitioning probe queries per table, per billion \
+     estimated rows; when exhausted, splitting stops early with coarser buckets.",
+);
+
 /// If the optimizer estimates the table has fewer rows than this, compute the exact row count
 /// with `COUNT(*)`. Otherwise, report the `information_schema` estimate directly.
 pub static MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS: Config<usize> = Config::new(
@@ -449,6 +461,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS)
         .add(&MYSQL_SOURCE_SNAPSHOT_PARALLELISM)
         .add(&MYSQL_SOURCE_SNAPSHOT_PARTITION_MIN_ROWS)
+        .add(&MYSQL_SOURCE_SNAPSHOT_PARTITION_REQUESTS_PER_BILLION_ROWS)
         .add(&ORE_OVERFLOWING_BEHAVIOR)
         .add(&PG_FETCH_SLOT_RESUME_LSN_INTERVAL)
         .add(&PG_SCHEMA_VALIDATION_INTERVAL)