polyc-query 2026.9.6

//! Resource ceilings applied to every query the projected read path runs.
//!
//! [`QueryLimits`] is the one resource-ceilings struct this crate exports:
//! `polyc-query-service` builds it from deployment configuration and passes
//! it through [`crate::core_service::ProjectedCoreComposition`]. Every
//! field is one the projected path reads:
//!
//! - `row_cap` and `timeout` bound every admitted statement when
//!   `crate::core_resolution` mints the effective per-request bounds;
//!   both are build constants today — the deployment sets neither from
//!   configuration, and no environment key feeds either (POLY-374).
//! - `spill_dir` and `spill_quota_bytes` configure the execution
//!   `RuntimeEnv` `crate::core_service` builds at composition.
//!
//! The execution memory pool is NOT one of these fields: it is sized from
//! `polyc-query-service`'s own policy (`POLYCHROME_QUERY_MEMORY_BYTES`),
//! which is why this struct carries no `memory_bytes`.

use std::path::PathBuf;
use std::time::Duration;

/// Resource ceilings applied to every query the projected read path runs.
///
/// Every query enforces a hard wall-clock timeout and a row cap whose
/// truncation is flagged in the response — [`QueryLimits::timeout`] and
/// [`QueryLimits::row_cap`], both consulted when `crate::core_resolution`
/// mints the effective bounds for one request. [`QueryLimits::spill_dir`]
/// and [`QueryLimits::spill_quota_bytes`] configure where and how much
/// `DataFusion`'s own execution-time spill is allowed to write.
///
/// Not `Copy` — [`QueryLimits::spill_dir`] owns a [`PathBuf`] — so a caller
/// holding one behind a shared reference clones it explicitly where an
/// owned copy is needed rather than relying on an implicit bitwise copy.
#[derive(Debug, Clone)]
pub struct QueryLimits {
    /// Hard wall-clock ceiling on one query's execution. `crate::core_resolution`
    /// intersects this with the per-request timeout when it mints the
    /// effective bounds, so a request can only ever tighten this ceiling.
    /// The enforcement seam is `crate::core_execution`'s operation
    /// cancellation: a query still running past the deadline is cancelled,
    /// and the refusal surfaces as a typed deadline error, never a hang.
    pub timeout: Duration,
    /// Maximum rows one query returns, across every collected batch
    /// combined. `crate::core_resolution` intersects this with the
    /// per-request row bound when it mints the effective bounds, and
    /// `crate::core_execution` asks `DataFusion` for one row past the
    /// effective bound so a truncated result is flagged rather than
    /// silently clipped.
    pub row_cap: usize,
    /// Directory `DataFusion`'s spillable operators write their spill files
    /// under. `crate::core_service` passes it to
    /// `RuntimeEnvBuilder::with_temp_file_path` when it builds the
    /// composition's one `RuntimeEnv` — consumed exactly ONCE, the same as
    /// [`QueryLimits::spill_quota_bytes`].
    ///
    /// This crate's OWN [`Self::default`] keeps this portable — a
    /// [`std::env::temp_dir`]-rooted path safe for any embedder (a bare
    /// test, a standalone tool) — deliberately NOT a Kubernetes-specific
    /// absolute path, since [`QueryLimits`] carries no notion of "this
    /// deployment's real mounts." A real deployment names its own,
    /// deployment-specific value through the query service's own
    /// configuration. The default is also suffixed with this process's id
    /// (`default_query_spill_dir`, this module) — see that private
    /// function's doc for the concurrent-process directory-creation race a
    /// bare, unsuffixed shared path invites.
    ///
    /// # This directory is created EAGERLY, not lazily — a bad path fails composition
    ///
    /// `RuntimeEnvBuilder::build_arc` (called once, from
    /// `crate::core_service::ProjectedCoreService`'s composition) creates
    /// this directory (and an initial working subdirectory inside it)
    /// IMMEDIATELY — verified against `datafusion-execution`'s own
    /// `disk_manager.rs::create_local_dirs`, which calls
    /// `std::fs::create_dir` synchronously during `DiskManager::try_new`. It
    /// is NOT deferred to first spill. A path the process cannot create or
    /// write to fails composition — service STARTUP — not merely the first
    /// spilling query. Any deployment pointing this field at a
    /// Kubernetes-specific mount MUST ensure that mount already exists
    /// before the service composes.
    pub spill_dir: PathBuf,
    /// Disk quota, in bytes, `DataFusion`'s `DiskManager` enforces against
    /// [`QueryLimits::spill_dir`], via
    /// `RuntimeEnvBuilder::with_max_temp_directory_size`. `crate::core_service`
    /// refuses a zero value at composition: `DataFusion` accepts it and
    /// then fails every spilling query at execution time, which reads as a
    /// healthy pod that answers some queries and not others.
    ///
    /// This crate's OWN [`Self::default`] sizes this as a generic backstop
    /// (16 GiB) against a spilling query filling whatever
    /// [`QueryLimits::spill_dir`] resolves to — appropriate for that field's
    /// own portable, non-Kubernetes-specific default. A real deployment
    /// names its own, tighter, TIER-ORDERED value through the query
    /// service's own configuration: the quota must sit below the volume's
    /// own limit, which must sit below the pod's.
    pub spill_quota_bytes: u64,
}

