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    /// Override the row and cumulative payload limits of each seeded expansion.
117    /// Live allocations still obey the operator, batch and shared query budgets.
118    /// `None` uses the optimizer's seed-window-derived limits.
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 let PhysicalPlan::VectorSeedScan {
328        resource_profile, ..
329    }
330    | PhysicalPlan::TextSeedScan {
331        resource_profile, ..
332    } = plan
333    {
334        shape.external_read_bytes = shape.external_read_bytes.saturating_add(usize_to_u64(
335            external_read_memory_budget(*resource_profile, memory).reserved_bytes(),
336        ));
337    }
338    shape
339}
340
341fn peak_shape_max(
342    left: ExecutionMemoryShape,
343    right: ExecutionMemoryShape,
344    memory: &ExecutionMemoryConfig,
345) -> ExecutionMemoryShape {
346    let shape_bytes = |shape: ExecutionMemoryShape| {
347        usize_to_u64(shape.pipeline_batch_count)
348            .saturating_mul(usize_to_u64(memory.batch_payload_bytes.get()))
349            .saturating_add(
350                usize_to_u64(shape.blocking_operator_count)
351                    .saturating_mul(usize_to_u64(memory.blocking_operator_bytes.get())),
352            )
353            .saturating_add(shape.external_read_bytes)
354            .saturating_add(shape.fixed_operator_bytes)
355    };
356    let left_total = shape_bytes(left);
357    let right_total = shape_bytes(right);
358    if left_total > right_total
359        || (left_total == right_total && left.fixed_operator_bytes >= right.fixed_operator_bytes)
360    {
361        left
362    } else {
363        right
364    }
365}
366
367fn retains_blocking_state(plan: &PhysicalPlan) -> bool {
368    matches!(
369        plan,
370        PhysicalPlan::GraphAlgorithm { .. }
371            | PhysicalPlan::VectorSeedScan { .. }
372            | PhysicalPlan::GraphSeedScan { .. }
373            | PhysicalPlan::TextSeedScan { .. }
374            | PhysicalPlan::SourceSegmentScan { .. }
375            | PhysicalPlan::NodeCartesianProductExec { .. }
376            | PhysicalPlan::HashJoinExec { .. }
377            | PhysicalPlan::AdjacencyExpandExec { .. }
378            | PhysicalPlan::OptionalRelationshipCountSumExec { .. }
379            | PhysicalPlan::ThreadRepairStatsExec { .. }
380            | PhysicalPlan::ShortestPathExec { .. }
381            | PhysicalPlan::AggregateExec { .. }
382            | PhysicalPlan::DistinctExec { .. }
383            | PhysicalPlan::SortExec { .. }
384            | PhysicalPlan::TopNExec { .. }
385            | PhysicalPlan::ScoringRerankExec { .. }
386            | PhysicalPlan::ScoringProgramExec { .. }
387            | PhysicalPlan::HostScoringExec { .. }
388    )
389}
390
391fn usize_to_u64(value: usize) -> u64 {
392    u64::try_from(value).unwrap_or(u64::MAX)
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398    use hawdb_core::RuntimeMemoryReservation;
399    use hawdb_plan_cypher::Predicate;
400
401    #[test]
402    fn result_memory_origin_tracks_the_effective_limit() {
403        let memory = ExecutionMemoryConfig {
404            query_memory_bytes: NonZeroUsize::new(128).unwrap(),
405            ..ExecutionMemoryConfig::default()
406        };
407        let ungoverned = RuntimeTaskContext::default();
408        for context in [None, Some(&ungoverned)] {
409            assert_eq!(
410                enforced_result_memory_budget_with_option(&memory, context).unwrap(),
411                (
412                    memory.query_memory_bytes,
413                    "ExecutionMemoryConfig.query_memory_bytes"
414                )
415            );
416        }
417        for (grant, expected, origin) in [
418            (64, 64, "RuntimeTaskContext.memory_reservation.result_bytes"),
419            (
420                128,
421                128,
422                "RuntimeTaskContext.memory_reservation.result_bytes",
423            ),
424            (256, 128, "ExecutionMemoryConfig.query_memory_bytes"),
425        ] {
426            let context = RuntimeTaskContext::default()
427                .with_memory_reservation(RuntimeMemoryReservation::new(1024, grant));
428            let (budget, option) =
429                enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap();
430            assert_eq!(budget.get(), expected);
431            assert_eq!(option, origin);
432            assert_eq!(
433                enforced_result_memory_budget(&memory, Some(&context)).unwrap(),
434                budget
435            );
436        }
437    }
438
439    #[test]
440    fn invalid_result_grants_do_not_fall_back_to_configuration() {
441        let memory = ExecutionMemoryConfig::default();
442        let context = RuntimeTaskContext::default()
443            .with_memory_reservation(RuntimeMemoryReservation::new(1024, 0));
444        let error = enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap_err();
445        assert!(matches!(error, HawDBError::Execution(_)));
446        assert!(error
447            .to_string()
448            .contains("query result reservation must be non-zero"));
449        if usize::BITS < u64::BITS {
450            let context = RuntimeTaskContext::default()
451                .with_memory_reservation(RuntimeMemoryReservation::new(1024, u64::MAX));
452            let error =
453                enforced_result_memory_budget_with_option(&memory, Some(&context)).unwrap_err();
454            assert!(matches!(error, HawDBError::Execution(_)));
455            assert!(error
456                .to_string()
457                .contains("does not fit the executor address space"));
458        }
459    }
460
461    fn admission_test_config() -> ExecutionMemoryConfig {
462        ExecutionMemoryConfig {
463            query_memory_bytes: NonZeroUsize::new(16 * 1024).unwrap(),
464            batch_rows: NonZeroUsize::new(8).unwrap(),
465            batch_payload_bytes: NonZeroUsize::new(1024).unwrap(),
466            blocking_operator_bytes: NonZeroUsize::new(4096).unwrap(),
467            graph_expansion_budget: None,
468            max_spill_bytes: NonZeroU64::new(1024 * 1024).unwrap(),
469            max_spill_runs: NonZeroUsize::new(8).unwrap(),
470            max_total_spill_bytes: NonZeroU64::new(4 * 1024 * 1024).unwrap(),
471            max_total_spill_runs: NonZeroUsize::new(32).unwrap(),
472            min_spill_free_bytes: NonZeroU64::new(1).unwrap(),
473            spill_free_space_probe_interval_bytes: NonZeroU64::new(1024).unwrap(),
474            spill_orphan_grace_period: Duration::from_secs(60),
475            spill_directory: std::env::temp_dir(),
476        }
477    }
478
479    #[test]
480    fn execution_admission_uses_configured_pipeline_and_blocking_budgets() {
481        let plan = PhysicalPlan::SortExec {
482            items: Vec::new(),
483            input: Box::new(PhysicalPlan::DistinctExec {
484                input: Box::new(PhysicalPlan::SeqNodeScan {
485                    variable: "n".to_string(),
486                    label: "Node".to_string(),
487                }),
488            }),
489        };
490
491        let estimate = estimated_execution_memory(&plan, &admission_test_config());
492
493        assert_eq!(estimate.pipeline_batch_count, 3);
494        assert_eq!(estimate.blocking_operator_count, 2);
495        assert_eq!(estimate.pipeline_bytes, 3 * 1024);
496        assert_eq!(estimate.blocking_bytes, 2 * 4096);
497        assert_eq!(estimate.external_read_bytes, 0);
498        assert_eq!(estimate.fixed_operator_bytes, 0);
499        assert_eq!(estimate.total_bytes, 11 * 1024);
500    }
501
502    #[test]
503    fn binary_admission_uses_the_higher_memory_child_path() {
504        let plan = PhysicalPlan::NodeCartesianProductExec {
505            left: Box::new(PhysicalPlan::ProjectExec {
506                items: Vec::new(),
507                input: Box::new(PhysicalPlan::ProjectExec {
508                    items: Vec::new(),
509                    input: Box::new(PhysicalPlan::SeqNodeScan {
510                        variable: "left".to_string(),
511                        label: "Left".to_string(),
512                    }),
513                }),
514            }),
515            right: Box::new(PhysicalPlan::DistinctExec {
516                input: Box::new(PhysicalPlan::SeqNodeScan {
517                    variable: "right".to_string(),
518                    label: "Right".to_string(),
519                }),
520            }),
521        };
522
523        let estimate = estimated_execution_memory(&plan, &admission_test_config());
524
525        assert_eq!(estimate.pipeline_batch_count, 3);
526        assert_eq!(estimate.blocking_operator_count, 2);
527        assert_eq!(estimate.total_bytes, 11 * 1024);
528    }
529
530    #[test]
531    fn source_segment_admission_includes_the_fixed_io_wave() {
532        let plan = PhysicalPlan::SourceSegmentScan {
533            variable: "n".to_string(),
534            predicate: Predicate::ConstantBool(true),
535        };
536
537        let estimate = estimated_execution_memory(&plan, &admission_test_config());
538
539        assert_eq!(estimate.pipeline_bytes, 1024);
540        assert_eq!(estimate.blocking_operator_count, 1);
541        assert_eq!(estimate.blocking_bytes, 4096);
542        assert_eq!(estimate.external_read_bytes, 0);
543        assert_eq!(
544            estimate.fixed_operator_bytes,
545            SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES
546        );
547        assert_eq!(
548            estimate.total_bytes,
549            1024 + 4096 + SOURCE_SEGMENT_SCAN_MAX_WAVE_BYTES
550        );
551    }
552
553    #[test]
554    fn vector_seed_admission_reserves_external_working_and_result_memory() {
555        let plan = PhysicalPlan::VectorSeedScan {
556            embedding_parameter: "embedding".to_string(),
557            output_external_id: true,
558            metadata_filters: Default::default(),
559            resource_profile: VectorExecutionResourceProfile {
560                priority: 128,
561                max_parallelism: 3,
562                max_working_memory_bytes: Some(2048),
563            },
564            vector_plan: hawdb_plan_cypher::VectorPhysicalPlan::TopK {
565                limit: 4,
566                input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::RawVectorRerank {
567                    embedding_dimension: 2,
568                    input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::VectorCandidateScan {
569                        source: hawdb_plan_cypher::VectorCandidateSource::Scalar,
570                        embedding_dimension: 2,
571                        candidate_limit: 4,
572                        input: Box::new(hawdb_plan_cypher::VectorPhysicalPlan::Filter {
573                            fields: Vec::new(),
574                        }),
575                    }),
576                }),
577            },
578        };
579
580        let estimate = estimated_execution_memory(&plan, &admission_test_config());
581
582        assert_eq!(estimate.pipeline_bytes, 1024);
583        assert_eq!(estimate.blocking_bytes, 4096);
584        assert_eq!(estimate.external_read_bytes, 2048 + 4096);
585        assert_eq!(estimate.total_bytes, 11 * 1024);
586        assert_eq!(max_external_read_parallelism(&plan), 3);
587    }
588
589    #[test]
590    fn mutation_admission_reserves_wal_staging_and_bounded_results() {
591        let limits = MutationLimits {
592            max_affected_rows: NonZeroUsize::new(3).unwrap(),
593            max_operations: NonZeroUsize::new(5).unwrap(),
594            max_result_rows: NonZeroUsize::new(7).unwrap(),
595            max_result_payload_bytes: NonZeroUsize::new(11).unwrap(),
596        };
597
598        assert_eq!(
599            estimated_mutation_memory_bytes(limits, Some(13)),
600            2 * 13
601                + 5 * MUTATION_OPERATION_BOOKKEEPING_BYTES
602                + 3 * MUTATION_AFFECTED_ROW_BOOKKEEPING_BYTES
603                + 7 * MUTATION_RESULT_ROW_BOOKKEEPING_BYTES
604                + 11
605        );
606        assert_eq!(estimated_mutation_memory_bytes(limits, None), u64::MAX);
607    }
608}