hawdb_executor/
morsel.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//! Storage-neutral morsel partitioning and resource admission.
16
17use crate::concurrent::{BoundedOrderedStreamControl, BoundedOrderedStreamReport};
18use crate::{BoundedExecutor, QueryMemoryAccount, QueryMemoryLease, SharedExecutorPool};
19use hawdb_core::{HawDBError, Result, RuntimeTaskContext};
20use std::num::NonZeroUsize;
21use std::panic::{catch_unwind, AssertUnwindSafe};
22use std::sync::atomic::{AtomicUsize, Ordering};
23use std::sync::Arc;
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
26pub struct PipelineId(pub u32);
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
29pub struct MorselOrdinal(pub u64);
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct Morsel {
33    pub pipeline_id: PipelineId,
34    pub ordinal: MorselOrdinal,
35    pub start_row: usize,
36    pub row_count: usize,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub struct MorselAdmissionRequest {
41    pub pipeline_id: PipelineId,
42    pub input_rows: usize,
43    pub target_rows: NonZeroUsize,
44    pub requested_parallelism: NonZeroUsize,
45    pub bytes_per_worker: NonZeroUsize,
46    pub memory_budget_bytes: NonZeroUsize,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct MorselAdmission {
51    pipeline_id: PipelineId,
52    input_rows: usize,
53    target_rows: NonZeroUsize,
54    morsel_count: usize,
55    max_workers: usize,
56    reserved_bytes: usize,
57}
58
59impl MorselAdmission {
60    pub fn try_new(request: MorselAdmissionRequest) -> Result<Self> {
61        if request.input_rows == 0 {
62            return Ok(Self {
63                pipeline_id: request.pipeline_id,
64                input_rows: 0,
65                target_rows: request.target_rows,
66                morsel_count: 0,
67                max_workers: 0,
68                reserved_bytes: 0,
69            });
70        }
71
72        let memory_workers = request.memory_budget_bytes.get() / request.bytes_per_worker.get();
73        if memory_workers == 0 {
74            return Err(HawDBError::Execution(format!(
75                "morsel pipeline {} requires {} bytes for one worker, exceeding memory budget {}",
76                request.pipeline_id.0, request.bytes_per_worker, request.memory_budget_bytes
77            )));
78        }
79
80        let morsel_count = request.input_rows.div_ceil(request.target_rows.get());
81        let max_workers = request
82            .requested_parallelism
83            .get()
84            .min(memory_workers)
85            .min(morsel_count);
86        let reserved_bytes = max_workers.saturating_mul(request.bytes_per_worker.get());
87        Ok(Self {
88            pipeline_id: request.pipeline_id,
89            input_rows: request.input_rows,
90            target_rows: request.target_rows,
91            morsel_count,
92            max_workers,
93            reserved_bytes,
94        })
95    }
96
97    pub fn morsels(&self) -> MorselIter {
98        MorselIter {
99            pipeline_id: self.pipeline_id,
100            input_rows: self.input_rows,
101            target_rows: self.target_rows,
102            next_ordinal: 0,
103            morsel_count: self.morsel_count,
104        }
105    }
106
107    pub fn morsel_count(&self) -> usize {
108        self.morsel_count
109    }
110
111    pub fn max_workers(&self) -> usize {
112        self.max_workers
113    }
114
115    pub fn reserved_bytes(&self) -> usize {
116        self.reserved_bytes
117    }
118
119    fn morsel_at(&self, index: usize) -> Option<Morsel> {
120        if index >= self.morsel_count {
121            return None;
122        }
123        let start_row = index.saturating_mul(self.target_rows.get());
124        Some(Morsel {
125            pipeline_id: self.pipeline_id,
126            ordinal: MorselOrdinal(index as u64),
127            start_row,
128            row_count: self
129                .target_rows
130                .get()
131                .min(self.input_rows.saturating_sub(start_row)),
132        })
133    }
134}
135
136pub struct MorselIter {
137    pipeline_id: PipelineId,
138    input_rows: usize,
139    target_rows: NonZeroUsize,
140    next_ordinal: usize,
141    morsel_count: usize,
142}
143
144impl Iterator for MorselIter {
145    type Item = Morsel;
146
147    fn next(&mut self) -> Option<Self::Item> {
148        if self.next_ordinal == self.morsel_count {
149            return None;
150        }
151        let ordinal = self.next_ordinal;
152        let start_row = ordinal.saturating_mul(self.target_rows.get());
153        let row_count = self
154            .target_rows
155            .get()
156            .min(self.input_rows.saturating_sub(start_row));
157        self.next_ordinal = self.next_ordinal.saturating_add(1);
158        Some(Morsel {
159            pipeline_id: self.pipeline_id,
160            ordinal: MorselOrdinal(ordinal as u64),
161            start_row,
162            row_count,
163        })
164    }
165
166    fn size_hint(&self) -> (usize, Option<usize>) {
167        let remaining = self.morsel_count.saturating_sub(self.next_ordinal);
168        (remaining, Some(remaining))
169    }
170}
171
172impl ExactSizeIterator for MorselIter {}
173
174pub fn admit_morsels(request: MorselAdmissionRequest) -> Result<MorselAdmission> {
175    MorselAdmission::try_new(request)
176}
177
178/// Materializing sequential baseline for scheduler qualification and benchmarks.
179/// Production bounded output uses `SharedPoolMorselScheduler::execute_accounted_ordered`.
180#[derive(Debug, Default, Clone, Copy)]
181pub struct SequentialMorselScheduler;
182
183impl SequentialMorselScheduler {
184    pub fn execute<T>(
185        self,
186        admission: &MorselAdmission,
187        mut execute: impl FnMut(Morsel) -> Result<T>,
188    ) -> Result<Vec<T>> {
189        let mut output = Vec::with_capacity(admission.morsel_count);
190        for morsel in admission.morsels() {
191            output.push(execute(morsel)?);
192        }
193        Ok(output)
194    }
195}
196
197#[derive(Debug, Clone)]
198pub struct SharedPoolMorselScheduler {
199    pool: SharedExecutorPool,
200}
201
202impl SharedPoolMorselScheduler {
203    pub fn new(pool: SharedExecutorPool) -> Self {
204        Self { pool }
205    }
206
207    pub fn execute<T, F>(&self, admission: &MorselAdmission, execute: F) -> Result<Vec<T>>
208    where
209        T: Send,
210        F: Fn(Morsel) -> Result<T> + Sync,
211    {
212        if admission.morsel_count() == 0 {
213            return Ok(Vec::new());
214        }
215        let max_workers = NonZeroUsize::new(admission.max_workers()).ok_or_else(|| {
216            HawDBError::Execution("non-empty morsel admission reserved no workers".to_string())
217        })?;
218        let morsels = admission.morsels().collect::<Vec<_>>();
219        BoundedExecutor::with_pool(max_workers, self.pool.clone())
220            .map_ordered(&morsels, |morsel| execute_catching_panic(&execute, *morsel))
221            .into_iter()
222            .collect()
223    }
224
225    pub fn execute_with_context<T, F>(
226        &self,
227        admission: &MorselAdmission,
228        context: &RuntimeTaskContext,
229        execute: F,
230    ) -> Result<Vec<T>>
231    where
232        T: Send,
233        F: Fn(Morsel) -> Result<T> + Sync,
234    {
235        if admission.morsel_count() == 0 {
236            context.checkpoint().map_err(|reason| {
237                HawDBError::Execution(format!("runtime task stopped: {reason}"))
238            })?;
239            return Ok(Vec::new());
240        }
241        let max_workers = NonZeroUsize::new(admission.max_workers()).ok_or_else(|| {
242            HawDBError::Execution("non-empty morsel admission reserved no workers".to_string())
243        })?;
244        let morsels = admission.morsels().collect::<Vec<_>>();
245        BoundedExecutor::with_pool(max_workers, self.pool.clone())
246            .map_ordered_with_context(&morsels, context, |morsel| {
247                execute_catching_panic(&execute, *morsel)
248            })
249            .map_err(|reason| HawDBError::Execution(format!("runtime task stopped: {reason}")))?
250            .into_iter()
251            .collect()
252    }
253
254    /// Streams worker outputs through a bounded channel and invokes the
255    /// coordinator consumer in stable morsel ordinal order.
256    ///
257    /// `output_reservation_bytes` is charged before the worker calls
258    /// `execute`. The returned resident estimate must fit that reservation.
259    /// The lease remains live through `consume`, so a consumer that retains
260    /// the value must transfer it into another query-owned account.
261    pub fn execute_accounted_ordered<T, F, C>(
262        &self,
263        admission: &MorselAdmission,
264        resources: MorselStreamResources<'_>,
265        execute: F,
266        mut consume: C,
267    ) -> Result<MorselStreamReport>
268    where
269        T: Send,
270        F: Fn(Morsel) -> Result<MorselOutput<T>> + Sync,
271        C: FnMut(Morsel, T) -> Result<MorselStreamControl>,
272    {
273        if admission.morsel_count() == 0 {
274            if let Some(context) = resources.task_context {
275                context.checkpoint().map_err(|reason| {
276                    HawDBError::Execution(format!("runtime task stopped: {reason}"))
277                })?;
278            }
279            return Ok(MorselStreamReport::default());
280        }
281        let max_workers = NonZeroUsize::new(admission.max_workers()).ok_or_else(|| {
282            HawDBError::Execution("non-empty morsel admission reserved no workers".to_string())
283        })?;
284        let counters = Arc::new(MorselOutputCounters::default());
285        let concurrent_report = BoundedExecutor::with_pool(max_workers, self.pool.clone())
286            .try_for_each_index_ordered(
287                admission.morsel_count(),
288                resources.task_context,
289                |index| {
290                    let morsel = admission
291                        .morsel_at(index)
292                        .expect("admitted morsel index is in range");
293                    let mut lease = resources
294                        .output_account
295                        .reserve(resources.output_reservation_bytes.get())?;
296                    let output = execute_catching_panic(&execute, morsel)?;
297                    if output.resident_bytes > resources.output_reservation_bytes.get() {
298                        return Err(HawDBError::Execution(format!(
299                            "morsel pipeline {} output at ordinal {} uses {} bytes, exceeding its {}-byte reservation",
300                            morsel.pipeline_id.0,
301                            morsel.ordinal.0,
302                            output.resident_bytes,
303                            resources.output_reservation_bytes
304                        )));
305                    }
306                    lease.shrink(
307                        resources.output_reservation_bytes.get() - output.resident_bytes,
308                    );
309                    Ok(PendingMorselOutput::new(
310                        output.value,
311                        output.resident_bytes,
312                        lease,
313                        Arc::clone(&counters),
314                    ))
315                },
316                |index, mut output| {
317                    let morsel = admission
318                        .morsel_at(index)
319                        .expect("admitted morsel index is in range");
320                    let value = output.take_value();
321                    let control = catch_unwind(AssertUnwindSafe(|| consume(morsel, value)))
322                        .unwrap_or_else(|_| {
323                            Err(HawDBError::Execution(format!(
324                                "morsel pipeline {} consumer panicked at ordinal {}",
325                                morsel.pipeline_id.0, morsel.ordinal.0
326                            )))
327                        })?;
328                    Ok(match control {
329                        MorselStreamControl::Continue => BoundedOrderedStreamControl::Continue,
330                        MorselStreamControl::Stop => BoundedOrderedStreamControl::Stop,
331                    })
332                },
333            )?;
334        Ok(MorselStreamReport::from_parts(concurrent_report, &counters))
335    }
336}
337
338#[derive(Debug, Clone, Copy)]
339pub struct MorselStreamResources<'a> {
340    pub task_context: Option<&'a RuntimeTaskContext>,
341    pub output_account: &'a QueryMemoryAccount,
342    pub output_reservation_bytes: NonZeroUsize,
343}
344
345#[derive(Debug)]
346pub struct MorselOutput<T> {
347    value: T,
348    resident_bytes: usize,
349}
350
351impl<T> MorselOutput<T> {
352    pub fn new(value: T, resident_bytes: usize) -> Self {
353        Self {
354            value,
355            resident_bytes,
356        }
357    }
358}
359
360#[derive(Debug, Clone, Copy, PartialEq, Eq)]
361pub enum MorselStreamControl {
362    Continue,
363    Stop,
364}
365
366#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
367pub struct MorselStreamReport {
368    pub peak_buffered_outputs: usize,
369    pub peak_buffered_output_bytes: usize,
370    pub peak_reorder_entries: usize,
371    pub stopped_early: bool,
372}
373
374impl MorselStreamReport {
375    fn from_parts(concurrent: BoundedOrderedStreamReport, counters: &MorselOutputCounters) -> Self {
376        Self {
377            peak_buffered_outputs: counters.peak_outputs.load(Ordering::Relaxed),
378            peak_buffered_output_bytes: counters.peak_bytes.load(Ordering::Relaxed),
379            peak_reorder_entries: concurrent.peak_reorder_entries,
380            stopped_early: concurrent.stopped_early,
381        }
382    }
383}
384
385#[derive(Debug, Default)]
386struct MorselOutputCounters {
387    current_outputs: AtomicUsize,
388    peak_outputs: AtomicUsize,
389    current_bytes: AtomicUsize,
390    peak_bytes: AtomicUsize,
391}
392
393struct PendingMorselOutput<T> {
394    value: Option<T>,
395    resident_bytes: usize,
396    _lease: QueryMemoryLease,
397    counters: Arc<MorselOutputCounters>,
398}
399
400impl<T> PendingMorselOutput<T> {
401    fn new(
402        value: T,
403        resident_bytes: usize,
404        lease: QueryMemoryLease,
405        counters: Arc<MorselOutputCounters>,
406    ) -> Self {
407        let outputs = counters.current_outputs.fetch_add(1, Ordering::Relaxed) + 1;
408        counters.peak_outputs.fetch_max(outputs, Ordering::Relaxed);
409        let bytes = counters
410            .current_bytes
411            .fetch_add(resident_bytes, Ordering::Relaxed)
412            .saturating_add(resident_bytes);
413        counters.peak_bytes.fetch_max(bytes, Ordering::Relaxed);
414        Self {
415            value: Some(value),
416            resident_bytes,
417            _lease: lease,
418            counters,
419        }
420    }
421
422    fn take_value(&mut self) -> T {
423        self.value
424            .take()
425            .expect("pending morsel output value is consumed once")
426    }
427}
428
429impl<T> Drop for PendingMorselOutput<T> {
430    fn drop(&mut self) {
431        self.counters
432            .current_outputs
433            .fetch_sub(1, Ordering::Relaxed);
434        self.counters
435            .current_bytes
436            .fetch_sub(self.resident_bytes, Ordering::Relaxed);
437    }
438}
439
440fn execute_catching_panic<T>(
441    execute: &(impl Fn(Morsel) -> Result<T> + Sync),
442    morsel: Morsel,
443) -> Result<T> {
444    catch_unwind(AssertUnwindSafe(|| execute(morsel))).unwrap_or_else(|_| {
445        Err(HawDBError::Execution(format!(
446            "morsel pipeline {} worker panicked at ordinal {}",
447            morsel.pipeline_id.0, morsel.ordinal.0
448        )))
449    })
450}
451
452/// Executes morsels in ordinal order, stopping at the first callback error.
453///
454/// This materializing baseline is consumed by the executor vectorization
455/// benchmark and `tests/morsel_oracle.rs`, which qualifies both materializing
456/// shared-pool methods and the production accounted stream against an independent
457/// row oracle. Parallel callbacks may already be running when one fails; their
458/// side effects need not match the sequential baseline's fail-fast behavior.
459pub fn execute_morsels_ordered<T>(
460    admission: &MorselAdmission,
461    mut execute: impl FnMut(Morsel) -> Result<T>,
462) -> Result<Vec<T>> {
463    SequentialMorselScheduler.execute(admission, &mut execute)
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469    use crate::{QueryMemoryClass, QueryMemoryLedger};
470
471    fn request(input_rows: usize) -> MorselAdmissionRequest {
472        MorselAdmissionRequest {
473            pipeline_id: PipelineId(7),
474            input_rows,
475            target_rows: NonZeroUsize::new(64).unwrap(),
476            requested_parallelism: NonZeroUsize::new(8).unwrap(),
477            bytes_per_worker: NonZeroUsize::new(1024).unwrap(),
478            memory_budget_bytes: NonZeroUsize::new(3 * 1024).unwrap(),
479        }
480    }
481
482    #[test]
483    fn admission_is_bounded_by_work_and_memory() {
484        let admission = admit_morsels(request(130)).unwrap();
485
486        assert_eq!(admission.max_workers(), 3);
487        assert_eq!(admission.reserved_bytes(), 3 * 1024);
488        assert_eq!(
489            admission
490                .morsels()
491                .map(|morsel| (morsel.start_row, morsel.row_count))
492                .collect::<Vec<_>>(),
493            vec![(0, 64), (64, 64), (128, 2)]
494        );
495    }
496
497    #[test]
498    fn admission_rejects_a_worker_that_cannot_fit() {
499        let mut request = request(1);
500        request.bytes_per_worker = NonZeroUsize::new(4096).unwrap();
501
502        assert!(admit_morsels(request).is_err());
503    }
504
505    #[test]
506    fn ordered_executor_preserves_morsel_ordinals() {
507        let admission = MorselAdmission::try_new(request(130)).unwrap();
508        let ordinals = SequentialMorselScheduler
509            .execute(&admission, |morsel| Ok(morsel.ordinal.0))
510            .unwrap();
511
512        assert_eq!(ordinals, vec![0, 1, 2]);
513    }
514
515    #[test]
516    fn empty_input_reserves_no_workers() {
517        let admission = admit_morsels(request(0)).unwrap();
518
519        assert_eq!(admission.morsel_count(), 0);
520        assert_eq!(admission.max_workers(), 0);
521        assert_eq!(admission.reserved_bytes(), 0);
522    }
523
524    #[test]
525    fn shared_pool_scheduler_matches_sequential_ordinal_order() {
526        let admission = admit_morsels(request(1025)).unwrap();
527        let sequential = execute_morsels_ordered(&admission, |morsel| {
528            Ok((morsel.ordinal.0, morsel.start_row, morsel.row_count))
529        })
530        .unwrap();
531        let pool = SharedExecutorPool::new(NonZeroUsize::new(3).unwrap()).unwrap();
532        let parallel = SharedPoolMorselScheduler::new(pool)
533            .execute(&admission, |morsel| {
534                std::thread::yield_now();
535                Ok((morsel.ordinal.0, morsel.start_row, morsel.row_count))
536            })
537            .unwrap();
538
539        assert_eq!(parallel, sequential);
540    }
541
542    #[test]
543    fn shared_pool_scheduler_converts_worker_panics_to_stable_errors() {
544        let admission = admit_morsels(request(130)).unwrap();
545        let pool = SharedExecutorPool::new(NonZeroUsize::new(2).unwrap()).unwrap();
546        let error = SharedPoolMorselScheduler::new(pool)
547            .execute(&admission, |morsel| -> Result<u64> {
548                assert_ne!(morsel.ordinal.0, 1, "injected worker panic");
549                Ok(morsel.ordinal.0)
550            })
551            .unwrap_err();
552
553        assert!(error.to_string().contains("ordinal 1"), "{error}");
554    }
555
556    #[test]
557    fn shared_pool_scheduler_observes_cancellation_between_morsels() {
558        let admission = admit_morsels(request(1025)).unwrap();
559        let pool = SharedExecutorPool::new(NonZeroUsize::MIN).unwrap();
560        let token = hawdb_core::RuntimeCancellationToken::new();
561        let context = RuntimeTaskContext::without_deadline(token.clone());
562        let error = SharedPoolMorselScheduler::new(pool)
563            .execute_with_context(&admission, &context, |morsel| {
564                if morsel.ordinal.0 == 0 {
565                    token.cancel();
566                }
567                Ok(morsel.ordinal.0)
568            })
569            .unwrap_err();
570
571        assert!(error.to_string().contains("cancelled"));
572    }
573
574    #[test]
575    fn accounted_stream_preserves_order_and_bounds_completed_outputs() {
576        let admission = admit_morsels(request(320)).unwrap();
577        let pool = SharedExecutorPool::new(NonZeroUsize::new(3).unwrap()).unwrap();
578        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
579        let account = ledger.account(
580            QueryMemoryClass::MorselOutput,
581            "test reorder",
582            NonZeroUsize::new(96).unwrap(),
583        );
584        let later_completed = AtomicUsize::new(0);
585        let mut ordinals = Vec::new();
586
587        let report = SharedPoolMorselScheduler::new(pool)
588            .execute_accounted_ordered(
589                &admission,
590                MorselStreamResources {
591                    task_context: None,
592                    output_account: &account,
593                    output_reservation_bytes: NonZeroUsize::new(32).unwrap(),
594                },
595                |morsel| {
596                    if morsel.ordinal.0 == 0 {
597                        while later_completed.load(Ordering::Acquire) < 2 {
598                            std::thread::yield_now();
599                        }
600                    } else if morsel.ordinal.0 <= 2 {
601                        later_completed.fetch_add(1, Ordering::Release);
602                    }
603                    Ok(MorselOutput::new(morsel.ordinal.0, 8))
604                },
605                |morsel, value| {
606                    assert_eq!(morsel.ordinal.0, value);
607                    ordinals.push(value);
608                    Ok(MorselStreamControl::Continue)
609                },
610            )
611            .unwrap();
612
613        assert_eq!(ordinals, vec![0, 1, 2, 3, 4]);
614        assert!((1..=admission.max_workers()).contains(&report.peak_buffered_outputs));
615        assert!(report.peak_reorder_entries <= admission.max_workers());
616        assert!(report.peak_buffered_output_bytes <= admission.max_workers() * 8);
617        assert_eq!(ledger.snapshot().used_bytes, 0);
618    }
619
620    #[test]
621    fn accounted_stream_reserves_before_worker_execution() {
622        let admission = admit_morsels(request(130)).unwrap();
623        let pool = SharedExecutorPool::new(NonZeroUsize::new(2).unwrap()).unwrap();
624        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
625        let account = ledger.account(
626            QueryMemoryClass::MorselOutput,
627            "test reservation",
628            NonZeroUsize::new(16).unwrap(),
629        );
630        let executed = AtomicUsize::new(0);
631
632        let error = SharedPoolMorselScheduler::new(pool)
633            .execute_accounted_ordered(
634                &admission,
635                MorselStreamResources {
636                    task_context: None,
637                    output_account: &account,
638                    output_reservation_bytes: NonZeroUsize::new(32).unwrap(),
639                },
640                |morsel| {
641                    executed.fetch_add(1, Ordering::SeqCst);
642                    Ok(MorselOutput::new(morsel.ordinal.0, 8))
643                },
644                |_morsel, _value| Ok(MorselStreamControl::Continue),
645            )
646            .unwrap_err();
647
648        assert!(error.to_string().contains("16-byte budget"));
649        assert_eq!(executed.load(Ordering::SeqCst), 0);
650        assert_eq!(ledger.snapshot().used_bytes, 0);
651    }
652
653    #[test]
654    fn accounted_stream_releases_outputs_after_worker_panic() {
655        let admission = admit_morsels(request(130)).unwrap();
656        let pool = SharedExecutorPool::new(NonZeroUsize::new(2).unwrap()).unwrap();
657        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
658        let account = ledger.account(
659            QueryMemoryClass::MorselOutput,
660            "test panic",
661            NonZeroUsize::new(64).unwrap(),
662        );
663
664        let error = SharedPoolMorselScheduler::new(pool)
665            .execute_accounted_ordered(
666                &admission,
667                MorselStreamResources {
668                    task_context: None,
669                    output_account: &account,
670                    output_reservation_bytes: NonZeroUsize::new(32).unwrap(),
671                },
672                |morsel| {
673                    assert_ne!(morsel.ordinal.0, 1, "injected worker panic");
674                    Ok(MorselOutput::new(morsel.ordinal.0, 8))
675                },
676                |_morsel, _value| Ok(MorselStreamControl::Continue),
677            )
678            .unwrap_err();
679
680        assert!(error.to_string().contains("ordinal 1"), "{error}");
681        assert_eq!(ledger.snapshot().used_bytes, 0);
682    }
683
684    #[test]
685    fn accounted_stream_stop_does_not_open_another_issue_window() {
686        let admission = admit_morsels(request(640)).unwrap();
687        let pool = SharedExecutorPool::new(NonZeroUsize::new(3).unwrap()).unwrap();
688        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
689        let account = ledger.account(
690            QueryMemoryClass::MorselOutput,
691            "test stop",
692            NonZeroUsize::new(96).unwrap(),
693        );
694        let executed = AtomicUsize::new(0);
695
696        let report = SharedPoolMorselScheduler::new(pool)
697            .execute_accounted_ordered(
698                &admission,
699                MorselStreamResources {
700                    task_context: None,
701                    output_account: &account,
702                    output_reservation_bytes: NonZeroUsize::new(32).unwrap(),
703                },
704                |morsel| {
705                    executed.fetch_add(1, Ordering::SeqCst);
706                    Ok(MorselOutput::new(morsel.ordinal.0, 8))
707                },
708                |_morsel, _value| Ok(MorselStreamControl::Stop),
709            )
710            .unwrap();
711
712        assert!(report.stopped_early);
713        assert!(executed.load(Ordering::SeqCst) <= admission.max_workers());
714        assert_eq!(ledger.snapshot().used_bytes, 0);
715    }
716
717    #[test]
718    fn accounted_stream_converts_consumer_panics_and_releases_outputs() {
719        let admission = admit_morsels(request(130)).unwrap();
720        let pool = SharedExecutorPool::new(NonZeroUsize::new(2).unwrap()).unwrap();
721        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
722        let account = ledger.account(
723            QueryMemoryClass::MorselOutput,
724            "test consumer panic",
725            NonZeroUsize::new(64).unwrap(),
726        );
727
728        let error = SharedPoolMorselScheduler::new(pool)
729            .execute_accounted_ordered(
730                &admission,
731                MorselStreamResources {
732                    task_context: None,
733                    output_account: &account,
734                    output_reservation_bytes: NonZeroUsize::new(32).unwrap(),
735                },
736                |morsel| Ok(MorselOutput::new(morsel.ordinal.0, 8)),
737                |morsel, _value| -> Result<MorselStreamControl> {
738                    assert_ne!(morsel.ordinal.0, 0, "injected consumer panic");
739                    Ok(MorselStreamControl::Continue)
740                },
741            )
742            .unwrap_err();
743
744        assert!(error.to_string().contains("consumer panicked at ordinal 0"));
745        assert_eq!(ledger.snapshot().used_bytes, 0);
746    }
747}
hawdb_executor/morsel.rs

hawdb_executor/
morsel.rs
