1use 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 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#[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 pub query_memory_bytes: NonZeroUsize,
108 pub batch_rows: NonZeroUsize,
110 pub batch_payload_bytes: NonZeroUsize,
112 pub blocking_operator_bytes: NonZeroUsize,
116 pub graph_expansion_budget: Option<GraphExpansionBudget>,
120 pub max_spill_bytes: NonZeroU64,
122 pub max_spill_runs: NonZeroUsize,
124 pub max_total_spill_bytes: NonZeroU64,
126 pub max_total_spill_runs: NonZeroUsize,
128 pub min_spill_free_bytes: NonZeroU64,
130 pub spill_free_space_probe_interval_bytes: NonZeroU64,
132 pub spill_orphan_grace_period: Duration,
134 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 #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
169 spill_directory: PathBuf::new(),
170 }
171 }
172}
173
174impl ExecutionMemoryConfig {
175 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}