Skip to main content

mongreldb_core/
spill.rs

1//! Spill manager for query execution (spec section 10.5, S1E-004).
2//!
3//! Implemented in the Stage 1E wave: when the node
4//! [`MemoryGovernor`](crate::memory::MemoryGovernor) enters escalation step 3,
5//! spill-eligible operators move working memory to disk through this manager.
6//! Per the spec, spill files are:
7//!
8//! - **Query-ID namespaced** — every [`QueryId`] spills under its own
9//!   subdirectory `temp/spill/q-<hex>/` of the database root, so one query's
10//!   files never interleave with another's and cancellation removes exactly
11//!   one directory.
12//! - **Checksummed** — every frame carries a CRC32C (the same Castagnoli CRC
13//!   as the WAL) over its kind, sequence, and stored payload; the sealing
14//!   trailer carries a SHA-256 over all plaintext data payloads plus the
15//!   frame count, verified on read.
16//! - **Bounded** — a per-query cap (fed from
17//!   [`ResourceGroup::temporary_disk_bytes`]) and a node-global cap, enforced
18//!   with the memory governor's add-then-validate-rollback protocol; overflow
19//!   is the typed [`SpillError::BudgetExceeded`].
20//! - **Deleted on success/error/cancel/startup cleanup** — the [`SpillHandle`]
21//!   RAII guard deletes its file on drop, an unfinished [`SpillWriter`]
22//!   deletes its partial file on drop (the error/cancel path), a dropped
23//!   [`SpillSession`] removes the whole per-query directory, and
24//!   [`SpillManager::open`] sweeps every stale entry left by a prior process
25//!   run (spill files never outlive the process that created them).
26//! - **Encrypted when database encryption is enabled** — with a meta DEK
27//!   present ([`crate::encryption::meta_dek_for`]) every frame payload is
28//!   sealed through the page-cipher stack (`encrypt_blob`: AES-256-GCM with a
29//!   fresh random nonce per frame), mirroring the catalog/PITR seal idiom;
30//!   otherwise frames are stored plaintext.
31//!
32//! Layout of one spill file:
33//!
34//! ```text
35//! [magic: 8B "MDBSPILL"][version: u16][enc flag: u8][reserved: u8]
36//! frame*: [payload len: u32][crc32c: u32][kind: u8][seq: u64][payload]
37//! trailer frame (kind = 1): [data frames: u64][data bytes: u64][sha256: 32B]
38//! ```
39//!
40//! The on-disk tree lives under a descriptor-pinned
41//! [`DurableRoot`](crate::durable_file::DurableRoot); every file operation is
42//! descriptor-relative and `temp/spill` itself is created lazily on first
43//! use. Open one `SpillManager` per database: opening a second manager on the
44//! same root sweeps the first manager's files, exactly like a restart.
45
46use std::fmt;
47use std::io::{self, Read, Write};
48use std::path::Path;
49use std::sync::atomic::{AtomicU64, Ordering};
50use std::sync::Arc;
51
52use crc::{Crc, CRC_32_ISCSI};
53use mongreldb_types::ids::QueryId;
54use sha2::{Digest, Sha256};
55
56use crate::durable_file::DurableRoot;
57use crate::encryption::DEK_LEN;
58use crate::resource::ResourceGroup;
59use crate::MongrelError;
60
61/// Spill tree location relative to the database root.
62const SPILL_DIR_REL: &str = "temp/spill";
63/// File magic (8 bytes, the [`MongrelError::MagicMismatch`] idiom).
64const SPILL_MAGIC: &[u8; 8] = b"MDBSPILL";
65/// On-disk format version this build reads and writes.
66const FORMAT_VERSION: u16 = 1;
67/// Encryption flag stored in the header: frames are plaintext.
68const ENC_PLAINTEXT: u8 = 0;
69/// Encryption flag stored in the header: frames are AES-256-GCM sealed.
70const ENC_AES_GCM: u8 = 1;
71/// Header length: magic + version + enc flag + reserved.
72const HEADER_LEN: usize = 8 + 2 + 1 + 1;
73/// Frame head length: payload len + crc + kind + seq.
74const FRAME_HEAD_LEN: usize = 4 + 4 + 1 + 8;
75/// Kind byte of a data frame.
76const FRAME_DATA: u8 = 0;
77/// Kind byte of the sealing trailer frame.
78const FRAME_TRAILER: u8 = 1;
79/// Trailer payload length: frame count + byte count + SHA-256.
80const TRAILER_LEN: usize = 8 + 8 + 32;
81/// Maximum bytes of one frame's plaintext payload (bounds reader allocation
82/// on corrupt input, mirroring the PITR chunk cap).
83const MAX_FRAME_PAYLOAD: u64 = 64 * 1024 * 1024;
84
85const CRC32C: Crc<u32> = Crc::<u32>::new(&CRC_32_ISCSI);
86
87/// Errors of spill configuration, I/O, budget enforcement, and verification.
88#[derive(Debug, thiserror::Error)]
89pub enum SpillError {
90    /// Filesystem failure on the spill tree.
91    #[error("io error: {0}")]
92    Io(#[from] std::io::Error),
93    /// The [`SpillConfig`] failed validation.
94    #[error("invalid spill configuration: {0}")]
95    InvalidConfig(&'static str),
96    /// A frame append would exceed the per-query or the node-global spill
97    /// budget (S1E-004 "bounded").
98    #[error(
99        "spill budget exceeded for query {query_id}: requested {requested} bytes \
100         ({query_remaining} per-query, {global_remaining} global remaining)"
101    )]
102    BudgetExceeded {
103        /// The query that hit its bound.
104        query_id: QueryId,
105        /// Bytes the frame needed.
106        requested: u64,
107        /// Per-query budget remaining before the attempt.
108        query_remaining: u64,
109        /// Node-global budget remaining before the attempt.
110        global_remaining: u64,
111    },
112    /// One frame's plaintext payload exceeded [`MAX_FRAME_PAYLOAD`].
113    #[error("spill frame of {bytes} bytes exceeds the {limit}-byte frame limit")]
114    FrameTooLarge {
115        /// Attempted payload size.
116        bytes: u64,
117        /// Configured bound.
118        limit: u64,
119    },
120    /// A frame's stored CRC32C did not match its bytes.
121    #[error("checksum mismatch for {context}: expected {expected}, got {actual}")]
122    ChecksumMismatch {
123        /// Which frame/file failed verification.
124        context: String,
125        /// Stored checksum.
126        expected: u32,
127        /// Computed checksum.
128        actual: u32,
129    },
130    /// Structural corruption: bad magic, version, framing, sequence gap, or
131    /// a trailer that does not match the streamed frames. Always fail-closed.
132    #[error("corrupt spill file: {0}")]
133    Corrupt(String),
134    /// An encrypted spill file was opened by a manager without a meta DEK.
135    #[error("encrypted spill file requires the database encryption key")]
136    EncryptionRequired,
137    /// A meta DEK was configured but the `encryption` feature is disabled.
138    #[error("spill encryption requires the `encryption` feature")]
139    EncryptionDisabled,
140    /// Sealing a frame failed.
141    #[error("spill encryption error: {0}")]
142    Encryption(String),
143    /// Opening a sealed frame failed (wrong key or tampering).
144    #[error("spill decryption error: {0}")]
145    Decryption(String),
146}
147
148impl From<SpillError> for MongrelError {
149    fn from(error: SpillError) -> Self {
150        match error {
151            SpillError::Io(error) => MongrelError::Io(error),
152            SpillError::InvalidConfig(message) => MongrelError::InvalidArgument(message.into()),
153            SpillError::BudgetExceeded {
154                requested,
155                query_remaining,
156                global_remaining,
157                ..
158            } => MongrelError::ResourceLimitExceeded {
159                resource: "spill temporary disk",
160                requested: usize::try_from(requested).unwrap_or(usize::MAX),
161                limit: usize::try_from(
162                    requested.saturating_add(query_remaining.min(global_remaining)),
163                )
164                .unwrap_or(usize::MAX),
165            },
166            SpillError::FrameTooLarge { bytes, limit } => MongrelError::ResourceLimitExceeded {
167                resource: "spill frame",
168                requested: usize::try_from(bytes).unwrap_or(usize::MAX),
169                limit: usize::try_from(limit).unwrap_or(usize::MAX),
170            },
171            SpillError::ChecksumMismatch {
172                context,
173                expected,
174                actual,
175            } => MongrelError::ChecksumMismatch {
176                expected: u64::from(expected),
177                actual: u64::from(actual),
178                context,
179            },
180            SpillError::Corrupt(message) => {
181                MongrelError::Other(format!("corrupt spill file: {message}"))
182            }
183            SpillError::EncryptionRequired | SpillError::EncryptionDisabled => {
184                MongrelError::EncryptionDisabled
185            }
186            SpillError::Encryption(message) => MongrelError::Encryption(message),
187            SpillError::Decryption(message) => MongrelError::Decryption(message),
188        }
189    }
190}
191
192/// Node-global spill configuration (S1E-004). The per-query bound is supplied
193/// per session from [`ResourceGroup::temporary_disk_bytes`].
194#[derive(Debug, Clone, Copy, PartialEq, Eq)]
195pub struct SpillConfig {
196    /// Total bytes of live spill files across every query on this node.
197    pub global_bytes: u64,
198}
199
200impl SpillConfig {
201    /// A config with the given node-global cap.
202    pub fn new(global_bytes: u64) -> Self {
203        Self { global_bytes }
204    }
205
206    fn validate(&self) -> Result<(), SpillError> {
207        if self.global_bytes == 0 {
208            return Err(SpillError::InvalidConfig(
209                "global spill budget must be nonzero",
210            ));
211        }
212        Ok(())
213    }
214}
215
216/// A point-in-time snapshot of spill-manager state (telemetry and tests).
217#[derive(Debug, Clone, PartialEq, Eq)]
218pub struct SpillStats {
219    /// Cumulative stored bytes written (headers, frames, trailers).
220    pub bytes_written: u64,
221    /// Cumulative stored bytes read back.
222    pub bytes_read: u64,
223    /// Finished spill files currently live (held by a [`SpillHandle`]).
224    pub files_live: u64,
225    /// Live bytes currently charged against the global budget.
226    pub global_used: u64,
227    /// Configured node-global budget.
228    pub global_budget_bytes: u64,
229    /// Global budget remaining (`global_budget_bytes - global_used`).
230    pub budget_remaining: u64,
231}
232
233struct ManagerInner {
234    db_root: DurableRoot,
235    /// Lazily created pinned `temp/spill` root (`None` until the first spill
236    /// file is created; the startup sweep never creates it).
237    spill_root: parking_lot::Mutex<Option<DurableRoot>>,
238    config: SpillConfig,
239    meta_dek: Option<[u8; DEK_LEN]>,
240    global_used: AtomicU64,
241    bytes_written: AtomicU64,
242    bytes_read: AtomicU64,
243    files_live: AtomicU64,
244}
245
246impl ManagerInner {
247    /// The pinned `temp/spill` root, creating it (and `temp/`) on first use.
248    fn spill_root(&self) -> Result<DurableRoot, SpillError> {
249        let mut guard = self.spill_root.lock();
250        if let Some(root) = guard.as_ref() {
251            return Ok(root.try_clone()?);
252        }
253        let root = self.db_root.create_directory_all_pinned(SPILL_DIR_REL)?;
254        *guard = Some(root.try_clone()?);
255        Ok(root)
256    }
257
258    /// The pinned `temp/spill` root only if it already exists (never
259    /// creates). Used by cleanup paths, which must not resurrect the tree.
260    fn spill_root_if_created(&self) -> Option<DurableRoot> {
261        self.spill_root
262            .lock()
263            .as_ref()
264            .and_then(|root| root.try_clone().ok())
265    }
266
267    /// Removes every stale entry a prior process run left in `temp/spill`.
268    /// Spill files are process-local by construction (live files are held by
269    /// handles), so everything present at open is garbage.
270    fn sweep_stale(&self) -> Result<(), SpillError> {
271        match self.db_root.entry_exists(SPILL_DIR_REL) {
272            Ok(true) => {}
273            // No `temp/spill` (or no `temp` yet): nothing to sweep. The tree
274            // is created lazily on the first spill.
275            Ok(false) => return Ok(()),
276            Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(()),
277            Err(error) => return Err(SpillError::Io(error)),
278        }
279        let root = self.db_root.open_directory(SPILL_DIR_REL)?;
280        // `io_path` is the descriptor-pinned operational path of `root`, so
281        // the enumeration itself cannot wander off the pinned directory.
282        for entry in std::fs::read_dir(root.io_path()?)? {
283            let entry = entry?;
284            let name = entry.file_name();
285            if entry.file_type()?.is_dir() {
286                root.remove_directory_all(Path::new(&name))?;
287            } else {
288                root.remove_file(Path::new(&name))?;
289            }
290        }
291        Ok(())
292    }
293
294    fn header(&self) -> [u8; HEADER_LEN] {
295        let mut header = [0u8; HEADER_LEN];
296        header[..8].copy_from_slice(SPILL_MAGIC);
297        header[8..10].copy_from_slice(&FORMAT_VERSION.to_le_bytes());
298        header[10] = if self.meta_dek.is_some() {
299            ENC_AES_GCM
300        } else {
301            ENC_PLAINTEXT
302        };
303        header
304    }
305
306    /// Charges `bytes` against the per-query and global budgets with the
307    /// governor's add-then-validate-rollback protocol, so a granted set never
308    /// exceeds either bound.
309    fn try_charge(&self, session: &SessionInner, bytes: u64) -> Result<(), SpillError> {
310        let new_query = session.used.fetch_add(bytes, Ordering::Relaxed) + bytes;
311        let new_global = self.global_used.fetch_add(bytes, Ordering::Relaxed) + bytes;
312        if new_query <= session.cap && new_global <= self.config.global_bytes {
313            Ok(())
314        } else {
315            session.used.fetch_sub(bytes, Ordering::Relaxed);
316            self.global_used.fetch_sub(bytes, Ordering::Relaxed);
317            Err(SpillError::BudgetExceeded {
318                query_id: session.query_id,
319                requested: bytes,
320                query_remaining: session.cap.saturating_sub(new_query - bytes),
321                global_remaining: self.config.global_bytes.saturating_sub(new_global - bytes),
322            })
323        }
324    }
325
326    /// Releases a previous charge (exact inverse of [`try_charge`](Self::try_charge)).
327    fn release(&self, session: &SessionInner, bytes: u64) {
328        session.used.fetch_sub(bytes, Ordering::Relaxed);
329        self.global_used.fetch_sub(bytes, Ordering::Relaxed);
330    }
331}
332
333/// The node-level spill manager (S1E-004). Cheap to clone (one `Arc`);
334/// thread-safe; budget accounting is lock-free.
335pub struct SpillManager {
336    inner: Arc<ManagerInner>,
337}
338
339impl Clone for SpillManager {
340    fn clone(&self) -> Self {
341        Self {
342            inner: Arc::clone(&self.inner),
343        }
344    }
345}
346
347impl fmt::Debug for SpillManager {
348    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
349        f.debug_struct("SpillManager")
350            .field("global_budget_bytes", &self.inner.config.global_bytes)
351            .field(
352                "global_used",
353                &self.inner.global_used.load(Ordering::Relaxed),
354            )
355            .field("files_live", &self.inner.files_live.load(Ordering::Relaxed))
356            .field("encrypted", &self.inner.meta_dek.is_some())
357            .finish()
358    }
359}
360
361impl SpillManager {
362    /// Opens the spill manager on a database root, sweeping every stale entry
363    /// a prior process run left in `temp/spill` (S1E-004 startup cleanup).
364    /// The `temp/spill` tree itself is created lazily on the first spill.
365    ///
366    /// `meta_dek` is the database meta DEK
367    /// ([`crate::encryption::meta_dek_for`]): `Some` seals every frame with
368    /// AES-256-GCM, `None` stores frames plaintext. Without the `encryption`
369    /// feature a `Some` DEK is rejected (fail closed).
370    pub fn open(
371        db_root: &DurableRoot,
372        config: SpillConfig,
373        meta_dek: Option<[u8; DEK_LEN]>,
374    ) -> Result<Self, SpillError> {
375        config.validate()?;
376        #[cfg(not(feature = "encryption"))]
377        if meta_dek.is_some() {
378            return Err(SpillError::EncryptionDisabled);
379        }
380        let manager = Self {
381            inner: Arc::new(ManagerInner {
382                db_root: db_root.try_clone()?,
383                spill_root: parking_lot::Mutex::new(None),
384                config,
385                meta_dek,
386                global_used: AtomicU64::new(0),
387                bytes_written: AtomicU64::new(0),
388                bytes_read: AtomicU64::new(0),
389                files_live: AtomicU64::new(0),
390            }),
391        };
392        manager.inner.sweep_stale()?;
393        Ok(manager)
394    }
395
396    /// The manager's configuration.
397    pub fn config(&self) -> &SpillConfig {
398        &self.inner.config
399    }
400
401    /// Starts a spill session for one query with an explicit per-query cap.
402    /// One session per `query_id` at a time: sessions share the per-query
403    /// directory name, and chunk creation is exclusive.
404    pub fn begin_query(
405        &self,
406        query_id: QueryId,
407        per_query_bytes: u64,
408    ) -> Result<SpillSession, SpillError> {
409        Ok(SpillSession {
410            inner: Arc::new(SessionInner {
411                manager: self.clone(),
412                query_id,
413                dir_name: format!("q-{}", query_id.to_hex()),
414                cap: per_query_bytes,
415                used: AtomicU64::new(0),
416                next_chunk: AtomicU64::new(0),
417            }),
418        })
419    }
420
421    /// Starts a spill session for one query admitted into `group`, with the
422    /// per-query cap fed from [`ResourceGroup::temporary_disk_bytes`]
423    /// (S1E-002/S1E-004).
424    pub fn begin_query_in_group(
425        &self,
426        query_id: QueryId,
427        group: &ResourceGroup,
428    ) -> Result<SpillSession, SpillError> {
429        self.begin_query(query_id, group.temporary_disk_bytes)
430    }
431
432    /// A point-in-time snapshot of spill state.
433    pub fn stats(&self) -> SpillStats {
434        let global_used = self.inner.global_used.load(Ordering::Relaxed);
435        SpillStats {
436            bytes_written: self.inner.bytes_written.load(Ordering::Relaxed),
437            bytes_read: self.inner.bytes_read.load(Ordering::Relaxed),
438            files_live: self.inner.files_live.load(Ordering::Relaxed),
439            global_used,
440            global_budget_bytes: self.inner.config.global_bytes,
441            budget_remaining: self.inner.config.global_bytes.saturating_sub(global_used),
442        }
443    }
444}
445
446struct SessionInner {
447    manager: SpillManager,
448    query_id: QueryId,
449    /// Per-query directory name under `temp/spill` (`q-<hex>`).
450    dir_name: String,
451    cap: u64,
452    used: AtomicU64,
453    next_chunk: AtomicU64,
454}
455
456impl SessionInner {
457    /// The pinned per-query directory, creating it on first use.
458    fn query_dir(&self) -> Result<DurableRoot, SpillError> {
459        Ok(self
460            .manager
461            .inner
462            .spill_root()?
463            .create_directory_all_pinned(&self.dir_name)?)
464    }
465}
466
467impl Drop for SessionInner {
468    /// Removes the whole per-query directory (S1E-004 cleanup on
469    /// success/error/cancel). Never creates the tree to delete it.
470    fn drop(&mut self) {
471        if let Some(root) = self.manager.inner.spill_root_if_created() {
472            let _ = root.remove_directory_all(Path::new(&self.dir_name));
473        }
474    }
475}
476
477/// A query's spill session: namespaced, budgeted factory of spill files.
478/// Dropping the session removes the whole per-query directory.
479pub struct SpillSession {
480    inner: Arc<SessionInner>,
481}
482
483impl fmt::Debug for SpillSession {
484    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
485        f.debug_struct("SpillSession")
486            .field("query_id", &self.inner.query_id)
487            .field("cap", &self.inner.cap)
488            .field("used", &self.inner.used.load(Ordering::Relaxed))
489            .finish()
490    }
491}
492
493impl SpillSession {
494    /// The query this session spills for.
495    pub fn query_id(&self) -> QueryId {
496        self.inner.query_id
497    }
498
499    /// The per-query spill cap in bytes.
500    pub fn cap(&self) -> u64 {
501        self.inner.cap
502    }
503
504    /// Live bytes currently charged to this query.
505    pub fn used(&self) -> u64 {
506        self.inner.used.load(Ordering::Relaxed)
507    }
508
509    /// Per-query budget remaining.
510    pub fn budget_remaining(&self) -> u64 {
511        self.inner.cap.saturating_sub(self.used())
512    }
513
514    /// Creates the next spill file of this query and returns its writer.
515    /// Dropping the writer before [`SpillWriter::finish`] deletes the partial
516    /// file (the error/cancel path).
517    pub fn new_writer(&self) -> Result<SpillWriter, SpillError> {
518        let seq = self.inner.next_chunk.fetch_add(1, Ordering::Relaxed);
519        let name = format!("chunk-{seq:06}.spill");
520        let dir = self.inner.query_dir()?;
521        let manager = &self.inner.manager;
522        manager.inner.try_charge(&self.inner, HEADER_LEN as u64)?;
523        let result = (|| {
524            let mut file = dir.create_regular_new(&name)?;
525            file.write_all(&manager.inner.header())?;
526            io::Result::Ok(file)
527        })();
528        match result {
529            Ok(file) => {
530                manager
531                    .inner
532                    .bytes_written
533                    .fetch_add(HEADER_LEN as u64, Ordering::Relaxed);
534                Ok(SpillWriter {
535                    inner: Some(WriterInner {
536                        session: Arc::clone(&self.inner),
537                        dir,
538                        name,
539                        file,
540                        bytes_on_disk: HEADER_LEN as u64,
541                        data_frames: 0,
542                        data_bytes: 0,
543                        digest: Sha256::new(),
544                        next_seq: 0,
545                    }),
546                })
547            }
548            Err(error) => {
549                manager.inner.release(&self.inner, HEADER_LEN as u64);
550                Err(SpillError::Io(error))
551            }
552        }
553    }
554}
555
556struct WriterInner {
557    session: Arc<SessionInner>,
558    /// Pinned per-query directory (descriptor-relative delete on abort).
559    dir: DurableRoot,
560    name: String,
561    file: std::fs::File,
562    bytes_on_disk: u64,
563    data_frames: u64,
564    data_bytes: u64,
565    digest: Sha256,
566    next_seq: u64,
567}
568
569impl WriterInner {
570    /// Seals (when encrypted), frames, charges, and appends one frame.
571    fn write_frame(&mut self, kind: u8, payload: &[u8]) -> Result<(), SpillError> {
572        if payload.len() as u64 > MAX_FRAME_PAYLOAD {
573            return Err(SpillError::FrameTooLarge {
574                bytes: payload.len() as u64,
575                limit: MAX_FRAME_PAYLOAD,
576            });
577        }
578        let stored = seal_payload(self.session.manager.inner.meta_dek.as_ref(), payload)?;
579        let seq = self.next_seq;
580        let mut digest = CRC32C.digest();
581        digest.update(&[kind]);
582        digest.update(&seq.to_le_bytes());
583        digest.update(&stored);
584        let crc = digest.finalize();
585        let frame_bytes = FRAME_HEAD_LEN as u64 + stored.len() as u64;
586        let manager = &self.session.manager;
587        manager.inner.try_charge(&self.session, frame_bytes)?;
588        let result = (|| {
589            self.file.write_all(&(stored.len() as u32).to_le_bytes())?;
590            self.file.write_all(&crc.to_le_bytes())?;
591            self.file.write_all(&[kind])?;
592            self.file.write_all(&seq.to_le_bytes())?;
593            self.file.write_all(&stored)
594        })();
595        match result {
596            Ok(()) => {
597                self.next_seq += 1;
598                self.bytes_on_disk += frame_bytes;
599                manager
600                    .inner
601                    .bytes_written
602                    .fetch_add(frame_bytes, Ordering::Relaxed);
603                Ok(())
604            }
605            Err(error) => {
606                manager.inner.release(&self.session, frame_bytes);
607                Err(SpillError::Io(error))
608            }
609        }
610    }
611
612    /// Writes the sealing trailer (frame count, byte count, SHA-256 over all
613    /// plaintext data payloads), then fsyncs file and directory.
614    fn write_trailer_and_sync(&mut self) -> Result<(), SpillError> {
615        let hash: [u8; 32] = self.digest.clone().finalize().into();
616        let mut trailer = Vec::with_capacity(TRAILER_LEN);
617        trailer.extend_from_slice(&self.data_frames.to_le_bytes());
618        trailer.extend_from_slice(&self.data_bytes.to_le_bytes());
619        trailer.extend_from_slice(&hash);
620        self.write_frame(FRAME_TRAILER, &trailer)?;
621        self.file.sync_all()?;
622        self.dir.sync_entry_parent(Path::new(&self.name))?;
623        Ok(())
624    }
625}
626
627/// Deletes the partial file of an unfinished writer and releases its budget
628/// (the error/cancel path; S1E-004 "deleted on error/cancel").
629fn abort_writer(inner: WriterInner) {
630    let WriterInner {
631        session,
632        dir,
633        name,
634        file,
635        bytes_on_disk,
636        ..
637    } = inner;
638    drop(file);
639    let _ = dir.remove_file(Path::new(&name));
640    session.manager.inner.release(&session, bytes_on_disk);
641}
642
643/// Streaming writer of one spill file. Append data frames with
644/// [`append`](Self::append); seal and fsync with [`finish`](Self::finish).
645/// Dropping an unfinished writer deletes the partial file.
646pub struct SpillWriter {
647    inner: Option<WriterInner>,
648}
649
650impl fmt::Debug for SpillWriter {
651    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
652        f.debug_struct("SpillWriter")
653            .field("name", &self.inner.as_ref().map(|inner| &inner.name))
654            .field(
655                "bytes_on_disk",
656                &self.inner.as_ref().map(|inner| inner.bytes_on_disk),
657            )
658            .finish()
659    }
660}
661
662impl SpillWriter {
663    /// Appends one data frame. `payload` is opaque to the manager (an
664    /// operator's serialized run); it is checksummed and, when the database
665    /// is encrypted, sealed before it touches disk.
666    pub fn append(&mut self, payload: &[u8]) -> Result<(), SpillError> {
667        let inner = self
668            .inner
669            .as_mut()
670            .expect("spill writer is live until finish");
671        inner.write_frame(FRAME_DATA, payload)?;
672        inner.digest.update(payload);
673        inner.data_frames += 1;
674        inner.data_bytes += payload.len() as u64;
675        Ok(())
676    }
677
678    /// Stored bytes written so far (header and frames, post-seal).
679    pub fn bytes_on_disk(&self) -> u64 {
680        self.inner.as_ref().map_or(0, |inner| inner.bytes_on_disk)
681    }
682
683    /// Seals the file with its checksum trailer, fsyncs it, and returns the
684    /// RAII handle. On error the partial file is deleted, exactly as if the
685    /// writer had been dropped.
686    pub fn finish(mut self) -> Result<SpillHandle, SpillError> {
687        let mut inner = self
688            .inner
689            .take()
690            .expect("spill writer is live until finish");
691        if let Err(error) = inner.write_trailer_and_sync() {
692            abort_writer(inner);
693            return Err(error);
694        }
695        inner
696            .session
697            .manager
698            .inner
699            .files_live
700            .fetch_add(1, Ordering::Relaxed);
701        let WriterInner {
702            session,
703            dir,
704            name,
705            bytes_on_disk,
706            data_frames,
707            ..
708        } = inner;
709        Ok(SpillHandle {
710            inner: Some(HandleInner {
711                session,
712                dir,
713                name,
714                bytes_on_disk,
715                data_frames,
716            }),
717        })
718    }
719
720    /// Abandons the file: deletes the partial write and releases its budget.
721    /// Equivalent to dropping the writer, but reports I/O failures.
722    pub fn abort(mut self) -> Result<(), SpillError> {
723        let inner = self
724            .inner
725            .take()
726            .expect("spill writer is live until finish");
727        let WriterInner {
728            session,
729            dir,
730            name,
731            file,
732            bytes_on_disk,
733            ..
734        } = inner;
735        drop(file);
736        dir.remove_file(Path::new(&name))?;
737        session.manager.inner.release(&session, bytes_on_disk);
738        Ok(())
739    }
740}
741
742impl Drop for SpillWriter {
743    fn drop(&mut self) {
744        if let Some(inner) = self.inner.take() {
745            abort_writer(inner);
746        }
747    }
748}
749
750struct HandleInner {
751    session: Arc<SessionInner>,
752    dir: DurableRoot,
753    name: String,
754    bytes_on_disk: u64,
755    data_frames: u64,
756}
757
758/// Deletes the finished file and releases its budget and live-file count.
759fn delete_file(inner: HandleInner) -> Result<(), SpillError> {
760    inner.dir.remove_file(Path::new(&inner.name))?;
761    inner
762        .session
763        .manager
764        .inner
765        .release(&inner.session, inner.bytes_on_disk);
766    inner
767        .session
768        .manager
769        .inner
770        .files_live
771        .fetch_sub(1, Ordering::Relaxed);
772    Ok(())
773}
774
775/// RAII guard for one sealed spill file: dropping it deletes the file and
776/// releases its budget (S1E-004 "deleted on success"). Open a verify-on-read
777/// [`SpillReader`] with [`reader`](Self::reader).
778#[must_use = "a spill handle deletes its file on drop"]
779pub struct SpillHandle {
780    inner: Option<HandleInner>,
781}
782
783impl fmt::Debug for SpillHandle {
784    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
785        f.debug_struct("SpillHandle")
786            .field("name", &self.inner.as_ref().map(|inner| &inner.name))
787            .field(
788                "bytes_on_disk",
789                &self.inner.as_ref().map(|inner| inner.bytes_on_disk),
790            )
791            .finish()
792    }
793}
794
795impl SpillHandle {
796    /// The query whose session created this file.
797    pub fn query_id(&self) -> QueryId {
798        self.inner
799            .as_ref()
800            .expect("spill handle is live")
801            .session
802            .query_id
803    }
804
805    /// Stored file size in bytes (header, frames, trailer).
806    pub fn bytes_on_disk(&self) -> u64 {
807        self.inner.as_ref().map_or(0, |inner| inner.bytes_on_disk)
808    }
809
810    /// Number of data frames sealed into the file.
811    pub fn frames(&self) -> u64 {
812        self.inner.as_ref().map_or(0, |inner| inner.data_frames)
813    }
814
815    /// Opens a streaming, verify-on-read reader over the file.
816    pub fn reader(&self) -> Result<SpillReader, SpillError> {
817        let inner = self.inner.as_ref().expect("spill handle is live");
818        let file = inner.dir.open_regular(Path::new(&inner.name))?;
819        reader_from(file, &inner.session.manager)
820    }
821
822    /// Deletes the file now, releasing its budget. Equivalent to dropping the
823    /// handle, but reports I/O failures.
824    pub fn delete(mut self) -> Result<(), SpillError> {
825        let inner = self.inner.take().expect("spill handle is live");
826        delete_file(inner)
827    }
828}
829
830impl Drop for SpillHandle {
831    fn drop(&mut self) {
832        if let Some(inner) = self.inner.take() {
833            let _ = delete_file(inner);
834        }
835    }
836}
837
838/// Parses and validates a spill file header and builds the reader (shared by
839/// [`SpillHandle::reader`] and tests). Fails closed on bad magic, an
840/// unsupported version, an unknown encryption flag, or a DEK mismatch with
841/// the manager (mirroring the PITR chunk envelope rules).
842fn reader_from(mut file: std::fs::File, manager: &SpillManager) -> Result<SpillReader, SpillError> {
843    let mut header = [0u8; HEADER_LEN];
844    file.read_exact(&mut header).map_err(|error| {
845        if error.kind() == io::ErrorKind::UnexpectedEof {
846            SpillError::Corrupt("spill file is shorter than its header".into())
847        } else {
848            SpillError::Io(error)
849        }
850    })?;
851    manager
852        .inner
853        .bytes_read
854        .fetch_add(HEADER_LEN as u64, Ordering::Relaxed);
855    if header[..8] != SPILL_MAGIC[..] {
856        return Err(SpillError::Corrupt("bad spill magic".into()));
857    }
858    let version = u16::from_le_bytes(header[8..10].try_into().expect("slice length"));
859    if version != FORMAT_VERSION {
860        return Err(SpillError::Corrupt(format!(
861            "unsupported spill format version {version}"
862        )));
863    }
864    let dek = match (header[10], manager.inner.meta_dek.as_ref()) {
865        (ENC_PLAINTEXT, None) => None,
866        (ENC_AES_GCM, Some(dek)) => Some(*dek),
867        (ENC_AES_GCM, None) => return Err(SpillError::EncryptionRequired),
868        (ENC_PLAINTEXT, Some(_)) => {
869            return Err(SpillError::Corrupt(
870                "plaintext spill file opened with an encryption key".into(),
871            ))
872        }
873        (other, _) => {
874            return Err(SpillError::Corrupt(format!(
875                "unknown spill encryption flag {other}"
876            )))
877        }
878    };
879    Ok(SpillReader {
880        file,
881        manager: manager.clone(),
882        dek,
883        next_seq: 0,
884        data_frames: 0,
885        data_bytes: 0,
886        digest: Sha256::new(),
887        done: false,
888    })
889}
890
891/// Streaming, verify-on-read reader of one sealed spill file. Every frame's
892/// CRC32C is verified before its payload is returned (and decrypted, when
893/// the file is sealed); the closing trailer re-checks the frame count, byte
894/// count, and the SHA-256 over all data payloads. Any failure is terminal.
895pub struct SpillReader {
896    file: std::fs::File,
897    manager: SpillManager,
898    dek: Option<[u8; DEK_LEN]>,
899    next_seq: u64,
900    data_frames: u64,
901    data_bytes: u64,
902    digest: Sha256,
903    done: bool,
904}
905
906impl fmt::Debug for SpillReader {
907    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
908        f.debug_struct("SpillReader")
909            .field("next_seq", &self.next_seq)
910            .field("data_frames", &self.data_frames)
911            .field("done", &self.done)
912            .finish()
913    }
914}
915
916impl SpillReader {
917    /// Reads the next data frame, or `None` once the sealing trailer has been
918    /// verified. Errors are terminal: after one, the reader yields no more
919    /// frames.
920    pub fn next_frame(&mut self) -> Result<Option<Vec<u8>>, SpillError> {
921        if self.done {
922            return Ok(None);
923        }
924        match self.next_frame_inner() {
925            Ok(frame) => Ok(frame),
926            Err(error) => {
927                self.done = true;
928                Err(error)
929            }
930        }
931    }
932
933    fn next_frame_inner(&mut self) -> Result<Option<Vec<u8>>, SpillError> {
934        let mut head = [0u8; FRAME_HEAD_LEN];
935        self.file.read_exact(&mut head).map_err(|error| {
936            if error.kind() == io::ErrorKind::UnexpectedEof {
937                SpillError::Corrupt("truncated spill file: missing trailer".into())
938            } else {
939                SpillError::Io(error)
940            }
941        })?;
942        let len = u64::from(u32::from_le_bytes(
943            head[0..4].try_into().expect("slice length"),
944        ));
945        let expected_crc = u32::from_le_bytes(head[4..8].try_into().expect("slice length"));
946        let kind = head[8];
947        let seq = u64::from_le_bytes(head[9..17].try_into().expect("slice length"));
948        if len > MAX_FRAME_PAYLOAD {
949            return Err(SpillError::Corrupt(format!(
950                "spill frame of {len} bytes exceeds the {MAX_FRAME_PAYLOAD}-byte limit"
951            )));
952        }
953        if seq != self.next_seq {
954            return Err(SpillError::Corrupt(format!(
955                "spill frame sequence gap: expected {}, found {seq}",
956                self.next_seq
957            )));
958        }
959        let mut stored = vec![0u8; len as usize];
960        self.file.read_exact(&mut stored).map_err(|error| {
961            if error.kind() == io::ErrorKind::UnexpectedEof {
962                SpillError::Corrupt("truncated spill frame payload".into())
963            } else {
964                SpillError::Io(error)
965            }
966        })?;
967        let mut digest = CRC32C.digest();
968        digest.update(&[kind]);
969        digest.update(&seq.to_le_bytes());
970        digest.update(&stored);
971        let actual_crc = digest.finalize();
972        if actual_crc != expected_crc {
973            return Err(SpillError::ChecksumMismatch {
974                context: format!("spill frame {seq}"),
975                expected: expected_crc,
976                actual: actual_crc,
977            });
978        }
979        self.manager
980            .inner
981            .bytes_read
982            .fetch_add(FRAME_HEAD_LEN as u64 + len, Ordering::Relaxed);
983        let plaintext = open_payload(self.dek.as_ref(), &stored)?;
984        self.next_seq += 1;
985        match kind {
986            FRAME_DATA => {
987                self.digest.update(&plaintext);
988                self.data_frames += 1;
989                self.data_bytes += plaintext.len() as u64;
990                Ok(Some(plaintext))
991            }
992            FRAME_TRAILER => {
993                if plaintext.len() != TRAILER_LEN {
994                    return Err(SpillError::Corrupt(
995                        "spill trailer has the wrong length".into(),
996                    ));
997                }
998                let frames = u64::from_le_bytes(plaintext[0..8].try_into().expect("slice length"));
999                let bytes = u64::from_le_bytes(plaintext[8..16].try_into().expect("slice length"));
1000                let hash: [u8; 32] = plaintext[16..48].try_into().expect("slice length");
1001                let actual_hash: [u8; 32] = self.digest.clone().finalize().into();
1002                if frames != self.data_frames || bytes != self.data_bytes || hash != actual_hash {
1003                    return Err(SpillError::Corrupt(
1004                        "spill trailer does not match the streamed frames".into(),
1005                    ));
1006                }
1007                self.done = true;
1008                Ok(None)
1009            }
1010            other => Err(SpillError::Corrupt(format!(
1011                "unknown spill frame kind {other}"
1012            ))),
1013        }
1014    }
1015}
1016
1017impl Iterator for SpillReader {
1018    type Item = Result<Vec<u8>, SpillError>;
1019
1020    fn next(&mut self) -> Option<Self::Item> {
1021        match self.next_frame() {
1022            Ok(Some(frame)) => Some(Ok(frame)),
1023            Ok(None) => None,
1024            Err(error) => Some(Err(error)),
1025        }
1026    }
1027}
1028
1029/// Maps the crypto stack's `MongrelError` onto the typed spill errors.
1030#[cfg(feature = "encryption")]
1031fn map_crypto(error: MongrelError) -> SpillError {
1032    match error {
1033        MongrelError::Encryption(message) => SpillError::Encryption(message),
1034        MongrelError::Decryption(message) => SpillError::Decryption(message),
1035        other => SpillError::Corrupt(other.to_string()),
1036    }
1037}
1038
1039/// Seals one frame's plaintext payload with the page-cipher stack when a meta
1040/// DEK is present (fresh random nonce per frame, the `encrypt_blob` idiom);
1041/// passes plaintext through otherwise.
1042#[cfg(feature = "encryption")]
1043fn seal_payload(dek: Option<&[u8; DEK_LEN]>, plaintext: &[u8]) -> Result<Vec<u8>, SpillError> {
1044    match dek {
1045        Some(dek) => crate::encryption::encrypt_blob(dek, plaintext).map_err(map_crypto),
1046        None => Ok(plaintext.to_vec()),
1047    }
1048}
1049
1050/// Fail-closed stub: [`SpillManager::open`] rejects a DEK without the
1051/// `encryption` feature, so this is unreachable in practice.
1052#[cfg(not(feature = "encryption"))]
1053fn seal_payload(dek: Option<&[u8; DEK_LEN]>, plaintext: &[u8]) -> Result<Vec<u8>, SpillError> {
1054    match dek {
1055        Some(_) => Err(SpillError::EncryptionDisabled),
1056        None => Ok(plaintext.to_vec()),
1057    }
1058}
1059
1060/// Inverse of [`seal_payload`]: authenticates and opens a sealed frame.
1061#[cfg(feature = "encryption")]
1062fn open_payload(dek: Option<&[u8; DEK_LEN]>, stored: &[u8]) -> Result<Vec<u8>, SpillError> {
1063    match dek {
1064        Some(dek) => crate::encryption::decrypt_blob(dek, stored).map_err(map_crypto),
1065        None => Ok(stored.to_vec()),
1066    }
1067}
1068
1069/// Fail-closed stub — see [`seal_payload`].
1070#[cfg(not(feature = "encryption"))]
1071fn open_payload(dek: Option<&[u8; DEK_LEN]>, stored: &[u8]) -> Result<Vec<u8>, SpillError> {
1072    match dek {
1073        Some(_) => Err(SpillError::EncryptionDisabled),
1074        None => Ok(stored.to_vec()),
1075    }
1076}
1077
1078#[cfg(test)]
1079mod tests {
1080    use super::*;
1081    use std::io::{Seek, SeekFrom};
1082    use std::path::PathBuf;
1083
1084    fn manager(dir: &tempfile::TempDir, global_bytes: u64) -> SpillManager {
1085        let root = DurableRoot::open(dir.path()).unwrap();
1086        SpillManager::open(&root, SpillConfig::new(global_bytes), None).unwrap()
1087    }
1088
1089    fn query_dir(dir: &tempfile::TempDir, query_id: QueryId) -> PathBuf {
1090        dir.path()
1091            .join("temp")
1092            .join("spill")
1093            .join(format!("q-{}", query_id.to_hex()))
1094    }
1095
1096    fn only_file(dir: &tempfile::TempDir, query_id: QueryId) -> PathBuf {
1097        let entries: Vec<_> = std::fs::read_dir(query_dir(dir, query_id))
1098            .unwrap()
1099            .map(|entry| entry.unwrap().path())
1100            .collect();
1101        assert_eq!(entries.len(), 1, "expected exactly one spill file");
1102        entries[0].clone()
1103    }
1104
1105    /// Serializes one frame exactly as `WriterInner::write_frame` does
1106    /// (plaintext mode), for crafting corrupt files.
1107    fn crafted_frame(kind: u8, seq: u64, payload: &[u8]) -> Vec<u8> {
1108        let mut digest = CRC32C.digest();
1109        digest.update(&[kind]);
1110        digest.update(&seq.to_le_bytes());
1111        digest.update(payload);
1112        let crc = digest.finalize();
1113        let mut out = Vec::with_capacity(FRAME_HEAD_LEN + payload.len());
1114        out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1115        out.extend_from_slice(&crc.to_le_bytes());
1116        out.push(kind);
1117        out.extend_from_slice(&seq.to_le_bytes());
1118        out.extend_from_slice(payload);
1119        out
1120    }
1121
1122    fn crafted_header(enc: u8) -> Vec<u8> {
1123        let mut header = vec![0u8; HEADER_LEN];
1124        header[..8].copy_from_slice(SPILL_MAGIC);
1125        header[8..10].copy_from_slice(&FORMAT_VERSION.to_le_bytes());
1126        header[10] = enc;
1127        header
1128    }
1129
1130    #[test]
1131    fn frame_round_trip_streams_in_order() {
1132        let dir = tempfile::tempdir().unwrap();
1133        let manager = manager(&dir, 1 << 20);
1134        let session = manager.begin_query(QueryId::new_random(), 1 << 20).unwrap();
1135        let payloads: Vec<Vec<u8>> = vec![
1136            b"first".to_vec(),
1137            Vec::new(), // empty frames are legal
1138            vec![0xAB; 1000],
1139            b"last".to_vec(),
1140        ];
1141        let mut writer = session.new_writer().unwrap();
1142        for payload in &payloads {
1143            writer.append(payload).unwrap();
1144        }
1145        let handle = writer.finish().unwrap();
1146        assert_eq!(handle.frames(), 4);
1147        assert_eq!(handle.query_id(), session.query_id());
1148        assert_eq!(handle.bytes_on_disk(), session.used());
1149
1150        let frames: Vec<Vec<u8>> = handle.reader().unwrap().collect::<Result<_, _>>().unwrap();
1151        assert_eq!(frames, payloads);
1152
1153        let stats = manager.stats();
1154        assert_eq!(stats.files_live, 1);
1155        assert_eq!(stats.bytes_written, handle.bytes_on_disk());
1156        assert_eq!(stats.bytes_read, handle.bytes_on_disk());
1157        assert_eq!(stats.global_used, handle.bytes_on_disk());
1158        assert_eq!(stats.budget_remaining, (1 << 20) - handle.bytes_on_disk());
1159    }
1160
1161    #[test]
1162    fn large_multi_chunk_round_trip() {
1163        let dir = tempfile::tempdir().unwrap();
1164        let manager = manager(&dir, 1 << 24);
1165        let session = manager.begin_query(QueryId::new_random(), 1 << 24).unwrap();
1166        // Two chunk files, three frames of mixed large sizes each.
1167        let mut handles = Vec::new();
1168        let mut expected = Vec::new();
1169        for chunk in 0..2u8 {
1170            let mut writer = session.new_writer().unwrap();
1171            let mut payloads = Vec::new();
1172            for (index, len) in [(1usize << 20), (1 << 20) + 7, 333_333].iter().enumerate() {
1173                let payload = vec![chunk * 16 + index as u8; *len];
1174                writer.append(&payload).unwrap();
1175                payloads.push(payload);
1176            }
1177            expected.push(payloads);
1178            handles.push(writer.finish().unwrap());
1179        }
1180        assert_eq!(manager.stats().files_live, 2);
1181        for (handle, payloads) in handles.iter().zip(expected.iter()) {
1182            let frames: Vec<Vec<u8>> = handle.reader().unwrap().collect::<Result<_, _>>().unwrap();
1183            assert_eq!(&frames, payloads);
1184        }
1185    }
1186
1187    #[test]
1188    fn checksum_mismatch_is_detected_on_read() {
1189        let dir = tempfile::tempdir().unwrap();
1190        let manager = manager(&dir, 1 << 20);
1191        let query_id = QueryId::new_random();
1192        let session = manager.begin_query(query_id, 1 << 20).unwrap();
1193        let mut writer = session.new_writer().unwrap();
1194        writer.append(&vec![0x11; 256]).unwrap();
1195        writer.append(b"second").unwrap();
1196        let handle = writer.finish().unwrap();
1197
1198        // Flip one byte in the first frame's payload (header 12 + head 17).
1199        let path = only_file(&dir, query_id);
1200        let mut file = std::fs::OpenOptions::new()
1201            .read(true)
1202            .write(true)
1203            .open(&path)
1204            .unwrap();
1205        file.seek(SeekFrom::Start((HEADER_LEN + FRAME_HEAD_LEN) as u64))
1206            .unwrap();
1207        file.write_all(&[0x99]).unwrap();
1208        drop(file);
1209
1210        let mut reader = handle.reader().unwrap();
1211        let error = reader.next_frame().unwrap_err();
1212        assert!(
1213            matches!(error, SpillError::ChecksumMismatch { .. }),
1214            "expected ChecksumMismatch, got {error:?}"
1215        );
1216        // The failure is terminal: no further frames are yielded.
1217        assert!(reader.next_frame().unwrap().is_none());
1218    }
1219
1220    #[test]
1221    fn trailer_mismatch_is_detected() {
1222        let dir = tempfile::tempdir().unwrap();
1223        let manager = manager(&dir, 1 << 20);
1224        // A crafted file: one valid data frame, then a trailer claiming two
1225        // frames — all CRCs valid, only the trailer semantics are wrong.
1226        let mut bytes = crafted_header(ENC_PLAINTEXT);
1227        bytes.extend_from_slice(&crafted_frame(FRAME_DATA, 0, b"abc"));
1228        let mut trailer = Vec::with_capacity(TRAILER_LEN);
1229        trailer.extend_from_slice(&2u64.to_le_bytes()); // wrong frame count
1230        trailer.extend_from_slice(&3u64.to_le_bytes());
1231        trailer.extend_from_slice(&[0u8; 32]);
1232        bytes.extend_from_slice(&crafted_frame(FRAME_TRAILER, 1, &trailer));
1233        let path = dir.path().join("crafted.spill");
1234        std::fs::write(&path, &bytes).unwrap();
1235
1236        let file = std::fs::File::open(&path).unwrap();
1237        let mut reader = reader_from(file, &manager).unwrap();
1238        assert_eq!(reader.next_frame().unwrap(), Some(b"abc".to_vec()));
1239        let error = reader.next_frame().unwrap_err();
1240        assert!(
1241            matches!(error, SpillError::Corrupt(_)),
1242            "expected Corrupt, got {error:?}"
1243        );
1244    }
1245
1246    #[test]
1247    fn bad_magic_and_version_are_rejected() {
1248        let dir = tempfile::tempdir().unwrap();
1249        let manager = manager(&dir, 1 << 20);
1250        for (label, mut header) in [
1251            ("magic", crafted_header(ENC_PLAINTEXT)),
1252            ("version", crafted_header(ENC_PLAINTEXT)),
1253            ("enc flag", crafted_header(ENC_PLAINTEXT)),
1254        ] {
1255            match label {
1256                "magic" => header[0] ^= 0xFF,
1257                "version" => header[8..10].copy_from_slice(&99u16.to_le_bytes()),
1258                _ => header[10] = 77,
1259            }
1260            let path = dir.path().join(format!("{label}.spill"));
1261            std::fs::write(&path, &header).unwrap();
1262            let file = std::fs::File::open(&path).unwrap();
1263            let error = reader_from(file, &manager).unwrap_err();
1264            assert!(
1265                matches!(error, SpillError::Corrupt(_)),
1266                "{label}: expected Corrupt, got {error:?}"
1267            );
1268        }
1269    }
1270
1271    #[test]
1272    fn per_query_budget_is_enforced() {
1273        let dir = tempfile::tempdir().unwrap();
1274        let manager = manager(&dir, 1 << 20);
1275        let query_id = QueryId::new_random();
1276        let session = manager.begin_query(query_id, 100).unwrap();
1277        let mut writer = session.new_writer().unwrap(); // header: 12 bytes
1278        writer.append(&[0u8; 50]).unwrap(); // frame: 17 + 50 = 67 → 79 used
1279        assert_eq!(session.used(), 79);
1280        assert_eq!(session.budget_remaining(), 21);
1281        let error = writer.append(&[0u8; 50]).unwrap_err();
1282        assert!(
1283            matches!(
1284                error,
1285                SpillError::BudgetExceeded {
1286                    query_id: id,
1287                    requested,
1288                    query_remaining: 21,
1289                    ..
1290                } if id == query_id && requested == 67
1291            ),
1292            "expected BudgetExceeded, got {error:?}"
1293        );
1294        // The failed frame charged nothing; a fitting frame still lands.
1295        assert_eq!(session.used(), 79);
1296        writer.append(&[0u8; 4]).unwrap(); // 17 + 4 = 21 → exactly the cap
1297        assert_eq!(session.used(), 100);
1298        assert_eq!(session.budget_remaining(), 0);
1299        let handle = writer.finish().unwrap_err();
1300        // Even the trailer no longer fits.
1301        assert!(matches!(handle, SpillError::BudgetExceeded { .. }));
1302    }
1303
1304    #[test]
1305    fn global_budget_is_enforced_across_queries() {
1306        let dir = tempfile::tempdir().unwrap();
1307        let manager = manager(&dir, 200);
1308        let a = manager.begin_query(QueryId::new_random(), 1 << 20).unwrap();
1309        let b = manager.begin_query(QueryId::new_random(), 1 << 20).unwrap();
1310        let mut writer_a = a.new_writer().unwrap();
1311        let mut writer_b = b.new_writer().unwrap();
1312        // Headers: 2 × 12 = 24 bytes of the global budget.
1313        writer_a.append(&[0u8; 100]).unwrap(); // +117 → 141 global
1314        let error = writer_b.append(&[0u8; 100]).unwrap_err(); // +117 > 200
1315        assert!(
1316            matches!(
1317                error,
1318                SpillError::BudgetExceeded {
1319                    global_remaining: 59,
1320                    ..
1321                }
1322            ),
1323            "expected global BudgetExceeded, got {error:?}"
1324        );
1325        assert_eq!(manager.stats().global_used, 141);
1326        // Releasing the first query's file re-opens the global budget.
1327        drop(writer_a);
1328        assert_eq!(manager.stats().global_used, 12);
1329        writer_b.append(&[0u8; 100]).unwrap();
1330    }
1331
1332    #[test]
1333    fn unfinished_writer_drop_deletes_file_and_releases_budget() {
1334        let dir = tempfile::tempdir().unwrap();
1335        let manager = manager(&dir, 1 << 20);
1336        let query_id = QueryId::new_random();
1337        let session = manager.begin_query(query_id, 1 << 20).unwrap();
1338        let mut writer = session.new_writer().unwrap();
1339        writer.append(&[0u8; 100]).unwrap();
1340        let path = only_file(&dir, query_id);
1341        assert!(path.exists());
1342        let used = session.used();
1343        assert!(used > 0);
1344        drop(writer); // the cancel path
1345        assert!(!path.exists(), "partial spill file must be deleted");
1346        assert_eq!(session.used(), 0);
1347        assert_eq!(manager.stats().global_used, 0);
1348        assert_eq!(manager.stats().files_live, 0);
1349    }
1350
1351    #[test]
1352    fn explicit_abort_deletes_and_reports() {
1353        let dir = tempfile::tempdir().unwrap();
1354        let manager = manager(&dir, 1 << 20);
1355        let query_id = QueryId::new_random();
1356        let session = manager.begin_query(query_id, 1 << 20).unwrap();
1357        let mut writer = session.new_writer().unwrap();
1358        writer.append(&[0u8; 64]).unwrap();
1359        let path = only_file(&dir, query_id);
1360        writer.abort().unwrap();
1361        assert!(!path.exists());
1362        assert_eq!(session.used(), 0);
1363        assert_eq!(manager.stats().global_used, 0);
1364    }
1365
1366    #[test]
1367    fn finished_handle_drop_and_delete_release_everything() {
1368        let dir = tempfile::tempdir().unwrap();
1369        let manager = manager(&dir, 1 << 20);
1370        let query_id = QueryId::new_random();
1371        let session = manager.begin_query(query_id, 1 << 20).unwrap();
1372
1373        let mut writer = session.new_writer().unwrap();
1374        writer.append(b"payload").unwrap();
1375        let handle = writer.finish().unwrap();
1376        let path = only_file(&dir, query_id);
1377        assert_eq!(manager.stats().files_live, 1);
1378        let used = session.used();
1379        drop(handle);
1380        assert!(!path.exists(), "sealed spill file must be deleted on drop");
1381        assert_eq!(session.used(), 0);
1382        assert_eq!(manager.stats().files_live, 0);
1383        assert_eq!(manager.stats().global_used, 0);
1384
1385        // Explicit delete reports and releases the same way.
1386        let mut writer = session.new_writer().unwrap();
1387        writer.append(b"payload").unwrap();
1388        let handle = writer.finish().unwrap();
1389        assert_eq!(session.used(), used);
1390        handle.delete().unwrap();
1391        assert_eq!(session.used(), 0);
1392        assert_eq!(manager.stats().files_live, 0);
1393    }
1394
1395    #[test]
1396    fn session_drop_removes_the_query_directory() {
1397        let dir = tempfile::tempdir().unwrap();
1398        let manager = manager(&dir, 1 << 20);
1399        let query_id = QueryId::new_random();
1400        let session = manager.begin_query(query_id, 1 << 20).unwrap();
1401        let mut writer = session.new_writer().unwrap();
1402        writer.append(b"x").unwrap();
1403        drop(writer);
1404        assert!(query_dir(&dir, query_id).exists());
1405        drop(session);
1406        assert!(
1407            !query_dir(&dir, query_id).exists(),
1408            "session drop must remove the per-query directory"
1409        );
1410        // A session that never spilled leaves nothing behind either.
1411        let quiet = manager.begin_query(QueryId::new_random(), 1 << 20).unwrap();
1412        drop(quiet);
1413        assert!(
1414            !dir.path().join("temp").join("spill").exists() || {
1415                std::fs::read_dir(dir.path().join("temp").join("spill"))
1416                    .unwrap()
1417                    .next()
1418                    .is_none()
1419            }
1420        );
1421    }
1422
1423    #[test]
1424    fn open_sweeps_stale_entries_from_prior_runs() {
1425        let dir = tempfile::tempdir().unwrap();
1426        let query_id = QueryId::new_random();
1427        let stale_path;
1428        let first = manager(&dir, 1 << 20);
1429        let handle;
1430        {
1431            let session = first.begin_query(query_id, 1 << 20).unwrap();
1432            let mut writer = session.new_writer().unwrap();
1433            writer.append(b"from a previous process").unwrap();
1434            handle = writer.finish().unwrap();
1435            stale_path = only_file(&dir, query_id);
1436            // A stray file directly in the spill root must be swept too.
1437            std::fs::write(dir.path().join("temp").join("spill").join("stray"), b"x").unwrap();
1438            // Simulate a crash: the session and handle leak away.
1439            std::mem::forget(session);
1440        }
1441        assert!(stale_path.exists());
1442        assert_eq!(first.stats().global_used, handle.bytes_on_disk());
1443
1444        // The next process opens its manager: everything stale is swept.
1445        let second = manager(&dir, 1 << 20);
1446        assert!(
1447            !stale_path.exists(),
1448            "startup sweep must remove stale files"
1449        );
1450        assert!(!dir.path().join("temp").join("spill").join("stray").exists());
1451        assert!(!query_dir(&dir, query_id).exists());
1452        assert_eq!(second.stats().global_used, 0);
1453        assert_eq!(second.stats().files_live, 0);
1454
1455        // The leaked handle's drop is still safe (no double delete) and its
1456        // accounting unwinds against the first manager.
1457        drop(handle);
1458        assert_eq!(first.stats().global_used, 0);
1459    }
1460
1461    #[test]
1462    fn frame_larger_than_the_limit_is_rejected() {
1463        let dir = tempfile::tempdir().unwrap();
1464        let manager = manager(&dir, u64::MAX);
1465        let session = manager
1466            .begin_query(QueryId::new_random(), u64::MAX)
1467            .unwrap();
1468        let mut writer = session.new_writer().unwrap();
1469        let huge = vec![0u8; MAX_FRAME_PAYLOAD as usize + 1];
1470        let error = writer.append(&huge).unwrap_err();
1471        assert!(
1472            matches!(error, SpillError::FrameTooLarge { .. }),
1473            "expected FrameTooLarge, got {error:?}"
1474        );
1475    }
1476
1477    #[test]
1478    fn invalid_config_is_rejected() {
1479        let dir = tempfile::tempdir().unwrap();
1480        let root = DurableRoot::open(dir.path()).unwrap();
1481        let error = SpillManager::open(&root, SpillConfig::new(0), None).unwrap_err();
1482        assert!(matches!(error, SpillError::InvalidConfig(_)));
1483    }
1484
1485    #[test]
1486    fn spill_error_maps_to_mongrel_error() {
1487        let io: MongrelError = SpillError::Io(io::Error::other("x")).into();
1488        assert!(matches!(io, MongrelError::Io(_)));
1489        let budget: MongrelError = SpillError::BudgetExceeded {
1490            query_id: QueryId::new_random(),
1491            requested: 10,
1492            query_remaining: 5,
1493            global_remaining: 90,
1494        }
1495        .into();
1496        assert!(matches!(budget, MongrelError::ResourceLimitExceeded { .. }));
1497        let checksum: MongrelError = SpillError::ChecksumMismatch {
1498            context: "frame".into(),
1499            expected: 1,
1500            actual: 2,
1501        }
1502        .into();
1503        assert!(matches!(checksum, MongrelError::ChecksumMismatch { .. }));
1504        let corrupt: MongrelError = SpillError::Corrupt("bad".into()).into();
1505        assert!(matches!(corrupt, MongrelError::Other(_)));
1506    }
1507
1508    #[cfg(feature = "encryption")]
1509    mod encrypted {
1510        use super::*;
1511        use crate::encryption::{meta_dek_for, Kek, SALT_LEN};
1512
1513        fn encrypted_manager(dir: &tempfile::TempDir, dek: [u8; DEK_LEN]) -> SpillManager {
1514            let root = DurableRoot::open(dir.path()).unwrap();
1515            SpillManager::open(&root, SpillConfig::new(1 << 20), Some(dek)).unwrap()
1516        }
1517
1518        fn test_dek(passphrase: &str) -> [u8; DEK_LEN] {
1519            let salt = [7u8; SALT_LEN];
1520            let kek = Kek::derive(passphrase, &salt).unwrap();
1521            meta_dek_for(Some(&kek)).unwrap()
1522        }
1523
1524        #[test]
1525        fn encrypted_round_trip_seals_every_frame_on_disk() {
1526            let dir = tempfile::tempdir().unwrap();
1527            let manager = encrypted_manager(&dir, test_dek("pw"));
1528            let query_id = QueryId::new_random();
1529            let session = manager.begin_query(query_id, 1 << 20).unwrap();
1530            let marker = b"highly-recognizable-plaintext-marker";
1531            let mut writer = session.new_writer().unwrap();
1532            writer.append(marker).unwrap();
1533            writer.append(&vec![0x5A; 4096]).unwrap();
1534            let handle = writer.finish().unwrap();
1535
1536            // Nothing on disk is the plaintext: marker and payload absent.
1537            let raw = std::fs::read(only_file(&dir, query_id)).unwrap();
1538            assert_eq!(raw[10], ENC_AES_GCM);
1539            assert!(!raw
1540                .windows(marker.len())
1541                .any(|window| window == marker.as_slice()));
1542            assert!(!raw.windows(64).any(|window| window == [0x5A; 64]));
1543
1544            let frames: Vec<Vec<u8>> = handle.reader().unwrap().collect::<Result<_, _>>().unwrap();
1545            assert_eq!(frames, vec![marker.to_vec(), vec![0x5A; 4096]]);
1546        }
1547
1548        #[test]
1549        fn encrypted_file_requires_the_key_and_detects_tampering() {
1550            let dir = tempfile::tempdir().unwrap();
1551            // Every manager is opened before any spill file exists: opening a
1552            // manager sweeps stale entries, so opening one after the file
1553            // lands would delete it (that behavior has its own test).
1554            let plaintext_manager = manager(&dir, 1 << 20);
1555            let wrong = encrypted_manager(&dir, test_dek("wrong"));
1556            let manager = encrypted_manager(&dir, test_dek("pw"));
1557            let query_id = QueryId::new_random();
1558            let session = manager.begin_query(query_id, 1 << 20).unwrap();
1559            let mut writer = session.new_writer().unwrap();
1560            writer.append(b"secret").unwrap();
1561            let handle = writer.finish().unwrap();
1562            let path = only_file(&dir, query_id);
1563            // A manager without the DEK refuses the encrypted file.
1564            let error =
1565                reader_from(std::fs::File::open(&path).unwrap(), &plaintext_manager).unwrap_err();
1566            assert!(
1567                matches!(error, SpillError::EncryptionRequired),
1568                "expected EncryptionRequired, got {error:?}"
1569            );
1570
1571            // A wrong DEK fails at the GCM tag (decryption), not at parse.
1572            let mut reader = reader_from(std::fs::File::open(&path).unwrap(), &wrong).unwrap();
1573            let error = reader.next_frame().unwrap_err();
1574            assert!(
1575                matches!(error, SpillError::Decryption(_)),
1576                "expected Decryption, got {error:?}"
1577            );
1578
1579            // The right key still reads.
1580            let frames: Vec<Vec<u8>> = handle.reader().unwrap().collect::<Result<_, _>>().unwrap();
1581            assert_eq!(frames, vec![b"secret".to_vec()]);
1582        }
1583    }
1584
1585    #[cfg(not(feature = "encryption"))]
1586    #[test]
1587    fn dek_without_the_encryption_feature_fails_closed() {
1588        let dir = tempfile::tempdir().unwrap();
1589        let root = DurableRoot::open(dir.path()).unwrap();
1590        let error =
1591            SpillManager::open(&root, SpillConfig::new(1 << 20), Some([1u8; DEK_LEN])).unwrap_err();
1592        assert!(matches!(error, SpillError::EncryptionDisabled));
1593    }
1594}