Skip to main content

rivet/tuning/
profile.rs

1//! Tuning profiles and the resolved `SourceTuning` config.
2//!
3//! Three baked-in profiles (`Fast`, `Balanced`, `Safe`) ship per-field defaults;
4//! a YAML `TuningConfig` can override individual fields and the chosen profile.
5//! `from_config_with_default_profile` is the production entry point and is
6//! wired in [`crate::plan::build`] to honour `source.environment:` —
7//! `Local` → `Fast`, `Replica`/`Production` → `Balanced`.
8
9use arrow::datatypes::SchemaRef;
10use schemars::JsonSchema;
11use serde::{Deserialize, Serialize};
12
13use super::memory::{compute_batch_size_from_memory, estimate_row_bytes};
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct SourceTuning {
17    pub batch_size: usize,
18    pub batch_size_memory_mb: Option<usize>,
19    pub throttle_ms: u64,
20    pub statement_timeout_s: u64,
21    pub max_retries: u32,
22    pub retry_backoff_ms: u64,
23    pub lock_timeout_s: u64,
24    /// RSS limit in MB before chunk processing throttles. `0` = no limit (disabled).
25    pub memory_threshold_mb: usize,
26    /// Hard cap on a single Arrow batch in MB. `None` = no cap.
27    pub max_batch_memory_mb: Option<usize>,
28    pub on_batch_memory_exceeded: BatchMemoryPolicy,
29    /// When true, Rivet samples DB pressure metrics every
30    /// [`super::ADAPTIVE_SAMPLE_INTERVAL`] batches and shrinks/restores the
31    /// fetch size in response. Default: false.
32    pub adaptive: bool,
33    /// Floor for the OPT-2 concurrency governor: the lowest worker/connection
34    /// count it will back parallelism down to under source pressure. `None`
35    /// ⇒ 1. The ceiling is the export's configured `parallel`. Only consulted
36    /// when `adaptive` is on and `parallel > 1`.
37    pub min_parallel: Option<usize>,
38    /// Hard ceiling on a single cell/value in MB (OPT-1 memory hardening): a
39    /// variable-length value (text/JSON/blob) larger than this aborts the run
40    /// with `RIVET_VALUE_TOO_LARGE` instead of risking OOM, since the
41    /// average-based batch cap can't bound one giant cell. `Some(0)` / `None`
42    /// disable the guard. Default: 256 MiB.
43    pub max_value_mb: Option<usize>,
44    configured_profile: TuningProfile,
45}
46
47#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
48#[serde(rename_all = "lowercase")]
49pub enum TuningProfile {
50    Fast,
51    Balanced,
52    Safe,
53}
54
55/// Action taken when a single Arrow batch exceeds `max_batch_memory_mb`.
56#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
57#[serde(rename_all = "snake_case")]
58pub enum BatchMemoryPolicy {
59    /// Log a warning and continue. (default)
60    #[default]
61    Warn,
62    /// Return an error — the export fails immediately.
63    Fail,
64    /// Split the oversized batch in half recursively until each sub-batch fits,
65    /// then process them individually. Transparent to the rest of the pipeline.
66    AutoShrink,
67}
68
69#[derive(Debug, Deserialize, Serialize, JsonSchema, Default, Clone)]
70#[serde(deny_unknown_fields)]
71pub struct TuningConfig {
72    pub profile: Option<TuningProfile>,
73    pub batch_size: Option<usize>,
74    /// Target memory per batch in MB. Mutually exclusive with batch_size.
75    pub batch_size_memory_mb: Option<usize>,
76    pub throttle_ms: Option<u64>,
77    pub statement_timeout_s: Option<u64>,
78    pub max_retries: Option<u32>,
79    pub retry_backoff_ms: Option<u64>,
80    pub lock_timeout_s: Option<u64>,
81    pub memory_threshold_mb: Option<usize>,
82    /// Hard cap on Arrow batch memory in MB. When a batch exceeds this limit,
83    /// `on_batch_memory_exceeded` determines the response.
84    pub max_batch_memory_mb: Option<usize>,
85    /// Policy applied when a batch exceeds `max_batch_memory_mb`. Default: `warn`.
86    pub on_batch_memory_exceeded: Option<BatchMemoryPolicy>,
87    /// Enable real-time batch size adaptation based on DB pressure metrics.
88    /// Postgres: samples `pg_stat_bgwriter`. MySQL: samples `Innodb_log_waits`.
89    /// Also arms the OPT-2 concurrency governor when `parallel > 1`.
90    pub adaptive: Option<bool>,
91    /// Floor for the concurrency governor (lowest parallelism under pressure).
92    /// Default 1. Ceiling is the export's `parallel`.
93    pub min_parallel: Option<usize>,
94    /// Hard per-value size ceiling in MB. A single text/JSON/blob cell larger
95    /// than this aborts the run with `RIVET_VALUE_TOO_LARGE`. `0` disables the
96    /// guard. Default: 256.
97    pub max_value_mb: Option<usize>,
98}
99
100/// Layer `export` on top of `source`: each field uses export when set, otherwise source.
101/// `None` only when both inputs are `None`.
102pub fn merge_tuning_config(
103    source: Option<&TuningConfig>,
104    export: Option<&TuningConfig>,
105) -> Option<TuningConfig> {
106    match (source, export) {
107        (None, None) => None,
108        (Some(s), None) => Some(s.clone()),
109        (None, Some(e)) => Some(e.clone()),
110        (Some(s), Some(e)) => Some(TuningConfig {
111            profile: e.profile.or(s.profile),
112            batch_size: e.batch_size.or(s.batch_size),
113            batch_size_memory_mb: e.batch_size_memory_mb.or(s.batch_size_memory_mb),
114            throttle_ms: e.throttle_ms.or(s.throttle_ms),
115            statement_timeout_s: e.statement_timeout_s.or(s.statement_timeout_s),
116            max_retries: e.max_retries.or(s.max_retries),
117            retry_backoff_ms: e.retry_backoff_ms.or(s.retry_backoff_ms),
118            lock_timeout_s: e.lock_timeout_s.or(s.lock_timeout_s),
119            memory_threshold_mb: e.memory_threshold_mb.or(s.memory_threshold_mb),
120            max_batch_memory_mb: e.max_batch_memory_mb.or(s.max_batch_memory_mb),
121            on_batch_memory_exceeded: e.on_batch_memory_exceeded.or(s.on_batch_memory_exceeded),
122            adaptive: e.adaptive.or(s.adaptive),
123            min_parallel: e.min_parallel.or(s.min_parallel),
124            max_value_mb: e.max_value_mb.or(s.max_value_mb),
125        }),
126    }
127}
128
129impl SourceTuning {
130    /// The per-value byte ceiling: `max_value_mb` converted to bytes, with
131    /// `Some(0)` / `None` meaning "disabled". Single source of truth shared by
132    /// the sink's post-materialization guard (`sink::check_value_ceiling`) and
133    /// the per-engine pre-allocation guard (`source::value_within_ceiling`), so
134    /// the two definitions of "0 disables" can never drift apart.
135    pub(crate) fn max_value_bytes(&self) -> Option<usize> {
136        self.max_value_mb
137            .filter(|&mb| mb > 0)
138            .map(|mb| mb * 1024 * 1024)
139    }
140
141    /// Build tuning with the legacy `Balanced` fallback. Public for downstream
142    /// callers and tests; production resolution in [`crate::plan::build`] uses
143    /// [`Self::from_config_with_default_profile`] so that `source.environment:`
144    /// can pick the right default.
145    #[allow(dead_code)]
146    pub fn from_config(config: Option<&TuningConfig>) -> Self {
147        Self::from_config_with_default_profile(config, TuningProfile::Balanced)
148    }
149
150    /// Like [`Self::from_config`] but lets the caller override the fallback
151    /// profile used when `config.profile` is unset.
152    pub fn from_config_with_default_profile(
153        config: Option<&TuningConfig>,
154        fallback_profile: TuningProfile,
155    ) -> Self {
156        let profile = config.and_then(|c| c.profile).unwrap_or(fallback_profile);
157
158        let mut tuning = Self::from_profile(profile);
159        tuning.configured_profile = profile;
160
161        if let Some(cfg) = config {
162            if let Some(v) = cfg.batch_size {
163                tuning.batch_size = v;
164            }
165            // Preserve the profile's memory-driven per-batch cap unless the user
166            // RESTATED it. A bare `tuning: {profile: fast}` — or any block that
167            // omits this field (e.g. `{throttle_ms: 100}`) — must NOT silently
168            // disable Balanced's Some(32) / Fast's Some(64), which is what caps
169            // PostgreSQL's Arrow FETCH memory (source/postgres/mod.rs); an
170            // unconditional assign treated the field's ABSENCE as "disable" and
171            // ~3x'd per-batch RSS on wide tables. Only an explicit `batch_size`
172            // (a row cap) drops it, since effective_batch_size prefers the memory
173            // cap when both are present. Mirrors the is_some()-guarded siblings
174            // (min_parallel, max_value_mb) below.
175            if cfg.batch_size_memory_mb.is_some() {
176                tuning.batch_size_memory_mb = cfg.batch_size_memory_mb;
177            } else if cfg.batch_size.is_some() {
178                tuning.batch_size_memory_mb = None;
179            }
180            if let Some(v) = cfg.throttle_ms {
181                tuning.throttle_ms = v;
182            }
183            if let Some(v) = cfg.statement_timeout_s {
184                tuning.statement_timeout_s = v;
185            }
186            if let Some(v) = cfg.max_retries {
187                tuning.max_retries = v;
188            }
189            if let Some(v) = cfg.retry_backoff_ms {
190                tuning.retry_backoff_ms = v;
191            }
192            if let Some(v) = cfg.lock_timeout_s {
193                tuning.lock_timeout_s = v;
194            }
195            if let Some(v) = cfg.memory_threshold_mb {
196                tuning.memory_threshold_mb = v;
197            }
198            tuning.max_batch_memory_mb = cfg.max_batch_memory_mb;
199            if let Some(v) = cfg.on_batch_memory_exceeded {
200                tuning.on_batch_memory_exceeded = v;
201            }
202            if let Some(v) = cfg.adaptive {
203                tuning.adaptive = v;
204            }
205            if cfg.min_parallel.is_some() {
206                tuning.min_parallel = cfg.min_parallel;
207            }
208            if cfg.max_value_mb.is_some() {
209                tuning.max_value_mb = cfg.max_value_mb;
210            }
211        }
212
213        tuning
214    }
215
216    fn from_profile(profile: TuningProfile) -> Self {
217        match profile {
218            TuningProfile::Fast => Self {
219                // Memory-driven, 2× Balanced's target (see its note). Narrow tables
220                // still clamp to 150k; the larger target only sizes up wide ones.
221                batch_size: 50_000,
222                batch_size_memory_mb: Some(64),
223                throttle_ms: 0,
224                statement_timeout_s: 0,
225                max_retries: 1,
226                retry_backoff_ms: 1_000,
227                lock_timeout_s: 0,
228                memory_threshold_mb: 0,
229                max_batch_memory_mb: None,
230                on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
231                adaptive: false,
232                min_parallel: None,
233                max_value_mb: Some(256),
234                configured_profile: TuningProfile::Fast,
235            },
236            TuningProfile::Balanced => Self {
237                // Memory-driven batch by default: a 32 MB-per-flush target sizes the
238                // batch to the row width — large for narrow tables (where the 150k
239                // hard-max in `compute_batch_size_from_memory`, not this target, binds
240                // and bounds the raw-row accumulator's RSS), small for wide ones. 32
241                // not 64 MB: a 64 MB target grew medium-wide tables (~2.6 KB/row) to
242                // ~21k rows — ~6% slower for zero flush-count gain — while 32 MB lands
243                // them near the old static 10k; narrow tables are unaffected (clamp
244                // binds either way). The static `batch_size` below is the no-schema
245                // fallback + advisory base.
246                batch_size: 10_000,
247                batch_size_memory_mb: Some(32),
248                throttle_ms: 50,
249                statement_timeout_s: 300,
250                max_retries: 3,
251                retry_backoff_ms: 2_000,
252                lock_timeout_s: 30,
253                memory_threshold_mb: 4_096,
254                max_batch_memory_mb: None,
255                on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
256                adaptive: false,
257                min_parallel: None,
258                max_value_mb: Some(256),
259                configured_profile: TuningProfile::Balanced,
260            },
261            TuningProfile::Safe => Self {
262                batch_size: 2_000,
263                batch_size_memory_mb: None,
264                throttle_ms: 500,
265                statement_timeout_s: 120,
266                max_retries: 5,
267                retry_backoff_ms: 5_000,
268                lock_timeout_s: 10,
269                memory_threshold_mb: 2_048,
270                max_batch_memory_mb: None,
271                on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
272                adaptive: false,
273                min_parallel: None,
274                max_value_mb: Some(256),
275                configured_profile: TuningProfile::Safe,
276            },
277        }
278    }
279
280    pub fn profile_name(&self) -> &'static str {
281        match self.configured_profile {
282            TuningProfile::Fast => "fast",
283            TuningProfile::Balanced => "balanced",
284            TuningProfile::Safe => "safe",
285        }
286    }
287
288    /// If `batch_size_memory_mb` is set, compute and return an adjusted batch_size
289    /// from the schema; otherwise return the configured `batch_size`.
290    pub fn effective_batch_size(&self, schema: Option<&SchemaRef>) -> usize {
291        if let (Some(mem_mb), Some(schema)) = (self.batch_size_memory_mb, schema) {
292            let computed = compute_batch_size_from_memory(mem_mb, schema);
293            log::info!(
294                "batch_size_memory_mb={}: estimated row ~{}B, computed batch_size={}",
295                mem_mb,
296                estimate_row_bytes(schema),
297                computed
298            );
299            computed
300        } else {
301            self.batch_size
302        }
303    }
304
305    /// Return the actual Arrow memory footprint of a batch in bytes.
306    ///
307    /// Sums `get_array_memory_size()` across all columns — includes buffers for
308    /// validity bitmaps, offsets, and value data. Does not include Arrow struct
309    /// overhead (~few hundred bytes) which is negligible at batch scale.
310    pub fn batch_memory_bytes(batch: &arrow::record_batch::RecordBatch) -> usize {
311        batch
312            .columns()
313            .iter()
314            .map(|col| col.get_array_memory_size())
315            .sum()
316    }
317
318    /// Produce a `ResourceSummary` from the resolved tuning settings.
319    ///
320    /// The summary requires no database connection. It reports two batch-memory
321    /// bounds based on narrow-table (~200 B/row) and wide-table (~10 KB/row)
322    /// heuristics. A `wide_table_risk` flag is set when the upper bound exceeds
323    /// 128 MB per batch.
324    pub fn resource_summary(&self) -> ResourceSummary {
325        const NARROW_BYTES: f64 = 200.0;
326        const WIDE_BYTES: f64 = 10_240.0;
327        let batch = self.batch_size as f64;
328        let batch_narrow_mb = batch * NARROW_BYTES / (1024.0 * 1024.0);
329        let batch_wide_mb = batch * WIDE_BYTES / (1024.0 * 1024.0);
330        ResourceSummary {
331            profile: self.profile_name().to_string(),
332            batch_size: self.batch_size,
333            batch_size_memory_mb: self.batch_size_memory_mb,
334            memory_threshold_mb: self.memory_threshold_mb,
335            throttle_ms: self.throttle_ms,
336            batch_narrow_mb,
337            batch_wide_mb,
338            wide_table_risk: batch_wide_mb > 128.0,
339        }
340    }
341}
342
343impl std::fmt::Display for SourceTuning {
344    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345        write!(
346            f,
347            "profile={}, batch_size={}, throttle={}ms, timeout={}s, retries={}, lock_timeout={}s",
348            self.profile_name(),
349            self.batch_size,
350            self.throttle_ms,
351            self.statement_timeout_s,
352            self.max_retries,
353            self.lock_timeout_s,
354        )
355    }
356}
357
358/// Resource estimate computed from tuning settings alone (no DB connection required).
359///
360/// `batch_narrow_mb` and `batch_wide_mb` bracket the expected per-batch memory:
361/// - narrow table: ~200 B/row (int-heavy, no text blobs)
362/// - wide table  : ~10 KB/row (many text/JSON/binary columns)
363///
364/// Use `wide_table_risk` to decide whether to recommend `adaptive_batch` or a
365/// lower `batch_size`.
366#[derive(Debug, Clone)]
367pub struct ResourceSummary {
368    #[allow(dead_code)]
369    pub profile: String,
370    pub batch_size: usize,
371    pub batch_size_memory_mb: Option<usize>,
372    pub memory_threshold_mb: usize,
373    pub throttle_ms: u64,
374    pub batch_narrow_mb: f64,
375    pub batch_wide_mb: f64,
376    pub wide_table_risk: bool,
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382
383    fn cfg_with_profile(profile: TuningProfile) -> TuningConfig {
384        TuningConfig {
385            profile: Some(profile),
386            ..Default::default()
387        }
388    }
389
390    #[test]
391    fn default_config_uses_balanced_profile() {
392        let t = SourceTuning::from_config(None);
393        assert_eq!(t.batch_size, 10_000);
394        assert_eq!(t.throttle_ms, 50);
395        assert_eq!(t.statement_timeout_s, 300);
396        assert_eq!(t.max_retries, 3);
397        assert_eq!(t.retry_backoff_ms, 2_000);
398        assert_eq!(t.lock_timeout_s, 30);
399    }
400
401    #[test]
402    fn fallback_profile_used_when_no_profile_in_config() {
403        let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Fast);
404        assert_eq!(t.batch_size, 50_000);
405        assert_eq!(t.throttle_ms, 0, "fallback to Fast must zero the throttle");
406        assert_eq!(t.profile_name(), "fast");
407
408        let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Safe);
409        assert_eq!(t.throttle_ms, 500);
410        assert_eq!(t.profile_name(), "safe");
411    }
412
413    #[test]
414    fn explicit_profile_wins_over_fallback() {
415        let cfg = cfg_with_profile(TuningProfile::Balanced);
416        let t = SourceTuning::from_config_with_default_profile(Some(&cfg), TuningProfile::Fast);
417        assert_eq!(
418            t.throttle_ms, 50,
419            "explicit balanced profile must keep its throttle"
420        );
421        assert_eq!(t.profile_name(), "balanced");
422    }
423
424    #[test]
425    fn fast_profile_favors_throughput() {
426        let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
427        assert_eq!(t.batch_size, 50_000);
428        assert_eq!(t.throttle_ms, 0);
429        assert_eq!(t.statement_timeout_s, 0);
430        assert_eq!(t.max_retries, 1);
431    }
432
433    #[test]
434    fn safe_profile_limits_impact() {
435        let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
436        assert_eq!(t.batch_size, 2_000);
437        assert_eq!(t.throttle_ms, 500);
438        assert_eq!(t.statement_timeout_s, 120);
439        assert_eq!(t.max_retries, 5);
440        assert_eq!(t.retry_backoff_ms, 5_000);
441        assert_eq!(t.lock_timeout_s, 10);
442    }
443
444    #[test]
445    fn explicit_fields_override_profile_defaults() {
446        let cfg = TuningConfig {
447            profile: Some(TuningProfile::Safe),
448            batch_size: Some(3_000),
449            throttle_ms: Some(250),
450            ..Default::default()
451        };
452        let t = SourceTuning::from_config(Some(&cfg));
453        assert_eq!(t.batch_size, 3_000, "explicit batch_size should win");
454        assert_eq!(t.throttle_ms, 250, "explicit throttle_ms should win");
455        assert_eq!(
456            t.statement_timeout_s, 120,
457            "non-overridden field stays at safe default"
458        );
459        assert_eq!(
460            t.max_retries, 5,
461            "non-overridden field stays at safe default"
462        );
463    }
464
465    #[test]
466    fn profile_name_fast() {
467        let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
468        assert_eq!(t.profile_name(), "fast");
469    }
470
471    #[test]
472    fn profile_name_balanced() {
473        let t = SourceTuning::from_config(None);
474        assert_eq!(t.profile_name(), "balanced");
475    }
476
477    #[test]
478    fn profile_name_safe() {
479        let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
480        assert_eq!(t.profile_name(), "safe");
481    }
482
483    #[test]
484    fn display_contains_all_fields() {
485        let t = SourceTuning::from_config(None);
486        let s = t.to_string();
487        assert!(s.contains("profile=balanced"), "missing profile in: {s}");
488        assert!(s.contains("batch_size=10000"), "missing batch_size in: {s}");
489        assert!(s.contains("throttle=50ms"), "missing throttle in: {s}");
490        assert!(s.contains("timeout=300s"), "missing timeout in: {s}");
491        assert!(s.contains("retries=3"), "missing retries in: {s}");
492        assert!(
493            s.contains("lock_timeout=30s"),
494            "missing lock_timeout in: {s}"
495        );
496    }
497
498    #[test]
499    fn merge_tuning_export_overrides_source_fields() {
500        let source = TuningConfig {
501            profile: Some(TuningProfile::Fast),
502            batch_size: Some(1_000),
503            throttle_ms: Some(0),
504            ..Default::default()
505        };
506        let export = TuningConfig {
507            profile: Some(TuningProfile::Safe),
508            batch_size: None,
509            ..Default::default()
510        };
511        let m = merge_tuning_config(Some(&source), Some(&export)).expect("merged");
512        assert_eq!(m.profile, Some(TuningProfile::Safe));
513        assert_eq!(
514            m.batch_size,
515            Some(1_000),
516            "export omitted batch_size -> keep source"
517        );
518        assert_eq!(m.throttle_ms, Some(0));
519    }
520
521    #[test]
522    fn merge_tuning_export_only() {
523        let e = cfg_with_profile(TuningProfile::Fast);
524        let m = merge_tuning_config(None, Some(&e)).expect("merged");
525        assert_eq!(m.profile, Some(TuningProfile::Fast));
526    }
527
528    #[test]
529    fn effective_batch_size_without_memory() {
530        let t = SourceTuning::from_config(None);
531        assert_eq!(t.effective_batch_size(None), 10_000);
532    }
533
534    #[test]
535    fn effective_batch_size_with_memory() {
536        use arrow::datatypes::{DataType, Field, Schema};
537        use std::sync::Arc;
538        let cfg = TuningConfig {
539            batch_size_memory_mb: Some(256),
540            ..Default::default()
541        };
542        let t = SourceTuning::from_config(Some(&cfg));
543        let schema = Arc::new(Schema::new(vec![
544            Field::new("id", DataType::Int64, false),
545            Field::new("name", DataType::Utf8, true),
546        ]));
547        let bs = t.effective_batch_size(Some(&schema));
548        assert!((1_000..=150_000).contains(&bs), "got {bs}");
549        // 256MB / 266B ≈ 1_009_022, clamped to 150_000
550        assert_eq!(bs, 150_000);
551    }
552
553    #[test]
554    fn a_tuning_block_omitting_batch_size_memory_mb_keeps_the_profile_default() {
555        // RED before the is_some() guard: a `tuning:` block that sets any UNRELATED
556        // field (here throttle_ms) clobbered Balanced's Some(32) / Fast's Some(64)
557        // memory-driven per-batch cap to None — treating the field's ABSENCE as
558        // "disable" — which ~3x'd PostgreSQL per-batch Arrow RSS on wide tables.
559        let only_throttle = TuningConfig {
560            throttle_ms: Some(50),
561            ..Default::default()
562        };
563        assert_eq!(
564            SourceTuning::from_config(Some(&only_throttle)).batch_size_memory_mb,
565            Some(32),
566            "an unrelated tuning field must not disable Balanced's memory cap"
567        );
568        let fast = TuningConfig {
569            profile: Some(TuningProfile::Fast),
570            ..Default::default()
571        };
572        assert_eq!(
573            SourceTuning::from_config(Some(&fast)).batch_size_memory_mb,
574            Some(64),
575            "naming the Fast profile must keep its Some(64) memory cap"
576        );
577        // An explicit row cap (batch_size) DOES drop the memory cap — the two are
578        // mutually exclusive and effective_batch_size prefers the memory cap.
579        let rowcap = TuningConfig {
580            batch_size: Some(5_000),
581            ..Default::default()
582        };
583        assert_eq!(
584            SourceTuning::from_config(Some(&rowcap)).batch_size_memory_mb,
585            None,
586            "an explicit batch_size drops the memory cap so the row cap wins"
587        );
588        // An explicit memory value is honored verbatim.
589        let memcap = TuningConfig {
590            batch_size_memory_mb: Some(16),
591            ..Default::default()
592        };
593        assert_eq!(
594            SourceTuning::from_config(Some(&memcap)).batch_size_memory_mb,
595            Some(16)
596        );
597    }
598
599    #[test]
600    fn resource_summary_balanced_profile() {
601        let t = SourceTuning::from_config(None);
602        let r = t.resource_summary();
603        assert_eq!(r.profile, "balanced");
604        assert_eq!(r.batch_size, 10_000);
605        // Balanced drives batch sizing from a 32 MB-per-flush target by default
606        // (the static batch_size above is the no-schema fallback). 32 not 64 MB:
607        // 64 overshot medium-wide tables; 32 keeps the narrow win (clamp-bound)
608        // while landing wide tables near the old static size — see profile note.
609        assert_eq!(r.batch_size_memory_mb, Some(32));
610        assert_eq!(r.memory_threshold_mb, 4_096);
611        assert_eq!(r.throttle_ms, 50);
612        // narrow: 10_000 × 200 B = ~1.9 MB
613        assert!(
614            r.batch_narrow_mb < 5.0,
615            "narrow too high: {}",
616            r.batch_narrow_mb
617        );
618        // wide: 10_000 × 10 KB = ~95 MB — no risk (< 128 MB)
619        assert!(
620            !r.wide_table_risk,
621            "balanced 10k should not trigger wide_table_risk"
622        );
623    }
624
625    #[test]
626    fn resource_summary_fast_profile_triggers_wide_table_risk() {
627        let t = SourceTuning::from_config(Some(&TuningConfig {
628            profile: Some(TuningProfile::Fast),
629            ..Default::default()
630        }));
631        let r = t.resource_summary();
632        assert_eq!(r.batch_size, 50_000);
633        // wide: 50_000 × 10 KB = ~476 MB → high risk
634        assert!(r.wide_table_risk, "fast 50k should trigger wide_table_risk");
635    }
636
637    #[test]
638    fn resource_summary_with_adaptive_batch() {
639        let cfg = TuningConfig {
640            batch_size_memory_mb: Some(64),
641            ..Default::default()
642        };
643        let t = SourceTuning::from_config(Some(&cfg));
644        let r = t.resource_summary();
645        assert_eq!(r.batch_size_memory_mb, Some(64));
646    }
647}