tableski_core/
limits.rs

1//! Per-query limits: a memory pool cap on the session, a wall-clock timeout, and caps on the
2//! rows and bytes a result may carry back. Unlimited by default for embedders; the CLI applies
3//! [`QueryLimits::safe`] unless told otherwise.
4
5use datafusion::arrow::record_batch::RecordBatch;
6use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
7use datafusion::execution::runtime_env::RuntimeEnvBuilder;
8use datafusion::prelude::{DataFrame, SessionConfig, SessionContext};
9use futures::StreamExt;
10use std::future::Future;
11use std::sync::OnceLock;
12use std::time::{Duration, Instant};
13
14/// The process-wide tokio runtime queries execute on when a timeout is set: one thread per
15/// core, separate from whatever runtime serves clients, so a busy query cannot starve the
16/// server. Built on first use.
17pub fn query_runtime() -> &'static tokio::runtime::Runtime {
18    static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
19    RUNTIME.get_or_init(|| {
20        tokio::runtime::Builder::new_multi_thread()
21            .worker_threads(std::thread::available_parallelism().map_or(2, |n| n.get()))
22            .thread_name("tableski-query")
23            .enable_all()
24            .build()
25            .expect("query runtime")
26    })
27}
28
29/// Bounds on what one query may consume and return. `None` means unlimited.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub struct QueryLimits {
32    /// Memory pool for the whole session, in bytes; every query draws from it and fails with
33    /// `Resources exhausted` instead of growing the process. With a cap set, spilling to disk
34    /// is disabled too, so a hostile query fails fast rather than filling the disk.
35    pub memory_bytes: Option<usize>,
36    /// Wall-clock budget for planning plus execution of one query.
37    pub timeout: Option<Duration>,
38    /// Most rows a result may carry; the rest is dropped and the result marked truncated.
39    pub max_rows: Option<usize>,
40    /// Most bytes (Arrow in-memory size) a result may carry; same truncation rule.
41    pub max_bytes: Option<usize>,
42}
43
44impl QueryLimits {
45    /// The CLI defaults: 2 GiB pool, 60 s, 10 000 rows, 16 MiB.
46    pub const SAFE_MEMORY_BYTES: usize = 2 * 1024 * 1024 * 1024;
47    /// See [`Self::safe`].
48    pub const SAFE_TIMEOUT: Duration = Duration::from_secs(60);
49    /// See [`Self::safe`].
50    pub const SAFE_MAX_ROWS: usize = 10_000;
51    /// See [`Self::safe`].
52    pub const SAFE_MAX_BYTES: usize = 16 * 1024 * 1024;
53
54    /// No limits at all: the pre-0.2 behaviour, and the default of [`AppState`](crate::AppState).
55    pub fn unlimited() -> Self {
56        Self::default()
57    }
58
59    /// Defaults that keep one query from taking the machine down; what the CLI uses.
60    pub fn safe() -> Self {
61        Self {
62            memory_bytes: Some(Self::SAFE_MEMORY_BYTES),
63            timeout: Some(Self::SAFE_TIMEOUT),
64            max_rows: Some(Self::SAFE_MAX_ROWS),
65            max_bytes: Some(Self::SAFE_MAX_BYTES),
66        }
67    }
68
69    /// A session whose memory pool honours [`Self::memory_bytes`] (unbounded when `None`).
70    /// Register tables into this context, then build the state with it.
71    pub fn session_context(&self) -> Result<SessionContext, String> {
72        let mut runtime = RuntimeEnvBuilder::new();
73        if let Some(bytes) = self.memory_bytes {
74            runtime = runtime
75                .with_memory_limit(bytes, 1.0)
76                .with_disk_manager_builder(
77                    DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled),
78                );
79        }
80        let runtime = runtime.build_arc().map_err(|e| e.to_string())?;
81        Ok(SessionContext::new_with_config_rt(
82            SessionConfig::new(),
83            runtime,
84        ))
85    }
86}
87
88/// What [`collect_limited`] brought back.
89#[derive(Debug, Clone, Default)]
90pub struct Collected {
91    /// The result batches, cut at the caps.
92    pub batches: Vec<RecordBatch>,
93    /// Rows in `batches`.
94    pub rows: usize,
95    /// Arrow in-memory size of `batches`.
96    pub bytes: usize,
97    /// Set when [`QueryLimits::max_rows`] or [`QueryLimits::max_bytes`] cut the result.
98    pub truncated: bool,
99}
100
101impl Collected {
102    /// One line telling the client what was cut and how to get the rest, or `None`.
103    pub fn truncation_note(&self, limits: &QueryLimits) -> Option<String> {
104        if !self.truncated {
105            return None;
106        }
107        let cap = match (limits.max_rows, limits.max_bytes) {
108            (Some(r), _) if self.rows >= r => format!("{r} rows"),
109            (_, Some(b)) => format!("{} MiB", b / (1024 * 1024)),
110            (Some(r), None) => format!("{r} rows"),
111            (None, None) => "the configured cap".to_string(),
112        };
113        Some(format!(
114            "-- result truncated at {cap} ({} rows returned); add a LIMIT / WHERE, aggregate, or use export_result",
115            self.rows
116        ))
117    }
118}
119
120/// Run `work` under `timeout` on the [`query_runtime`]: the caller's runtime (HTTP, timers)
121/// stays responsive however busy the work is, and on expiry the task is aborted (it stops at
122/// its next yield point). Without a timeout, `work` runs inline.
123pub async fn run_bounded<T, F>(work: F, timeout: Option<Duration>, what: &str) -> Result<T, String>
124where
125    T: Send + 'static,
126    F: Future<Output = Result<T, String>> + Send + 'static,
127{
128    let Some(budget) = timeout else {
129        return work.await;
130    };
131    let started = Instant::now();
132    let mut task = query_runtime().spawn(work);
133    match tokio::time::timeout(budget, &mut task).await {
134        Ok(joined) => joined.map_err(|e| format!("{what} task failed: {e}"))?,
135        Err(_) => {
136            task.abort();
137            Err(format!(
138                "{what} cancelled after {:.1}s (wall-clock limit {:.0}s; narrow the query or raise --query-timeout-secs)",
139                started.elapsed().as_secs_f64(),
140                budget.as_secs_f64()
141            ))
142        }
143    }
144}
145
146/// Execute `df` as a stream, stop at the row/byte caps, and give up at the timeout.
147///
148/// Stopping early means a runaway `SELECT *` never materialises past the cap. With a timeout
149/// the query runs on [`query_runtime`] via [`run_bounded`]; the caller's runtime and the
150/// session stay usable. A plan that never yields keeps its query thread busy until it does
151/// (see [`check_plan`](crate::guard::check_plan) for why cartesian products are refused).
152pub async fn collect_limited(df: DataFrame, limits: &QueryLimits) -> Result<Collected, String> {
153    let (max_rows, max_bytes) = (limits.max_rows, limits.max_bytes);
154    run_bounded(
155        collect_capped(df, max_rows, max_bytes),
156        limits.timeout,
157        "query",
158    )
159    .await
160    .map_err(|e| {
161        if e.contains("Resources exhausted") {
162            format!("{e} (memory pool limit; narrow the query or raise --max-memory-mb)")
163        } else {
164            e
165        }
166    })
167}
168
169async fn collect_capped(
170    df: DataFrame,
171    max_rows: Option<usize>,
172    max_bytes: Option<usize>,
173) -> Result<Collected, String> {
174    let mut stream = df.execute_stream().await.map_err(|e| e.to_string())?;
175    let mut out = Collected::default();
176    while let Some(batch) = stream.next().await {
177        let mut batch = batch.map_err(|e| e.to_string())?;
178        if let Some(cap) = max_rows {
179            let room = cap.saturating_sub(out.rows);
180            if batch.num_rows() > room {
181                batch = batch.slice(0, room);
182                out.truncated = true;
183            }
184        }
185        if let Some(cap) = max_bytes {
186            let size = batch.get_array_memory_size();
187            if out.bytes + size > cap && !out.batches.is_empty() {
188                out.truncated = true;
189                break;
190            }
191            if out.bytes + size > cap {
192                out.truncated = true;
193            }
194        }
195        out.rows += batch.num_rows();
196        out.bytes += batch.get_array_memory_size();
197        if batch.num_rows() > 0 {
198            out.batches.push(batch);
199        }
200        if out.truncated {
201            break;
202        }
203    }
204    Ok(out)
205}
tableski_core/limits.rs

tableski_core/
limits.rs
