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

misc/python/materialize/mzcompose/__init__.py

@@ -701,6 +701,9 @@ def get_default_system_parameters(
     # The estimated path is covered explicitly in mysql-cdc/statistics.td and
     # by parallel-workload.
     "mysql_source_snapshot_exact_count_max_rows",
+    # Not varied here because the 256-request floor dominates for test-sized
+    # tables. parallel-workload flips it.
+    "mysql_source_snapshot_partition_requests_per_billion_rows",
     "postgres_fetch_slot_resume_lsn_interval",
     "pg_schema_validation_interval",
     "pg_source_validate_timeline",

misc/python/materialize/parallel_workload/action.py

@@ -3046,6 +3046,14 @@ def __init__(
         self.flags_with_values["mysql_source_snapshot_parallelism"] = (
             BOOLEAN_FLAG_VALUES
         )
+        # 0 leaves only the 256-request floor, the default scales with table
+        # size.
+        self.flags_with_values[
+            "mysql_source_snapshot_partition_requests_per_billion_rows"
+        ] = [
+            "0",
+            "2500",
+        ]
 
         # If you are adding a new config flag in Materialize, consider using it
         # here instead of just marking it as uninteresting to silence the

src/mysql-util/src/partition.rs

@@ -16,13 +16,18 @@ use crate::{KeyProber, MySqlError, QualifiedTableRef};
 /// into `num_workers` roughly even partitions.
 /// This should be run in a repeatable read transaction against a primary key varchar/char column
 /// with the `utf8mb4_bin` collation.
+///
+/// At most `max_requests` probes are issued against the server. When the
+/// budget runs out, remaining prefixes 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);
@@ -31,6 +36,7 @@ pub async fn partition_table(
         num_workers,
         estimated_row_count,
         min_rows_per_worker,
+        max_requests,
     )
     .await?;
     tracing::trace!(
@@ -60,6 +66,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());
@@ -75,15 +82,24 @@ async fn partition<D: PrimaryKeyProber>(
     .max(f64::cast_lossy(min_rows_per_worker))
     .max(1.0);
 
-    compute_boundaries(db, workers, estimated_row_count, target_max_rows_per_prefix).await
+    compute_boundaries(
+        db,
+        workers,
+        estimated_row_count,
+        target_max_rows_per_prefix,
+        max_requests,
+    )
+    .await
 }
 
 async fn compute_boundaries<D: PrimaryKeyProber>(
     db: &mut D,
     workers: usize,
     estimated_row_count: u64,
     target_rows_per_prefix: f64,
+    max_requests: u64,
 ) -> Result<Vec<String>, MySqlError> {
+    let mut budget = max_requests;
     // BFS of prefixes, splitting until estimates fall under the target.
     let mut ordered_prefixes = vec![Prefix {
         prefix: String::new(),
@@ -96,11 +112,14 @@ async fn compute_boundaries<D: PrimaryKeyProber>(
         let mut next_ordered_prefixes: Vec<Prefix> = vec![];
         let mut split_any = false;
         for prefix in ordered_prefixes {
-            if f64::cast_lossy(prefix.estimated_rows) > target_rows_per_prefix {
+            // Entering a split costs one probe for the first prefix plus two
+            // for its first walk step, a smaller budget keeps the prefix as a
+            // leaf.
+            if budget >= 3 && f64::cast_lossy(prefix.estimated_rows) > target_rows_per_prefix {
                 split_any = true;
                 // Partitioning children can drop some rows from the parent prefix range. This
                 // is acceptable given the approximate nature of the algorithm.
-                let children = children_prefixes(db, &prefix).await?;
+                let children = children_prefixes(db, &prefix, &mut budget).await?;
                 next_ordered_prefixes.extend(children);
             } else {
                 next_ordered_prefixes.push(prefix);
@@ -111,6 +130,11 @@ async fn compute_boundaries<D: PrimaryKeyProber>(
             break;
         }
     }
+    tracing::debug!(
+        prefixes = ordered_prefixes.len(),
+        requests_spent = max_requests - budget,
+        "split key space into prefixes"
+    );
 
     // Recompute the total after partitioning the table to get more even splits because the actual row count and the
     // granularly estimated row count can diverge from the original top level estimate.
@@ -147,13 +171,19 @@ async fn compute_boundaries<D: PrimaryKeyProber>(
 ///
 /// Note: This will drop the key "a" on the floor, along with any keys
 /// sorting below their own prefix (below-space characters at this depth).
+///
+/// `budget` is decremented once per probe. The caller must provide at least
+/// 3, one for the first prefix and two for a walk step. The walk closes out
+/// with a tail child once it cannot afford another step.
 async fn children_prefixes<D: PrimaryKeyProber>(
     db: &mut D,
     parent: &Prefix,
+    budget: &mut u64,
 ) -> Result<Vec<Prefix>, 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, parent.end.as_deref(), depth)
         .await?
@@ -162,10 +192,25 @@ async fn children_prefixes<D: PrimaryKeyProber>(
     };
 
     loop {
+        // A walk step costs two probes. When they are unaffordable, close out
+        // with a tail child so the parent's key space stays covered,
+        // estimated as the parent's unconsumed mass.
+        if *budget < 2 {
+            let consumed: u64 = children.iter().map(|c| c.estimated_rows).sum();
+            children.push(Prefix {
+                prefix: cur,
+                end: parent.end.clone(),
+                estimated_rows: parent.estimated_rows.saturating_sub(consumed).max(1),
+                depth,
+            });
+            return Ok(children);
+        }
+        *budget -= 1;
         let next = db
             .prefix_of_first_row_not_matching_prefix(&cur, parent.end.as_deref(), depth)
             .await?;
         let end = next.clone().or_else(|| parent.end.clone());
+        *budget -= 1;
         let estimated_rows = db.estimate_range_rows(&cur, end.as_deref()).await?;
         children.push(Prefix {
             prefix: cur,
@@ -247,9 +292,15 @@ mod tests {
     /// cover them.
     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: &str, end: Option<&str>) -> (usize, usize) {
             // The lower bound is exclusive, a key equal to `start` is skipped.
             let lo = self.keys.partition_point(|k| k.as_str() <= start);
@@ -267,6 +318,7 @@ mod tests {
             start: &str,
             end: Option<&str>,
         ) -> Result<u64, MySqlError> {
+            self.requests += 1;
             let (lo, hi) = self.bounds(start, end);
             Ok(u64::cast_from(hi - lo))
         }
@@ -277,6 +329,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);
@@ -290,6 +343,7 @@ mod tests {
             end: Option<&str>,
             len: usize,
         ) -> Result<Option<String>, MySqlError> {
+            self.requests += 1;
             let (_, hi) = self.bounds("", end);
             // Find the last key matching `cur`, byte prefixes stand in for
             // the collation's LIKE matching.
@@ -310,9 +364,9 @@ 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_ROWS_PER_WORKER).await?;
+        let boundaries = partition(&mut db, 1, count, MIN_ROWS_PER_WORKER, u64::MAX).await?;
         assert!(boundaries.is_empty());
         Ok(())
     }
@@ -321,28 +375,26 @@ mod tests {
     async fn small_table_gets_no_boundaries() -> Result<(), MySqlError> {
         // All keys share one depth-1 prefix and fit under `min_rows_per_worker`,
         // so the single open-ended range yields no boundary.
-        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_ROWS_PER_WORKER).await?;
+        let boundaries = partition(&mut db, 4, count, MIN_ROWS_PER_WORKER, 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_ROWS_PER_WORKER).await?;
+        let mut db = MockDb::new(vec![]);
+        let boundaries = partition(&mut db, 4, 0, MIN_ROWS_PER_WORKER, 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_ROWS_PER_WORKER).await?;
+        let boundaries = partition(&mut db, 4, count, MIN_ROWS_PER_WORKER, u64::MAX).await?;
         assert_eq!(boundaries.len(), 3);
         // Boundaries must be sorted and split the keys into ~50k chunks.
         let mut prev = 0;
@@ -361,9 +413,9 @@ mod tests {
 
     #[mz_ore::test(tokio::test)]
     async fn low_min_rows_per_worker_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 {
@@ -386,9 +438,9 @@ mod tests {
         // stalling on the all-encompassing "U" prefix.
         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!(
@@ -403,8 +455,8 @@ mod tests {
     async fn fractional_target_still_terminates() -> Result<(), MySqlError> {
         // count / (workers * 4) is fractional and the minimum is zero, so
         // the target floors at one row instead of splitting forever.
-        let mut db = MockDb { keys: keys(3) };
-        let boundaries = partition(&mut db, 4, 3, 0).await?;
+        let mut db = MockDb::new(keys(3));
+        let boundaries = partition(&mut db, 4, 3, 0, u64::MAX).await?;
         assert_eq!(boundaries, vec!["000001", "000002"]);
         Ok(())
     }
@@ -437,12 +489,13 @@ mod tests {
         let total = u64::cast_from(all_keys.len());
 
         // A minimum above the table size yields no boundaries at all.
-        let bounds = partition_table(&mut conn, table.clone(), "id", 4, total, 50_000).await?;
+        let bounds =
+            partition_table(&mut conn, table.clone(), "id", 4, total, 50_000, u64::MAX).await?;
         assert!(bounds.is_empty(), "{bounds:?}");
 
         // A low minimum splits inside the 'a' extensions rather than stopping
         // at the exact key.
-        let bounds = partition_table(&mut conn, table, "id", 4, total, 10).await?;
+        let bounds = partition_table(&mut conn, table, "id", 4, total, 10, u64::MAX).await?;
         assert_eq!(bounds.len(), 3, "{bounds:?}");
 
         // MySQL agrees the boundaries are strictly increasing.
@@ -488,7 +541,7 @@ mod tests {
         let total = u64::cast_from(all_keys.len());
 
         // Partition for 4 workers with a minimum split size around 250.
-        let bounds = partition_table(&mut conn, table, "id", 4, total, 250).await?;
+        let bounds = partition_table(&mut conn, table, "id", 4, total, 250, u64::MAX).await?;
         assert_eq!(bounds.len(), 3);
         let counts = partition_counts(&mut conn, DB, &bounds, total).await?;
         // ~8k keys are visible, so each count gets at least 2k under perfect

src/storage-types/src/dyncfgs.rs

@@ -215,6 +215,18 @@ pub static MYSQL_SOURCE_SNAPSHOT_PARALLELISM: Config<bool> = Config::new(
     "Whether to split MySQL snapshot reads across workers by primary-key ranges.",
 );
 
+/// 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(
@@ -438,6 +450,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&MYSQL_REPLICATION_HEARTBEAT_INTERVAL)
         .add(&MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS)
         .add(&MYSQL_SOURCE_SNAPSHOT_PARALLELISM)
+        .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)