impl Default for QueryLimits {
    /// 30 s / 10,000 rows / a portable temp-dir spill path / 16 GiB spill
    /// quota.
    ///
    /// These are portable, conservative starting points, not measured
    /// optima. `row_cap` and `timeout` are generous for an interactive
    /// query and bound the wire payload and the wait without a caller
    /// needing to remember to add `LIMIT`. `spill_dir`/`spill_quota_bytes`
    /// are deliberately portable, NOT the Kubernetes-specific values a
    /// real deployment should use — see each field's own doc for where the
    /// deployment-specific values live instead.
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(30),
            row_cap: 10_000,
            spill_dir: default_query_spill_dir(),
            spill_quota_bytes: DEFAULT_QUERY_SPILL_QUOTA_BYTES,
        }
    }
}

/// Default [`QueryLimits::spill_dir`] directory name, created directly under
/// [`std::env::temp_dir`] — portable, always-writable, no Kubernetes
/// assumption. See that field's own doc for why a real deployment overrides
/// this through its own config layer instead of relying on it as-is.
const DEFAULT_QUERY_SPILL_DIR_NAME: &str = "polychrome-query-spill";

/// [`DEFAULT_QUERY_SPILL_DIR_NAME`], suffixed with this process's id.
///
/// `datafusion-execution`'s `disk_manager.rs::create_local_dirs` (the
/// function [`QueryLimits::spill_dir`]'s own doc names as creating this
/// directory eagerly) is `if !root.exists() { std::fs::create_dir(root)? }`
/// — a check-then-create with no atomicity between the two. A BARE
/// `DEFAULT_QUERY_SPILL_DIR_NAME` under `std::env::temp_dir()` is one fixed
/// path shared by every consumer that takes this default, and this crate's
/// own test suite alone constructs a `QueryLimits` from many
/// `#[tokio::test]`s. `cargo nextest` runs each as its OWN process, so two
/// tests racing to compose a service at the same instant can both observe
/// the shared path absent, then have one process's `create_dir` fail with
/// `AlreadyExists` — a nondeterministic startup refusal. Per-process
/// uniqueness removes the shared path two DIFFERENT, concurrently-running
/// processes could ever race on — the failure mode this exists to prevent —
/// without changing the "portable, always-writable, no Kubernetes
/// assumption" contract the field's own doc promises: a process id is
/// exactly as portable as the bare name it replaces.
fn default_query_spill_dir() -> PathBuf {
    std::env::temp_dir().join(format!(
        "{DEFAULT_QUERY_SPILL_DIR_NAME}-{}",
        std::process::id()
    ))
}

/// Default [`QueryLimits::spill_quota_bytes`] — 16 GiB, a generic backstop
/// sized for [`DEFAULT_QUERY_SPILL_DIR_NAME`]'s own portable default, not a
/// Kubernetes deployment's tier-ordered quota. See that field's own doc for
/// where the deployment-specific value lives instead.
const DEFAULT_QUERY_SPILL_QUOTA_BYTES: u64 = 16 * 1024 * 1024 * 1024;

/// Owns [`QueryLimits::spill_dir`] for the life of the
/// [`crate::core_service::ProjectedCoreService`] that built its `RuntimeEnv`,
/// and removes it on drop.
///
/// `RuntimeEnvBuilder::build_arc` creates `spill_dir` eagerly at composition
/// (see [`QueryLimits::spill_dir`]'s own doc). `DataFusion`'s `DiskManager`
/// never removes its own working directory afterward. Nothing else reclaims
/// it either. A Kubernetes deployment points `spill_dir` at a dedicated
/// `emptyDir`, which the pod lifecycle cleans up on its own. The portable
/// default above has no such reclaim. One `ProjectedCoreService` is composed
/// once per process, and normally held for the process's whole life. So this
/// guard's drop lines up with process shutdown — `polychrome stop`, a
/// graceful restart — rather than firing mid-run (POLY-314).
///
/// Like every `Drop`-based guard, this only runs on a graceful shutdown
/// path. It does nothing against `SIGKILL`. But that covers the common case
/// a leaked spill directory is actually seen in.
pub(crate) struct SpillDirGuard(PathBuf);

impl SpillDirGuard {
    pub(crate) const fn new(spill_dir: PathBuf) -> Self {
        Self(spill_dir)
    }
}

impl Drop for SpillDirGuard {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.0);
    }
}

#[cfg(test)]
mod tests {
    use super::SpillDirGuard;

    /// Dropping the guard removes the directory it was built with, the same
    /// way `RuntimeEnvBuilder::build_arc` creates it eagerly at composition
    /// (POLY-314).
    #[test]
    fn drop_removes_the_spill_directory() {
        let dir = std::env::temp_dir().join(format!(
            "polyc-query-spill-guard-test-{}",
            std::process::id()
        ));
        std::fs::create_dir_all(&dir).expect("create scratch spill dir");
        assert!(dir.exists(), "precondition: scratch dir exists");

        drop(SpillDirGuard::new(dir.clone()));

        assert!(!dir.exists(), "guard removes the spill directory on drop");
    }

    /// A guard over a directory that never existed drops without panicking —
    /// `remove_dir_all`'s error is intentionally swallowed.
    #[test]
    fn drop_tolerates_a_missing_directory() {
        let dir = std::env::temp_dir().join(format!(
            "polyc-query-spill-guard-missing-test-{}",
            std::process::id()
        ));
        assert!(!dir.exists(), "precondition: scratch dir absent");

        drop(SpillDirGuard::new(dir));
    }
}