Skip to main content

hawdb_executor/
memory.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Execution memory defaults and admission estimates.
16
17use hawdb_core::{HawDBError, Result, RuntimeTaskContext};
18use hawdb_plan_cypher::{
19    GraphExpansionBudget, PhysicalPlan, PlanChildren, VectorExecutionResourceProfile,
20};
21use hawdb_storage::mutation::MutationLimits;
22use std::num::{NonZeroU64, NonZeroUsize};
23use std::path::PathBuf;
24use std::time::Duration;
25
26#[doc(hidden)]
27pub fn enforced_query_memory_budget(
28    memory: &ExecutionMemoryConfig,
29    task_context: Option<&RuntimeTaskContext>,
30) -> Result<NonZeroUsize> {
31    let Some(reservation) = task_context.and_then(RuntimeTaskContext::memory_reservation) else {
32        return Ok(memory.query_memory_bytes);
33    };
34    // Never widen an undersized admission to an operator-configured floor.
35    // The shared root remains the admitted reservation and the operator fails
36    // closed when its first charge cannot fit.
37    runtime_memory_budget("query memory", reservation.memory_bytes())
38}
39
40#[doc(hidden)]
41pub fn enforced_result_memory_budget(
42    memory: &ExecutionMemoryConfig,
43    task_context: Option<&RuntimeTaskContext>,
44) -> Result<NonZeroUsize> {
45    enforced_result_memory_budget_with_option(memory, task_context).map(|(budget, _)| budget)
46}
47
48/// Returns the effective result limit and the configuration or admission that supplied it.
49#[doc(hidden)]
50pub fn enforced_result_memory_budget_with_option(
51    memory: &ExecutionMemoryConfig,
52    task_context: Option<&RuntimeTaskContext>,
53) -> Result<(NonZeroUsize, &'static str)> {
54    let configured = (
55        memory.query_memory_bytes,
56        "ExecutionMemoryConfig.query_memory_bytes",
57    );
58    let Some(reservation) = task_context.and_then(RuntimeTaskContext::memory_reservation) else {
59        return Ok(configured);
60    };
61    let admitted = runtime_memory_budget("query result", reservation.result_bytes())?;
62    if admitted <= memory.query_memory_bytes {
63        Ok((
64            admitted,
65            "RuntimeTaskContext.memory_reservation.result_bytes",
66        ))
67    } else {
68        Ok(configured)
69    }
70}
71
72fn runtime_memory_budget(owner: &str, bytes: u64) -> Result<NonZeroUsize> {
73    let bytes = usize::try_from(bytes).map_err(|_| {
74        HawDBError::Execution(format!(
75            "runtime-admitted {owner} reservation {bytes} does not fit the executor address space"
76        ))
77    })?;
78    NonZeroUsize::new(bytes).ok_or_else(|| {
79        HawDBError::Execution(format!(
80            "runtime-admitted {owner} reservation must be non-zero"
81        ))
82    })
83}
84
85#[doc(hidden)]
86pub const DEFAULT_EXECUTION_BATCH_ROWS: usize = 256;
87const DEFAULT_EXECUTION_BATCH_PAYLOAD_BYTES: usize = 16 * 1024 * 1024;
88const DEFAULT_QUERY_MEMORY_BYTES: usize = 256 * 1024 * 1024;
89#[doc(hidden)]
90pub const DEFAULT_BLOCKING_OPERATOR_MEMORY_BYTES: usize = 64 * 1024 * 1024;
91const DEFAULT_EXECUTION_MAX_SPILL_BYTES: u64 = 4 * 1024 * 1024 * 1024;
92const DEFAULT_EXECUTION_MAX_SPILL_RUNS: usize = 128;
93const DEFAULT_EXECUTION_MAX_TOTAL_SPILL_BYTES: u64 = 16 * 1024 * 1024 * 1024;
94const DEFAULT_EXECUTION_MAX_TOTAL_SPILL_RUNS: usize = 512;
95const DEFAULT_EXECUTION_MIN_SPILL_FREE_BYTES: u64 = 1024 * 1024 * 1024;
96const DEFAULT_SPILL_FREE_SPACE_PROBE_INTERVAL_BYTES: u64 = 64 * 1024 * 1024;
97const DEFAULT_SPILL_ORPHAN_GRACE_PERIOD: Duration = Duration::from_secs(24 * 60 * 60);
98const MUTATION_OPERATION_BOOKKEEPING_BYTES: u64 = 64;
99const MUTATION_AFFECTED_ROW_BOOKKEEPING_BYTES: u64 = 16;
100const MUTATION_RESULT_ROW_BOOKKEEPING_BYTES: u64 = 64;
101#[doc(hidden)]
102pub const SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES: u64 = 2 * 1024 * 1024;
103
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct ExecutionMemoryConfig {
106    /// Maximum tracked resident bytes owned by one query across all operators.
107    pub query_memory_bytes: NonZeroUsize,
108    /// Maximum row count in an executor-owned transfer batch.
109    pub batch_rows: NonZeroUsize,
110    /// Maximum estimated resident bytes in an executor-owned transfer batch.
111    pub batch_payload_bytes: NonZeroUsize,
112    /// Maximum estimated bytes retained for blocking state, including sort,
113    /// aggregation, deduplication, graph traversal state and selected top-K.
114    /// Source reads use the shared query budget independently of this cap.
115    pub blocking_operator_bytes: NonZeroUsize,
116    /// Optional row and cumulative payload limits for each seeded expansion.
117    /// `None` adds no cumulative cap; live allocations still obey the operator,
118    /// batch and shared query budgets, including runtime-admitted reservations.
119    pub graph_expansion_budget: Option<GraphExpansionBudget>,
120    /// Maximum cumulative serialized spill bytes, including merge passes.
121    pub max_spill_bytes: NonZeroU64,
122    /// Maximum cumulative spill runs created, including merge passes.
123    pub max_spill_runs: NonZeroUsize,
124    /// Maximum live spill bytes shared by all queries using the spill directory.
125    pub max_total_spill_bytes: NonZeroU64,
126    /// Maximum live spill runs shared by all queries using the spill directory.
127    pub max_total_spill_runs: NonZeroUsize,
128    /// Free space preserved on the filesystem containing the spill directory.
129    pub min_spill_free_bytes: NonZeroU64,
130    /// Reserved bytes between filesystem free-space probes for one spill pool.
131    pub spill_free_space_probe_interval_bytes: NonZeroU64,
132    /// Minimum age before a spill file from an earlier process is removed.
133    pub spill_orphan_grace_period: Duration,
134    /// Directory governed as one shared spill pool.
135    pub spill_directory: PathBuf,
136}
137
138impl Default for ExecutionMemoryConfig {
139    fn default() -> Self {
140        Self {
141            query_memory_bytes: NonZeroUsize::new(DEFAULT_QUERY_MEMORY_BYTES)
142                .expect("default query memory budget is non-zero"),
143            batch_rows: NonZeroUsize::new(DEFAULT_EXECUTION_BATCH_ROWS)
144                .expect("default execution batch size is non-zero"),
145            batch_payload_bytes: NonZeroUsize::new(DEFAULT_EXECUTION_BATCH_PAYLOAD_BYTES)
146                .expect("default execution batch byte size is non-zero"),
147            blocking_operator_bytes: NonZeroUsize::new(DEFAULT_BLOCKING_OPERATOR_MEMORY_BYTES)
148                .expect("default blocking operator memory budget is non-zero"),
149            graph_expansion_budget: None,
150            max_spill_bytes: NonZeroU64::new(DEFAULT_EXECUTION_MAX_SPILL_BYTES)
151                .expect("default spill byte budget is non-zero"),
152            max_spill_runs: NonZeroUsize::new(DEFAULT_EXECUTION_MAX_SPILL_RUNS)
153                .expect("default spill run budget is non-zero"),
154            max_total_spill_bytes: NonZeroU64::new(DEFAULT_EXECUTION_MAX_TOTAL_SPILL_BYTES)
155                .expect("default total spill byte budget is non-zero"),
156            max_total_spill_runs: NonZeroUsize::new(DEFAULT_EXECUTION_MAX_TOTAL_SPILL_RUNS)
157                .expect("default total spill run budget is non-zero"),
158            min_spill_free_bytes: NonZeroU64::new(DEFAULT_EXECUTION_MIN_SPILL_FREE_BYTES)
159                .expect("default spill free-space reserve is non-zero"),
160            spill_free_space_probe_interval_bytes: NonZeroU64::new(
161                DEFAULT_SPILL_FREE_SPACE_PROBE_INTERVAL_BYTES,
162            )
163            .expect("default spill free-space probe interval is non-zero"),
164            spill_orphan_grace_period: DEFAULT_SPILL_ORPHAN_GRACE_PERIOD,
165            #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
166            spill_directory: std::env::temp_dir().join("hawdb-spill"),
167            // Browser queries are memory-only. Do not call the unsupported temp_dir().
168            #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
169            spill_directory: PathBuf::new(),
170        }
171    }
172}
173
174impl ExecutionMemoryConfig {
175    /// Returns process-wide spill usage for this configured directory.
176    pub fn spill_pool_snapshot(&self) -> hawdb_core::Result<crate::SpillPoolSnapshot> {
177        crate::spill::spill_pool_snapshot(self)
178    }
179}
180
181#[derive(Debug, Clone, Copy, PartialEq, Eq)]
182#[doc(hidden)]
183pub struct ExecutionMemoryEstimate {
184    pub pipeline_batch_count: usize,
185    pub blocking_operator_count: usize,
186    pub pipeline_bytes: u64,
187    pub blocking_bytes: u64,
188    pub external_read_bytes: u64,
189    pub fixed_operator_bytes: u64,
190    pub total_bytes: u64,
191}
192
193#[doc(hidden)]
194pub fn estimated_execution_memory(
195    plan: &PhysicalPlan,
196    memory: &ExecutionMemoryConfig,
197) -> ExecutionMemoryEstimate {
198    let shape = peak_execution_memory_shape(plan, memory);
199    let pipeline_bytes = usize_to_u64(shape.pipeline_batch_count)
200        .saturating_mul(usize_to_u64(memory.batch_payload_bytes.get()));
201    let blocking_bytes = usize_to_u64(shape.blocking_operator_count)
202        .saturating_mul(usize_to_u64(memory.blocking_operator_bytes.get()));
203    let external_read_bytes = shape.external_read_bytes;
204    let fixed_operator_bytes = shape.fixed_operator_bytes;
205    ExecutionMemoryEstimate {
206        pipeline_batch_count: shape.pipeline_batch_count,
207        blocking_operator_count: shape.blocking_operator_count,
208        pipeline_bytes,
209        blocking_bytes,
210        external_read_bytes,
211        fixed_operator_bytes,
212        total_bytes: pipeline_bytes
213            .saturating_add(blocking_bytes)
214            .saturating_add(external_read_bytes)
215            .saturating_add(fixed_operator_bytes),
216    }
217}
218
219#[derive(Debug, Clone, Copy, PartialEq, Eq)]
220#[doc(hidden)]
221pub struct ExternalReadMemoryBudget {
222    pub max_working_bytes: NonZeroUsize,
223    pub max_result_bytes: NonZeroUsize,
224}
225
226impl ExternalReadMemoryBudget {
227    pub fn reserved_bytes(self) -> usize {
228        self.max_working_bytes
229            .get()
230            .saturating_add(self.max_result_bytes.get())
231    }
232}
233
234#[doc(hidden)]
235pub fn external_read_memory_budget(
236    profile: VectorExecutionResourceProfile,
237    memory: &ExecutionMemoryConfig,
238) -> ExternalReadMemoryBudget {
239    let configured_working_bytes = memory.blocking_operator_bytes.get();
240    let planned_working_bytes = profile
241        .max_working_memory_bytes
242        .map(|bytes| usize::try_from(bytes).unwrap_or(usize::MAX))
243        .unwrap_or(configured_working_bytes)
244        .max(1);
245    ExternalReadMemoryBudget {
246        max_working_bytes: NonZeroUsize::new(planned_working_bytes.min(configured_working_bytes))
247            .expect("configured external read working budget is non-zero"),
248        max_result_bytes: memory.blocking_operator_bytes,
249    }
250}
251
252#[doc(hidden)]
253pub fn max_external_read_parallelism(plan: &PhysicalPlan) -> usize {
254    let own_parallelism = match plan {
255        PhysicalPlan::VectorSeedScan {
256            resource_profile, ..
257        }
258        | PhysicalPlan::TextSeedScan {
259            resource_profile, ..
260        } => resource_profile.max_parallelism.max(1),
261        _ => 1,
262    };
263    match plan.children() {
264        PlanChildren::None => own_parallelism,
265        PlanChildren::Unary(input) => own_parallelism.max(max_external_read_parallelism(input)),
266        PlanChildren::Binary(left, right) => own_parallelism
267            .max(max_external_read_parallelism(left))
268            .max(max_external_read_parallelism(right)),
269    }
270}
271
272#[doc(hidden)]
273pub fn estimated_mutation_memory_bytes(
274    limits: MutationLimits,
275    max_wal_record_bytes: Option<usize>,
276) -> u64 {
277    let Some(max_wal_record_bytes) = max_wal_record_bytes else {
278        return u64::MAX;
279    };
280    usize_to_u64(max_wal_record_bytes)
281        .saturating_mul(2)
282        .saturating_add(
283            usize_to_u64(limits.max_operations.get())
284                .saturating_mul(MUTATION_OPERATION_BOOKKEEPING_BYTES),
285        )
286        .saturating_add(
287            usize_to_u64(limits.max_affected_rows.get())
288                .saturating_mul(MUTATION_AFFECTED_ROW_BOOKKEEPING_BYTES),
289        )
290        .saturating_add(
291            usize_to_u64(limits.max_result_rows.get())
292                .saturating_mul(MUTATION_RESULT_ROW_BOOKKEEPING_BYTES),
293        )
294        .saturating_add(usize_to_u64(limits.max_result_payload_bytes.get()))
295}
296
297#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
298struct ExecutionMemoryShape {
299    pipeline_batch_count: usize,
300    blocking_operator_count: usize,
301    external_read_bytes: u64,
302    fixed_operator_bytes: u64,
303}
304
305fn peak_execution_memory_shape(
306    plan: &PhysicalPlan,
307    memory: &ExecutionMemoryConfig,
308) -> ExecutionMemoryShape {
309    let mut shape = match plan.children() {
310        PlanChildren::None => ExecutionMemoryShape::default(),
311        PlanChildren::Unary(input) => peak_execution_memory_shape(input, memory),
312        PlanChildren::Binary(left, right) => peak_shape_max(
313            peak_execution_memory_shape(left, memory),
314            peak_execution_memory_shape(right, memory),
315            memory,
316        ),
317    };
318    shape.pipeline_batch_count = shape.pipeline_batch_count.saturating_add(1);
319    if retains_blocking_state(plan) {
320        shape.blocking_operator_count = shape.blocking_operator_count.saturating_add(1);
321    }
322    if matches!(plan, PhysicalPlan::SourceSegmentScan { .. }) {
323        shape.fixed_operator_bytes = shape
324            .fixed_operator_bytes
325            .saturating_add(SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES);
326    }
327    if matches!(plan, PhysicalPlan::VectorIndexScan(_)) {
328        shape.external_read_bytes = shape
329            .external_read_bytes
330            .saturating_add(usize_to_u64(memory.blocking_operator_bytes.get()).saturating_mul(2));
331    }
332    if let PhysicalPlan::VectorSeedScan {
333        resource_profile, ..
334    }
335    | PhysicalPlan::TextSeedScan {
336        resource_profile, ..
337    } = plan
338    {
339        shape.external_read_bytes = shape.external_read_bytes.saturating_add(usize_to_u64(
340            external_read_memory_budget(*resource_profile, memory).reserved_bytes(),
341        ));
342    }
343    shape
344}
345
346fn peak_shape_max(
347    left: ExecutionMemoryShape,
348    right: ExecutionMemoryShape,
349    memory: &ExecutionMemoryConfig,
350) -> ExecutionMemoryShape {
351    let shape_bytes = |shape: ExecutionMemoryShape| {
352        usize_to_u64(shape.pipeline_batch_count)
353            .saturating_mul(usize_to_u64(memory.batch_payload_bytes.get()))
354            .saturating_add(
355                usize_to_u64(shape.blocking_operator_count)
356                    .saturating_mul(usize_to_u64(memory.blocking_operator_bytes.get())),
357            )
358            .saturating_add(shape.external_read_bytes)
359            .saturating_add(shape.fixed_operator_bytes)
360    };
361    let left_total = shape_bytes(left);
362    let right_total = shape_bytes(right);
363    if left_total > right_total
364        || (left_total == right_total && left.fixed_operator_bytes >= right.fixed_operator_bytes)
365    {
366        left
367    } else {
368        right
369    }
370}
371
372fn retains_blocking_state(plan: &PhysicalPlan) -> bool {
373    matches!(
374        plan,
375        PhysicalPlan::GraphAlgorithm { .. }
376            | PhysicalPlan::GraphAlgorithmOutput(_)
377            | PhysicalPlan::VectorIndexScan(_)
378            | PhysicalPlan::VectorSeedScan { .. }
379            | PhysicalPlan::GraphSeedScan { .. }
380            | PhysicalPlan::TextSeedScan { .. }
381            | PhysicalPlan::SourceSegmentScan { .. }
382            | PhysicalPlan::NodeCartesianProductExec { .. }
383            | PhysicalPlan::HashJoinExec { .. }
384            | PhysicalPlan::AdjacencyExpandExec { .. }
385            | PhysicalPlan::OptionalRelationshipCountSumExec { .. }
386            | PhysicalPlan::ThreadRepairStatsExec { .. }
387            | PhysicalPlan::ShortestPathExec { .. }
388            | PhysicalPlan::AggregateExec { .. }
389            | PhysicalPlan::DistinctExec { .. }
390            | PhysicalPlan::SortExec { .. }
391            | PhysicalPlan::TopNExec { .. }
392            | PhysicalPlan::ScoringRerankExec { .. }
393            | PhysicalPlan::ScoringProgramExec { .. }
394            | PhysicalPlan::HostScoringExec { .. }
395    )
396}
397
398fn usize_to_u64(value: usize) -> u64 {
399    u64::try_from(value).unwrap_or(u64::MAX)
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use hawdb_core::RuntimeMemoryReservation;
406    use hawdb_plan_cypher::Predicate;
407
408    #[test]
409    fn result_memory_origin_tracks_the_effective_limit() {
410        let memory = ExecutionMemoryConfig {
411            query_memory_bytes: NonZeroUsize::new(128).unwrap(),
412            ..ExecutionMemoryConfig::default()
413        };
414        let ungoverned = RuntimeTaskContext::default();
415        for context in [None, Some(&ungoverned)] {
416            assert_eq!(
417                enforced_result_memory_budget_with_option(&memory, context).unwrap(),
418                (
419                    memory.query_memory_bytes,
420                    "ExecutionMemoryConfig.query_memory_bytes"
421                )
422            );
423        }
424        for (grant, expected, origin) in [
425            (64, 64, "RuntimeTaskContext.memory_reservation.result_bytes"),
426            (
427                128,
428                128,
429                "RuntimeTaskContext.memory_reservation.result_bytes",
430            ),
431            (256, 128, "ExecutionMemoryConfig.query_memory_bytes"),
432        ] {
433            let context = RuntimeTaskContext::default()
434                .with_memory_reservation(RuntimeMemoryReservation::new(1024, grant));
435            let (budget, option) =
436                enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap();
437            assert_eq!(budget.get(), expected);
438            assert_eq!(option, origin);
439            assert_eq!(
440                enforced_result_memory_budget(&memory, Some(&context)).unwrap(),
441                budget
442            );
443        }
444    }
445
446    #[test]
447    fn invalid_result_grants_do_not_fall_back_to_configuration() {
448        let memory = ExecutionMemoryConfig::default();
449        let context = RuntimeTaskContext::default()
450            .with_memory_reservation(RuntimeMemoryReservation::new(1024, 0));
451        let error = enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap_err();
452        assert!(matches!(error, HawDBError::Execution(_)));
453        assert!(error
454            .to_string()
455            .contains("query result reservation must be non-zero"));
456        if usize::BITS < u64::BITS {
457            let context = RuntimeTaskContext::default()
458                .with_memory_reservation(RuntimeMemoryReservation::new(1024, u64::MAX));
459            let error =
460                enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap_err();
461            assert!(matches!(error, HawDBError::Execution(_)));
462            assert!(error
463                .to_string()
464                .contains("does not fit the executor address space"));
465        }
466    }
467
468    fn admission_test_config() -> ExecutionMemoryConfig {
469        ExecutionMemoryConfig {
470            query_memory_bytes: NonZeroUsize::new(16 * 1024).unwrap(),
471            batch_rows: NonZeroUsize::new(8).unwrap(),
472            batch_payload_bytes: NonZeroUsize::new(1024).unwrap(),
473            blocking_operator_bytes: NonZeroUsize::new(4096).unwrap(),
474            graph_expansion_budget: None,
475            max_spill_bytes: NonZeroU64::new(1024 * 1024).unwrap(),
476            max_spill_runs: NonZeroUsize::new(8).unwrap(),
477            max_total_spill_bytes: NonZeroU64::new(4 * 1024 * 1024).unwrap(),
478            max_total_spill_runs: NonZeroUsize::new(32).unwrap(),
479            min_spill_free_bytes: NonZeroU64::new(1).unwrap(),
480            spill_free_space_probe_interval_bytes: NonZeroU64::new(1024).unwrap(),
481            spill_orphan_grace_period: Duration::from_secs(60),
482            spill_directory: std::env::temp_dir(),
483        }
484    }
485
486    #[test]
487    fn execution_admission_uses_configured_pipeline_and_blocking_budgets() {
488        let plan = PhysicalPlan::SortExec {
489            items: Vec::new(),
490            input: Box::new(PhysicalPlan::DistinctExec {
491                input: Box::new(PhysicalPlan::SeqNodeScan {
492                    variable: "n".to_string(),
493                    label: "Node".to_string(),
494                }),
495            }),
496        };
497
498        let estimate = estimated_execution_memory(&plan, &admission_test_config());
499
500        assert_eq!(estimate.pipeline_batch_count, 3);
501        assert_eq!(estimate.blocking_operator_count, 2);
502        assert_eq!(estimate.pipeline_bytes, 3 * 1024);
503        assert_eq!(estimate.blocking_bytes, 2 * 4096);
504        assert_eq!(estimate.external_read_bytes, 0);
505        assert_eq!(estimate.fixed_operator_bytes, 0);
506        assert_eq!(estimate.total_bytes, 11 * 1024);
507    }
508
509    #[test]
510    fn binary_admission_uses_the_higher_memory_child_path() {
511        let plan = PhysicalPlan::NodeCartesianProductExec {
512            left: Box::new(PhysicalPlan::ProjectExec {
513                items: Vec::new(),
514                input: Box::new(PhysicalPlan::ProjectExec {
515                    items: Vec::new(),
516                    input: Box::new(PhysicalPlan::SeqNodeScan {
517                        variable: "left".to_string(),
518                        label: "Left".to_string(),
519                    }),
520                }),
521            }),
522            right: Box::new(PhysicalPlan::DistinctExec {
523                input: Box::new(PhysicalPlan::SeqNodeScan {
524                    variable: "right".to_string(),
525                    label: "Right".to_string(),
526                }),
527            }),
528        };
529
530        let estimate = estimated_execution_memory(&plan, &admission_test_config());
531
532        assert_eq!(estimate.pipeline_batch_count, 3);
533        assert_eq!(estimate.blocking_operator_count, 2);
534        assert_eq!(estimate.total_bytes, 11 * 1024);
535    }
536
537    #[test]
538    fn source_segment_admission_includes_the_fixed_io_wave() {
539        let plan = PhysicalPlan::SourceSegmentScan {
540            variable: "n".to_string(),
541            predicate: Predicate::ConstantBool(true),
542        };
543
544        let estimate = estimated_execution_memory(&plan, &admission_test_config());
545
546        assert_eq!(estimate.pipeline_bytes, 1024);
547        assert_eq!(estimate.blocking_operator_count, 1);
548        assert_eq!(estimate.blocking_bytes, 4096);
549        assert_eq!(estimate.external_read_bytes, 0);
550        assert_eq!(
551            estimate.fixed_operator_bytes,
552            SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES
553        );
554        assert_eq!(
555            estimate.total_bytes,
556            1024 + 4096 + SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES
557        );
558    }
559
560    #[test]
561    fn vector_seed_admission_reserves_external_working_and_result_memory() {
562        let plan = PhysicalPlan::VectorSeedScan {
563            embedding_parameter: "embedding".to_string(),
564            output_external_id: true,
565            metadata_filters: Default::default(),
566            resource_profile: VectorExecutionResourceProfile {
567                priority: 128,
568                max_parallelism: 3,
569                max_working_memory_bytes: Some(2048),
570            },
571            vector_plan: hawdb_plan_cypher::VectorPhysicalPlan::TopK {
572                limit: 4,
573                input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::RawVectorRerank {
574                    embedding_dimension: 2,
575                    input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::VectorCandidateScan {
576                        source: hawdb_plan_cypher::VectorCandidateSource::Scalar,
577                        embedding_dimension: 2,
578                        candidate_limit: 4,
579                        input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::Filter {
580                            fields: Vec::new(),
581                        }),
582                    }),
583                }),
584            },
585        };
586
587        let estimate = estimated_execution_memory(&plan, &admission_test_config());
588
589        assert_eq!(estimate.pipeline_bytes, 1024);
590        assert_eq!(estimate.blocking_bytes, 4096);
591        assert_eq!(estimate.external_read_bytes, 2048 + 4096);
592        assert_eq!(estimate.total_bytes, 11 * 1024);
593        assert_eq!(max_external_read_parallelism(&plan), 3);
594    }
595
596    #[test]
597    fn mutation_admission_reserves_wal_staging_and_bounded_results() {
598        let limits = MutationLimits {
599            max_affected_rows: NonZeroUsize::new(3).unwrap(),
600            max_operations: NonZeroUsize::new(5).unwrap(),
601            max_result_rows: NonZeroUsize::new(7).unwrap(),
602            max_result_payload_bytes: NonZeroUsize::new(11).unwrap(),
603        };
604
605        assert_eq!(
606            estimated_mutation_memory_bytes(limits, Some(13)),
607            2 * 13
608                + 5 * MUTATION_OPERATION_BOOKKEEPING_BYTES
609                + 3 * MUTATION_AFFECTED_ROW_BOOKKEEPING_BYTES
610                + 7 * MUTATION_RESULT_ROW_BOOKKEEPING_BYTES
611                + 11
612        );
613        assert_eq!(estimated_mutation_memory_bytes(limits, None), u64::MAX);
614    }
615}