diff --git a/libs/@local/integrations/engine/src/byte_size.rs b/libs/@local/integrations/engine/src/byte_size.rs new file mode 100644 index 00000000000..1214e036b55 --- /dev/null +++ b/libs/@local/integrations/engine/src/byte_size.rs @@ -0,0 +1,18 @@ +/// An amount of memory or disk space, in bytes. +/// +/// It displays as ` bytes`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, derive_more::Display)] +#[display("{_0} bytes")] +pub struct ByteSize(u64); + +impl ByteSize { + #[must_use] + pub const fn from_bytes(bytes: u64) -> Self { + Self(bytes) + } + + /// Returns the size in the form DuckDB reads for a size setting. + pub(crate) fn setting(self) -> Box { + format!("{} bytes", self.0).into_boxed_str() + } +} diff --git a/libs/@local/integrations/engine/src/lib.rs b/libs/@local/integrations/engine/src/lib.rs index ff4376b05ca..fdb9775bf5a 100644 --- a/libs/@local/integrations/engine/src/lib.rs +++ b/libs/@local/integrations/engine/src/lib.rs @@ -15,11 +15,15 @@ extern crate alloc; +mod byte_size; mod sql; mod variant_name; mod workspace; -pub use self::workspace::{ - CloseError, DataErrorType, OpenError, QueryErrorType, Rows, Setting, SnapshotError, - StatementError, Workspace, WorkspaceConfig, +pub use self::{ + byte_size::ByteSize, + workspace::{ + CloseError, DataErrorType, DiskUsageError, OpenError, QueryErrorType, Rows, Setting, + SnapshotError, StatementError, Workspace, WorkspaceConfig, + }, }; diff --git a/libs/@local/integrations/engine/src/workspace/connection.rs b/libs/@local/integrations/engine/src/workspace/connection.rs index 1a13a9738da..ccf46a51eef 100644 --- a/libs/@local/integrations/engine/src/workspace/connection.rs +++ b/libs/@local/integrations/engine/src/workspace/connection.rs @@ -11,9 +11,11 @@ use error_stack::{Report, ResultExt as _}; use tokio::sync::mpsc; use crate::{ + ByteSize, WorkspaceConfig, sql::Identifier, workspace::{ - CloseError, Command, OpenError, Rows, Setting, SnapshotError, StatementError, + CloseError, Command, DiskUsageError, OpenError, Rows, Setting, SnapshotError, + StatementError, files::WorkspaceFiles, interrupt::{Interrupt, InterruptGuard, RequestId}, reply::Reply, @@ -21,6 +23,12 @@ use crate::{ }, }; +/// The disk space the workspace files use, and the limit they are checked against. +struct DiskUsage { + used: ByteSize, + limit: ByteSize, +} + /// The database's only connection, owned by the workspace thread. /// /// Fields are dropped in declaration order, so `interrupt` stops interrupts before `duckdb` closes @@ -29,6 +37,7 @@ pub(super) struct Connection { interrupt: InterruptGuard, duckdb: duckdb::Connection, files: WorkspaceFiles, + disk_limit: Option, /// An exclusive lock on the workspace's lock file, held until the database is closed. lock: File, } @@ -60,37 +69,126 @@ impl Connection { } } - /// Opens the database at `path` with its settings locked. - pub(super) fn open(path: PathBuf) -> Result> { + /// Returns the size of the file at `path`, the total size of the files under it if it is a + /// directory, or zero if it does not exist. + fn file_size(path: &Path) -> io::Result { + let metadata = match fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(0), + Err(error) => return Err(error), + }; + if !metadata.is_dir() { + return Ok(metadata.len()); + } + + let mut total = 0_u64; + for entry in fs::read_dir(path)? { + total = total.saturating_add(Self::file_size(&entry?.path())?); + } + Ok(total) + } + + /// Returns the disk space the workspace files use. + fn workspace_size(&self) -> Result> { + let mut used = 0_u64; + for path in self.files.measured() { + let size = Self::file_size(path).change_context_lazy(|| DiskUsageError::Measure { + path: path.to_owned(), + })?; + used = used.saturating_add(size); + } + Ok(ByteSize::from_bytes(used)) + } + + /// Returns the disk space the workspace files use, or `None` if there is no disk limit. + /// + /// When the files are over the limit, a checkpoint writes the log into the database file and + /// truncates the log, and the files are measured again. + fn disk_usage(&self) -> Result, Report> { + let Some(limit) = self.disk_limit else { + return Ok(None); + }; + let mut used = self.workspace_size()?; + if used > limit { + self.duckdb + .execute_batch("CHECKPOINT") + .change_context(DiskUsageError::Checkpoint)?; + used = self.workspace_size()?; + } + Ok(Some(DiskUsage { used, limit })) + } + + /// Checks the workspace files against the disk limit after a request. + fn check_disk_limit(&self) -> Result<(), Report> { + match self + .disk_usage() + .change_context(StatementError::DiskUsage)? + { + Some(DiskUsage { used, limit }) if used > limit => { + Err(Report::new(StatementError::DiskLimit { used, limit })) + } + Some(_) | None => Ok(()), + } + } + + /// Opens the database at `path` with the limits in `config` and its settings locked. + pub(super) fn open(path: PathBuf, config: &WorkspaceConfig) -> Result> { let files = WorkspaceFiles::new(path); let lock = Self::lock(&files)?; - let mut config = Config::default(); - for (setting, value) in [ - (Setting::AllowCommunityExtensions, "false"), - (Setting::AutoinstallKnownExtensions, "false"), - (Setting::AutoloadKnownExtensions, "false"), - (Setting::EnableExternalAccess, "false"), - (Setting::LockConfiguration, "true"), - ] { - config = config - .with(setting.to_string(), value) - .change_context_lazy(|| OpenError::Configure { - setting, - value: value.into(), - })?; + let limits = [ + ( + Setting::MemoryLimit, + config.memory_limit.map(ByteSize::setting), + ), + ( + Setting::MaxTempDirectorySize, + config.temp_directory_limit.map(ByteSize::setting), + ), + ( + Setting::Threads, + config + .threads + .map(|threads| threads.to_string().into_boxed_str()), + ), + ]; + let settings = [ + (Setting::AllowCommunityExtensions, Some("false".into())), + (Setting::AutoinstallKnownExtensions, Some("false".into())), + (Setting::AutoloadKnownExtensions, Some("false".into())), + (Setting::EnableExternalAccess, Some("false".into())), + ] + .into_iter() + .chain(limits) + .chain([(Setting::LockConfiguration, Some("true".into()))]); + + let mut duckdb_config = Config::default(); + for (setting, value) in settings { + let Some(value) = value else { + continue; + }; + duckdb_config = duckdb_config + .with(setting.to_string(), &value) + .change_context_lazy(|| OpenError::Configure { setting, value })?; } - let duckdb = duckdb::Connection::open_with_flags(files.database(), config) + let duckdb = duckdb::Connection::open_with_flags(files.database(), duckdb_config) .change_context_lazy(|| OpenError::Open { path: files.database().to_owned(), })?; - Ok(Self { + let opened = Self { interrupt: InterruptGuard::new(Arc::new(Interrupt::new(duckdb.interrupt_handle()))), duckdb, files, + disk_limit: config.disk_limit, lock, - }) + }; + match opened.disk_usage().change_context(OpenError::DiskUsage)? { + Some(DiskUsage { used, limit }) if used > limit => { + Err(Report::new(OpenError::DiskLimit { used, limit })) + } + Some(_) | None => Ok(opened), + } } /// Returns the [`Interrupt`] that cancels the requests running on this connection. @@ -147,6 +245,7 @@ impl Connection { .into_iter() .map(String::into_boxed_str) .collect(); + self.check_disk_limit()?; Ok(Rows::new(columns, rows)) } @@ -162,7 +261,7 @@ impl Connection { )) .and_then(|mut statement| statement.execute([])) .map_err(StatementError::from_duckdb)?; - Ok(()) + self.check_disk_limit() } /// Copies the database file to `partial`, syncs it and renames it to `target`. @@ -300,16 +399,18 @@ impl Connection { } } - /// Opens the database at `path`, sends the result on `ready`, and runs `commands`. + /// Opens the database at `path` with the limits in `config`, sends the result on `ready`, and + /// runs `commands`. /// /// A panic while the commands run is logged. The connection is dropped while the panic /// unwinds, so nothing it holds is used after the panic. pub(super) fn serve( path: PathBuf, + config: &WorkspaceConfig, ready: Reply, OpenError>, commands: &mut mpsc::Receiver<(RequestId, Command)>, ) { - let connection = match Self::open(path) { + let connection = match Self::open(path, config) { Ok(connection) => connection, Err(report) => { ready.send(Err(report)); diff --git a/libs/@local/integrations/engine/src/workspace/error.rs b/libs/@local/integrations/engine/src/workspace/error.rs index becec93c3e2..ff429a4da7b 100644 --- a/libs/@local/integrations/engine/src/workspace/error.rs +++ b/libs/@local/integrations/engine/src/workspace/error.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use error_stack::Report; -use crate::variant_name::VariantName; +use crate::{ByteSize, variant_name::VariantName}; /// A DuckDB setting the workspace applies when it opens the database. /// @@ -15,6 +15,9 @@ pub enum Setting { AutoloadKnownExtensions, EnableExternalAccess, LockConfiguration, + MemoryLimit, + MaxTempDirectorySize, + Threads, } impl Setting { @@ -31,6 +34,16 @@ impl fmt::Display for Setting { } } +/// Reports why the disk space of the workspace files could not be checked against +/// [`WorkspaceConfig::disk_limit`](crate::WorkspaceConfig::disk_limit). +#[derive(Debug, derive_more::Display, derive_more::Error)] +pub enum DiskUsageError { + #[display("could not read the disk usage of `{}`", path.display())] + Measure { path: PathBuf }, + #[display("could not checkpoint the database")] + Checkpoint, +} + /// Reports why [`Workspace::open`](crate::Workspace::open) failed. #[derive(Debug, derive_more::Display, derive_more::Error)] pub enum OpenError { @@ -42,6 +55,10 @@ pub enum OpenError { Configure { setting: Setting, value: Box }, #[display("could not open the database at `{}`", path.display())] Open { path: PathBuf }, + #[display("the workspace files use {used}, over the disk limit of {limit}")] + DiskLimit { used: ByteSize, limit: ByteSize }, + #[display("could not check the workspace files against the disk limit")] + DiskUsage, #[display("could not start the workspace thread")] Spawn, #[display("the workspace thread stopped before the database opened")] @@ -180,6 +197,12 @@ pub enum StatementError { /// could not be read. #[display("the statement failed")] Failed, + #[display( + "after the statement, the workspace files use {used}, over the disk limit of {limit}" + )] + DiskLimit { used: ByteSize, limit: ByteSize }, + #[display("could not check the workspace files against the disk limit after the statement")] + DiskUsage, #[display("the workspace thread stopped before the statement finished")] Stopped, } diff --git a/libs/@local/integrations/engine/src/workspace/files.rs b/libs/@local/integrations/engine/src/workspace/files.rs index a62a6cc372c..77a3666485a 100644 --- a/libs/@local/integrations/engine/src/workspace/files.rs +++ b/libs/@local/integrations/engine/src/workspace/files.rs @@ -60,6 +60,12 @@ impl WorkspaceFiles { &self.lock } + /// Returns the files that the disk limit counts: the database file, its log and the + /// temporary directory. + pub(super) fn measured(&self) -> [&Path; 3] { + [&self.database, &self.log, &self.temp] + } + /// Returns whether `path` names one of the files, once both directories are resolved. pub(super) fn contains(&self, path: &Path) -> bool { let Some(path) = Self::resolve(path) else { diff --git a/libs/@local/integrations/engine/src/workspace/mod.rs b/libs/@local/integrations/engine/src/workspace/mod.rs index fe68e4ebde1..2928904550e 100644 --- a/libs/@local/integrations/engine/src/workspace/mod.rs +++ b/libs/@local/integrations/engine/src/workspace/mod.rs @@ -15,23 +15,42 @@ use tokio::sync::mpsc; pub use self::{ error::{ - CloseError, DataErrorType, OpenError, QueryErrorType, Setting, SnapshotError, - StatementError, + CloseError, DataErrorType, DiskUsageError, OpenError, QueryErrorType, Setting, + SnapshotError, StatementError, }, rows::Rows, }; -use crate::workspace::{ - connection::Connection, - interrupt::{Interrupt, Pending, RequestId}, - reply::Reply, +use crate::{ + ByteSize, + workspace::{ + connection::Connection, + interrupt::{Interrupt, Pending, RequestId}, + reply::Reply, + }, }; /// Settings for a [`Workspace`]. -#[derive(Debug, Clone, Copy)] +#[derive(Debug)] pub struct WorkspaceConfig { /// The number of requests that can wait while another runs. Defaults to 64. /// Callers wait for queue space when it is full. pub queue_capacity: NonZeroUsize, + /// The most memory DuckDB uses. A query that needs more spills to the temporary directory, + /// or fails if it cannot. + pub memory_limit: Option, + /// The most disk space DuckDB uses for the data it spills to its temporary directory. + pub temp_directory_limit: Option, + /// The disk space the workspace files can use before a request fails. The workspace files + /// are the database file, its log and DuckDB's temporary directory. + /// + /// The workspace checks the limit when it opens and after each [`query`](Workspace::query) + /// and [`create_table`](Workspace::create_table), so a single request can write past it. + /// When the files are over the limit, it checkpoints the database, which writes the log into + /// the database file, and measures them again. A table created by a request that takes the + /// files over the limit stays in the database. + pub disk_limit: Option, + /// The number of threads DuckDB runs a query on. + pub threads: Option, } impl Default for WorkspaceConfig { @@ -39,6 +58,10 @@ impl Default for WorkspaceConfig { Self { queue_capacity: NonZeroUsize::new(64) .expect("the default queue capacity should be nonzero"), + memory_limit: None, + temp_directory_limit: None, + disk_limit: None, + threads: None, } } } @@ -97,6 +120,8 @@ impl Workspace { /// - [`Locked`] if another workspace has the database open /// - [`Configure`] if DuckDB refuses a setting /// - [`Open`] if DuckDB cannot open the database + /// - [`DiskLimit`] if the workspace files are over [`WorkspaceConfig::disk_limit`] + /// - [`DiskUsage`] if the workspace files cannot be checked against the disk limit /// - [`Spawn`] if the workspace thread cannot be started /// - [`Stopped`] if the workspace thread stopped before it opened the database /// @@ -104,6 +129,8 @@ impl Workspace { /// [`Locked`]: OpenError::Locked /// [`Configure`]: OpenError::Configure /// [`Open`]: OpenError::Open + /// [`DiskLimit`]: OpenError::DiskLimit + /// [`DiskUsage`]: OpenError::DiskUsage /// [`Spawn`]: OpenError::Spawn /// [`Stopped`]: OpenError::Stopped pub async fn open(path: PathBuf, config: WorkspaceConfig) -> Result> { @@ -115,7 +142,7 @@ impl Workspace { .name("integrations-workspace".to_owned()) .spawn(move || { let _entered = span.entered(); - Connection::serve(path, ready, &mut receiver); + Connection::serve(path, &config, ready, &mut receiver); }) .change_context(OpenError::Spawn)?; @@ -159,6 +186,9 @@ impl Workspace { /// - [`InvalidQuery`] if DuckDB cannot run the query as written /// - [`InvalidData`] if a value in the rows causes an error /// - [`Failed`] if DuckDB fails for another reason or a value cannot be read + /// - [`DiskLimit`] if the workspace files are over [`WorkspaceConfig::disk_limit`] after the + /// query + /// - [`DiskUsage`] if the workspace files cannot be checked against the disk limit /// - [`Stopped`] if the workspace thread stopped /// /// [`NotSelect`]: StatementError::NotSelect @@ -168,6 +198,8 @@ impl Workspace { /// [`InvalidQuery`]: StatementError::InvalidQuery /// [`InvalidData`]: StatementError::InvalidData /// [`Failed`]: StatementError::Failed + /// [`DiskLimit`]: StatementError::DiskLimit + /// [`DiskUsage`]: StatementError::DiskUsage /// [`Stopped`]: StatementError::Stopped pub async fn query(&self, sql: Box) -> Result> { self.request( @@ -187,6 +219,9 @@ impl Workspace { /// - [`InvalidQuery`] if DuckDB cannot run the query as written /// - [`InvalidData`] if a value in the rows causes an error /// - [`Failed`] if DuckDB fails for another reason + /// - [`DiskLimit`] if the workspace files are over [`WorkspaceConfig::disk_limit`] after the + /// table is created + /// - [`DiskUsage`] if the workspace files cannot be checked against the disk limit /// - [`Stopped`] if the workspace thread stopped /// /// [`NotSelect`]: StatementError::NotSelect @@ -195,6 +230,8 @@ impl Workspace { /// [`InvalidQuery`]: StatementError::InvalidQuery /// [`InvalidData`]: StatementError::InvalidData /// [`Failed`]: StatementError::Failed + /// [`DiskLimit`]: StatementError::DiskLimit + /// [`DiskUsage`]: StatementError::DiskUsage /// [`Stopped`]: StatementError::Stopped pub async fn create_table( &self, diff --git a/libs/@local/integrations/engine/tests/workspace.rs b/libs/@local/integrations/engine/tests/workspace.rs index 09234840a7b..87bb698c66e 100644 --- a/libs/@local/integrations/engine/tests/workspace.rs +++ b/libs/@local/integrations/engine/tests/workspace.rs @@ -9,22 +9,28 @@ use std::{ use duckdb::types::{TimeUnit, Value}; use error_stack::Report; use hash_integrations_engine::{ - DataErrorType, OpenError, QueryErrorType, SnapshotError, StatementError, Workspace, + ByteSize, DataErrorType, OpenError, QueryErrorType, SnapshotError, StatementError, Workspace, WorkspaceConfig, }; use pretty_assertions::assert_eq; use tempfile::TempDir; +const MIB: u64 = 1024 * 1024; + fn database_path(directory: &TempDir) -> PathBuf { directory.path().join("workspace.duckdb") } -async fn open(directory: &TempDir) -> Workspace { - Workspace::open(database_path(directory), WorkspaceConfig::default()) +async fn open_with(directory: &TempDir, config: WorkspaceConfig) -> Workspace { + Workspace::open(database_path(directory), config) .await .expect("the workspace should open") } +async fn open(directory: &TempDir) -> Workspace { + open_with(directory, WorkspaceConfig::default()).await +} + async fn first_value(workspace: &Workspace, sql: &str) -> Result> { let rows = workspace.query(sql.into()).await?; let row = rows.iter().next().expect("the query should return a row"); @@ -398,6 +404,7 @@ async fn dropping_queued_request_skips_it() { database_path(&directory), WorkspaceConfig { queue_capacity: NonZeroUsize::MIN, + ..WorkspaceConfig::default() }, ) .await @@ -597,3 +604,170 @@ async fn snapshot_replaces_link_at_partial_path() { "the snapshot should hold the table" ); } + +#[tokio::test] +async fn open_applies_limits() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open_with( + &directory, + WorkspaceConfig { + memory_limit: Some(ByteSize::from_bytes(256 * MIB)), + temp_directory_limit: Some(ByteSize::from_bytes(1024 * MIB)), + threads: NonZeroUsize::new(2), + ..WorkspaceConfig::default() + }, + ) + .await; + + let rows = workspace + .query( + "SELECT current_setting('memory_limit'), current_setting('max_temp_directory_size'), \ + current_setting('threads')" + .into(), + ) + .await + .expect("the settings should be readable"); + + assert_eq!( + rows.iter().next(), + Some( + [ + Value::Text("256.0 MiB".to_owned()), + Value::Text("1.0 GiB".to_owned()), + Value::BigInt(2), + ] + .as_slice() + ), + "DuckDB should use the configured limits" + ); +} + +#[tokio::test] +async fn open_refuses_files_over_disk_limit() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open(&directory).await; + workspace + .create_table("flights".into(), "SELECT 'BA117' AS code".into()) + .await + .expect("the table should be created"); + workspace.close().await.expect("the workspace should close"); + + let report = Workspace::open( + database_path(&directory), + WorkspaceConfig { + disk_limit: Some(ByteSize::from_bytes(1)), + ..WorkspaceConfig::default() + }, + ) + .await + .expect_err("files over the disk limit should not open"); + + assert!( + matches!( + report.current_context(), + OpenError::DiskLimit { used, limit } + if *limit == ByteSize::from_bytes(1) && used > limit + ), + "the size and the limit should be reported, not {report:?}" + ); +} + +#[tokio::test] +async fn create_table_over_disk_limit_keeps_table() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open_with( + &directory, + WorkspaceConfig { + disk_limit: Some(ByteSize::from_bytes(2 * MIB)), + ..WorkspaceConfig::default() + }, + ) + .await; + + let report = workspace + .create_table( + "numbers".into(), + "SELECT random() AS value FROM range(300000)".into(), + ) + .await + .expect_err("a table that takes the files over the limit should be reported"); + + assert!( + matches!( + report.current_context(), + StatementError::DiskLimit { used, limit } if used > limit + ), + "the disk limit should be reported, not {report:?}" + ); + workspace.close().await.expect("the workspace should close"); + assert_eq!( + count_rows(&database_path(&directory), "numbers"), + 300_000, + "the table should stay in the database" + ); +} + +#[tokio::test] +async fn disk_limit_counts_temp_directory() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open_with( + &directory, + WorkspaceConfig { + disk_limit: Some(ByteSize::from_bytes(2 * MIB)), + ..WorkspaceConfig::default() + }, + ) + .await; + let temp = directory.path().join("workspace.duckdb.tmp"); + fs::create_dir_all(&temp).expect("should create the temporary directory"); + let spilled = vec![0_u8; usize::try_from(3 * MIB).expect("the size should fit")]; + fs::write(temp.join("spilled"), spilled).expect("should write a file"); + + let report = workspace + .query("SELECT 1".into()) + .await + .expect_err("a file in the temporary directory should count towards the limit"); + + assert!( + matches!(report.current_context(), StatementError::DiskLimit { .. }), + "the disk limit should be reported, not {report:?}" + ); +} + +#[tokio::test] +async fn disk_limit_checkpoints_before_refusing() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open_with( + &directory, + WorkspaceConfig { + disk_limit: Some(ByteSize::from_bytes(MIB)), + ..WorkspaceConfig::default() + }, + ) + .await; + + // Each table adds about 120 KB to the log, so the log alone passes the limit by the ninth. + for count in 0..20 { + workspace + .create_table( + "counters".into(), + format!("SELECT range AS id, {count} AS count FROM range(10000)").into(), + ) + .await + .expect("each table should fit once the log is checkpointed"); + } + + let log = fs::metadata(directory.path().join("workspace.duckdb.wal")) + .map_or(0, |metadata| metadata.len()); + assert!( + log < MIB, + "the log should be checkpointed before it passes the limit, but it holds {log} bytes" + ); + assert_eq!( + first_value(&workspace, "SELECT min(count) FROM counters") + .await + .expect("the query should run"), + Value::Int(19), + "the last table should be kept" + ); +}