From fa3ee2e71072571863588bee4ecbad964a20e6fb Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Fri, 4 Sep 2026 14:12:36 +0200 Subject: [PATCH] feat: run queries on their own Tokio runtime One runtime served the API and ran every query. A scan holds a worker until a partition yields, so one long query took every thread and the API stopped answering. Queries now run on a second runtime. BEACON_API_THREADS sizes the API runtime. The result stream crosses between the two over a bounded channel. --- CHANGELOG.md | 9 + beacon-db/beacon-core/src/lib.rs | 1 + beacon-db/beacon-core/src/query_executor.rs | 297 ++++++++++++++++++ beacon-db/beacon-core/src/runtime.rs | 53 ++++ beacon-db/beacon-core/src/runtime_builder.rs | 10 +- .../beacon-server-config/src/error.rs | 4 + beacon-server/beacon-server-config/src/lib.rs | 52 ++- beacon-server/beacon-server/src/main.rs | 45 ++- beacon-server/beacon-server/src/server/mod.rs | 12 +- .../beacon-server/tests/common/mod.rs | 3 +- docs/docs/2.0.0-rc5/server/configuration.md | 3 +- .../2.0.0-rc5/server/performance-tuning.md | 12 +- 12 files changed, 483 insertions(+), 18 deletions(-) create mode 100644 beacon-db/beacon-core/src/query_executor.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 63d883fc..0c3e8b06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -293,6 +293,15 @@ tag. Releases before 2.0.0 are recorded in the ### Fixed +- **A long query no longer makes the API unreachable.** The HTTP API, Flight SQL and every query + shared one Tokio runtime of `BEACON_WORKER_THREADS` threads. A scan holds a thread until a + partition yields, and one query starts as many partitions as the machine has cores, so a long + query took every thread and a login, a health check or the admin UI waited for it to finish. + Queries now run on a runtime of their own. The API runs on a second runtime, sized by the new + `BEACON_API_THREADS` (default `4`), so it always has a thread to answer with. The result stream + crosses between the two over a bounded channel: a slow client holds at most a few batches ahead + of what it has read, and a client that disconnects cancels its query at once. A panic inside a + query reaches the client as an error instead of a result that ends early. - **A query could lose the rows of a file that changed after its analysis.** File statistics let a `WHERE` drop whole files before the scan opens them, on the column ranges a background pass recorded. That pass compares the size, the modification time and the etag of every listed file diff --git a/beacon-db/beacon-core/src/lib.rs b/beacon-db/beacon-core/src/lib.rs index f0fc0242..24353f90 100644 --- a/beacon-db/beacon-core/src/lib.rs +++ b/beacon-db/beacon-core/src/lib.rs @@ -6,6 +6,7 @@ pub mod file_stats; pub mod metrics; pub mod parser; pub mod query; +pub mod query_executor; pub(crate) mod query_metrics_store; pub mod query_result; pub mod runtime; diff --git a/beacon-db/beacon-core/src/query_executor.rs b/beacon-db/beacon-core/src/query_executor.rs new file mode 100644 index 00000000..72b0a193 --- /dev/null +++ b/beacon-db/beacon-core/src/query_executor.rs @@ -0,0 +1,297 @@ +//! The runtime that runs queries, kept apart from the runtime that serves them. +//! +//! A partition decode holds a worker until it yields, and Tokio cannot preempt +//! it. When queries and the API share one runtime, a long scan takes every +//! worker and every other request waits. The executor pins each query to a +//! runtime of its own. +//! +//! A plain `spawn` is not enough. DataFusion builds its stream lazily, and +//! `RepartitionExec` spawns its tasks onto the runtime that polls the stream. +//! [`QueryExecutor::bridge_stream`] therefore polls the stream on the query +//! runtime and hands each batch to the caller over a bounded channel. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{Stream, StreamExt}; +use tokio::runtime::Handle; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tracing::Instrument; + +/// Batches in flight between the query runtime and the caller. +/// +/// The producer waits when the channel is full, so a slow client holds at most +/// this many batches ahead of what it has read. +const BRIDGE_CAPACITY: usize = 4; + +/// A handle to the runtime that runs queries. +/// +/// Cheap to clone. The embedded database passes the one runtime it has. The +/// query and the caller then share a runtime, and the bridge costs one channel +/// hop per batch. +#[derive(Clone, Debug)] +pub struct QueryExecutor { + handle: Handle, +} + +impl QueryExecutor { + pub fn new(handle: Handle) -> Self { + Self { handle } + } + + /// The runtime this executor spawns onto. + pub fn handle(&self) -> &Handle { + &self.handle + } + + /// Runs `future` to completion on the query runtime. + /// + /// A panic in the future comes back as an error, not as a panic in the + /// caller. The current tracing span follows the future. + pub async fn run(&self, future: F) -> anyhow::Result + where + F: Future> + Send + 'static, + T: Send + 'static, + { + self.handle + .spawn(future.in_current_span()) + .await + .map_err(|error| anyhow::anyhow!("query task failed: {error}"))? + } + + /// Polls `stream` on the query runtime and yields its batches to the caller. + /// + /// The producer task ends when the stream ends or when the caller drops the + /// returned stream. A panic in the producer reaches the caller as an error, + /// so a short result never looks complete. + pub fn bridge_stream(&self, stream: SendableRecordBatchStream) -> SendableRecordBatchStream { + let schema = stream.schema(); + let (sender, receiver) = mpsc::channel(BRIDGE_CAPACITY); + let producer = self.handle.spawn( + async move { + let mut stream = stream; + while let Some(item) = stream.next().await { + // A closed channel means the caller is gone. Stop the query. + if sender.send(item).await.is_err() { + break; + } + } + } + .in_current_span(), + ); + Box::pin(RecordBatchStreamAdapter::new( + schema, + BridgedStream { + receiver, + producer: Some(producer), + }, + )) + } +} + +/// The receiving end of a bridged stream. +struct BridgedStream { + receiver: mpsc::Receiver>, + /// The producer task. `None` once it is joined. + producer: Option>, +} + +impl Stream for BridgedStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.receiver.poll_recv(cx) { + Poll::Ready(Some(item)) => Poll::Ready(Some(item)), + Poll::Pending => Poll::Pending, + // The channel is closed and empty, so the producer is done or about + // to be. Join it: a panic must reach the caller, not end the stream. + Poll::Ready(None) => { + let Some(producer) = self.producer.as_mut() else { + return Poll::Ready(None); + }; + match Pin::new(producer).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(outcome) => { + self.producer = None; + Poll::Ready(outcome.err().map(|error| { + Err(DataFusionError::Execution(format!( + "query execution task failed: {error}" + ))) + })) + } + } + } + } + } +} + +impl Drop for BridgedStream { + fn drop(&mut self) { + // A dropped result cancels the query at once, even inside a read. + if let Some(producer) = &self.producer { + producer.abort(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, Mutex}; + use std::time::{Duration, Instant}; + + use arrow::array::{AsArray, Int32Array}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; + use futures::TryStreamExt; + use tokio::runtime::{Builder, Runtime}; + + use super::*; + + const QUERY_THREAD: &str = "test-query"; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])) + } + + fn batch(value: i32) -> RecordBatch { + RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![value]))]).unwrap() + } + + fn source(stream: S) -> SendableRecordBatchStream + where + S: Stream> + Send + 'static, + { + Box::pin(RecordBatchStreamAdapter::new(schema(), stream)) + } + + /// An API runtime and a query runtime, as the server has them. + fn two_runtimes() -> (Runtime, QueryExecutor, Runtime) { + let api = Builder::new_current_thread().enable_all().build().unwrap(); + let query = Builder::new_multi_thread() + .worker_threads(2) + .thread_name(QUERY_THREAD) + .enable_all() + .build() + .unwrap(); + let executor = QueryExecutor::new(query.handle().clone()); + (api, executor, query) + } + + #[test] + fn batches_cross_runtimes_in_order() { + let (api, executor, _query) = two_runtimes(); + let batches = api.block_on(async { + let stream = executor.bridge_stream(source(futures::stream::iter( + (0..10).map(|value| Ok(batch(value))), + ))); + stream.try_collect::>().await.unwrap() + }); + let values: Vec = batches + .iter() + .map(|batch| batch.column(0).as_primitive::().value(0)) + .collect(); + assert_eq!(values, (0..10).collect::>()); + } + + #[test] + fn the_source_is_polled_on_the_query_runtime() { + let (api, executor, _query) = two_runtimes(); + let seen = Arc::new(Mutex::new(None)); + let record = seen.clone(); + api.block_on(async { + let stream = executor.bridge_stream(source(futures::stream::once(async move { + *record.lock().unwrap() = std::thread::current().name().map(String::from); + Ok(batch(1)) + }))); + stream.try_collect::>().await.unwrap(); + }); + assert_eq!(seen.lock().unwrap().as_deref(), Some(QUERY_THREAD)); + } + + #[test] + fn an_error_item_passes_through() { + let (api, executor, _query) = two_runtimes(); + let error = api.block_on(async { + let stream = executor.bridge_stream(source(futures::stream::iter([ + Ok(batch(1)), + Err(DataFusionError::Execution("scan failed".into())), + ]))); + stream.try_collect::>().await.unwrap_err() + }); + assert!(error.to_string().contains("scan failed"), "{error}"); + } + + #[test] + fn a_panic_in_the_producer_is_an_error_not_an_end() { + let (api, executor, _query) = two_runtimes(); + let error = api.block_on(async { + let stream = executor.bridge_stream(source(futures::stream::poll_fn( + |_: &mut Context<'_>| -> Poll>> { + panic!("decoder bug") + }, + ))); + stream.try_collect::>().await.unwrap_err() + }); + assert!( + error.to_string().contains("query execution task failed"), + "{error}" + ); + } + + #[test] + fn dropping_the_stream_cancels_the_producer() { + struct Guard(Arc); + impl Drop for Guard { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + let (api, executor, _query) = two_runtimes(); + let dropped = Arc::new(AtomicBool::new(false)); + let guard = Guard(dropped.clone()); + api.block_on(async { + // The source never yields, so only a cancel can release the guard. + let mut stream = executor.bridge_stream(source(futures::stream::poll_fn( + move |_: &mut Context<'_>| -> Poll>> { + let _held = &guard; + Poll::Pending + }, + ))); + assert!(futures::poll!(stream.next()).is_pending()); + drop(stream); + }); + + let deadline = Instant::now() + Duration::from_secs(5); + while !dropped.load(Ordering::SeqCst) { + assert!(Instant::now() < deadline, "the producer kept running"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[test] + fn run_reports_a_panic_as_an_error() { + async fn faulty() -> anyhow::Result<()> { + panic!("planner bug") + } + + let (api, executor, _query) = two_runtimes(); + let error = api.block_on(executor.run(faulty())).unwrap_err(); + assert!(error.to_string().contains("query task failed"), "{error}"); + } + + #[test] + fn run_returns_the_value() { + let (api, executor, _query) = two_runtimes(); + let value = api + .block_on(executor.run(async { Ok::<_, anyhow::Error>(41 + 1) })) + .unwrap(); + assert_eq!(value, 42); + } +} diff --git a/beacon-db/beacon-core/src/runtime.rs b/beacon-db/beacon-core/src/runtime.rs index 7c11d6ee..b1d00e07 100644 --- a/beacon-db/beacon-core/src/runtime.rs +++ b/beacon-db/beacon-core/src/runtime.rs @@ -10,6 +10,7 @@ use futures::TryStreamExt; use crate::{ parser::{beacon_parser::BeaconParser, statement::BeaconStatement}, + query_executor::QueryExecutor, query_metrics_store::QueryMetricsStore, query_result::{ArrowOutputStream, QueryOutput, QueryResult}, }; @@ -29,7 +30,15 @@ use crate::{ /// ([`Self::get_query_metrics`]). Those answer in Arrow and in beacon's own /// types; turning them into a wire format is the transport's job, not this /// one's. +/// +/// Every field is shared, so a clone is cheap. [`Self::run_query`] clones the +/// runtime into the task it spawns on the query executor. +#[derive(Clone)] pub struct Runtime { + /// The runtime queries run on. The server gives it a runtime of its own, so + /// a long scan never holds an API worker. The embedded database gives it the + /// one runtime it has. + pub(crate) executor: QueryExecutor, pub(crate) session_ctx: Arc, /// Where a completed query's metrics are written: the managed table that /// `beacon.system.query_metrics` exposes, which is how callers read them. @@ -95,6 +104,19 @@ impl Runtime { &self, query: crate::query::Query, identity: beacon_auth::AuthIdentity, + ) -> anyhow::Result { + // Planning reads metadata too, so the whole query moves to the executor. + let runtime = self.clone(); + self.executor + .run(async move { runtime.run_query_on_executor(query, identity).await }) + .await + } + + /// [`Self::run_query`], on the query executor. + async fn run_query_on_executor( + &self, + query: crate::query::Query, + identity: beacon_auth::AuthIdentity, ) -> anyhow::Result { let query_id = uuid::Uuid::new_v4(); let query_json = serde_json::to_value(&query)?; @@ -179,6 +201,9 @@ impl Runtime { if let Some(physical_plan) = physical_plan { metrics.set_physical_plan(physical_plan); } + // The caller drains the result on its own runtime. The bridge keeps the + // scan, and the tasks DataFusion spawns on first poll, on the executor. + let stream = self.executor.bridge_stream(stream); let output_stream = ArrowOutputStream::new(stream, metrics, self.query_metrics.clone()); Ok(QueryResult { query_output: QueryOutput::Stream(output_stream), @@ -445,6 +470,18 @@ impl Runtime { &self, query: crate::query::Query, identity: beacon_auth::AuthIdentity, + ) -> anyhow::Result { + let runtime = self.clone(); + self.executor + .run(async move { runtime.explain_query_on_executor(query, identity).await }) + .await + } + + /// [`Self::explain_query`], on the query executor. + async fn explain_query_on_executor( + &self, + query: crate::query::Query, + identity: beacon_auth::AuthIdentity, ) -> anyhow::Result { // `output` (file format) is meaningless here: only the query body is planned. let plan = self.lower_query(query.inner).await?; @@ -476,6 +513,22 @@ impl Runtime { &self, query: crate::query::Query, identity: beacon_auth::AuthIdentity, + ) -> anyhow::Result { + let runtime = self.clone(); + self.executor + .run(async move { + runtime + .explain_analyze_query_on_executor(query, identity) + .await + }) + .await + } + + /// [`Self::explain_analyze_query`], on the query executor. + async fn explain_analyze_query_on_executor( + &self, + query: crate::query::Query, + identity: beacon_auth::AuthIdentity, ) -> anyhow::Result { let plan = self.lower_query(query.inner).await?; crate::statement_plan::validate_query_plan(&plan, identity.is_super_user)?; diff --git a/beacon-db/beacon-core/src/runtime_builder.rs b/beacon-db/beacon-core/src/runtime_builder.rs index d3e10f38..5e5aef86 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -53,6 +53,7 @@ use tokio::runtime::Handle; use crate::{ auth_store::TablesAuthStore, + query_executor::QueryExecutor, runtime::Runtime, settings::{SqlSettings, SqlStreamCoalesceSettings}, statement_plan::{new_session_cell, BeaconQueryPlanner, CoalesceSqlStream, SessionCell}, @@ -354,7 +355,7 @@ impl RuntimeBuilder { // only a hydrated catalog lets their `CREATE TABLE IF NOT EXISTS` see them — // otherwise the managed-table create path re-creates them and Lance rejects // the duplicate dataset. - register_schema_provider(&self, &session_ctx).await?; + register_schema_provider(&runtime_handle, &session_ctx).await?; // The session and its tables now exist and the session cell is filled, so the auth store can // reach its tables. Ensure they exist, hydrate the in-memory user/role copies from whatever @@ -390,6 +391,7 @@ impl RuntimeBuilder { crate::query::temp_object::sweep_stale_outputs(&tmp_dir); Ok(Runtime { + executor: QueryExecutor::new(runtime_handle), session_ctx, query_metrics, auth: auth_context, @@ -676,11 +678,13 @@ async fn init_session_ctx( } async fn register_schema_provider( - _runtime_builder: &RuntimeBuilder, + runtime_handle: &Handle, session_ctx: &Arc, ) -> anyhow::Result<()> { + // The provider blocks on this handle from inside DataFusion's sync API. It + // must be the query runtime, not the runtime `build` happens to run on. let schema_provider = Arc::new(PersistentSchemaProvider::new( - tokio::runtime::Handle::current(), + runtime_handle.clone(), session_ctx.clone(), DEFAULT_DB_STORE_URL_OBJECT_URL.clone(), )); diff --git a/beacon-server/beacon-server-config/src/error.rs b/beacon-server/beacon-server-config/src/error.rs index 15472328..81bf6168 100644 --- a/beacon-server/beacon-server-config/src/error.rs +++ b/beacon-server/beacon-server-config/src/error.rs @@ -29,6 +29,10 @@ pub enum ConfigError { #[error("invalid BEACON_LOG_LEVEL: {0}")] InvalidLogLevel(String), + /// `BEACON_WORKER_THREADS` or `BEACON_API_THREADS` was set to zero. + #[error("invalid thread count: {0}")] + InvalidThreads(String), + /// `BEACON_SECRETS_KEY` was set but is not a base64-encoded 32-byte key. #[error("invalid BEACON_SECRETS_KEY: {0}")] InvalidSecretsKey(String), diff --git a/beacon-server/beacon-server-config/src/lib.rs b/beacon-server/beacon-server-config/src/lib.rs index 2f8f248f..8a183936 100644 --- a/beacon-server/beacon-server-config/src/lib.rs +++ b/beacon-server/beacon-server-config/src/lib.rs @@ -80,7 +80,11 @@ pub struct OidcConfig { pub struct ServerConfig { pub port: u16, pub host: String, + /// Threads of the query runtime. From `BEACON_WORKER_THREADS`. pub worker_threads: usize, + /// Threads of the runtime that serves HTTP and Flight SQL. From + /// `BEACON_API_THREADS`. Its own runtime, so a long query does not block it. + pub api_threads: usize, /// URL prefix for all HTTP routes, e.g. `/base-path`. Empty string means serve at `/`. pub base_path: String, /// Directory holding the built admin web UI (Vite `dist/`). Served at @@ -348,6 +352,8 @@ struct RawConfig { sql_stream_coalesce_max_rows: usize, #[envconfig(from = "BEACON_WORKER_THREADS", default = "8")] worker_threads: usize, + #[envconfig(from = "BEACON_API_THREADS", default = "4")] + api_threads: usize, #[envconfig(from = "BEACON_BASE_PATH", default = "")] base_path: String, /// Directory containing the built admin web UI. Defaults to `web` (resolved @@ -614,6 +620,7 @@ impl From for Config { port: raw.port, host: raw.host, worker_threads: raw.worker_threads, + api_threads: raw.api_threads, base_path: raw.base_path, web_ui_dir: raw.web_ui_dir, max_upload_bytes: raw.max_upload_bytes, @@ -756,6 +763,24 @@ fn validate_storage(s3: &S3Config) -> Result<()> { Ok(()) } +/// Rejects a runtime with no threads. +/// +/// Tokio panics on a zero thread count. A clean error at startup names the +/// variable instead. +fn validate_threads(server: &ServerConfig) -> Result<()> { + for (name, count) in [ + ("BEACON_WORKER_THREADS", server.worker_threads), + ("BEACON_API_THREADS", server.api_threads), + ] { + if count == 0 { + return Err(ConfigError::InvalidThreads(format!( + "{name} must be at least 1" + ))); + } + } + Ok(()) +} + /// Levels accepted by `BEACON_LOG_LEVEL`, in the spelling `tracing` expects. const LOG_LEVELS: [&str; 6] = ["trace", "debug", "info", "warn", "error", "off"]; @@ -841,6 +866,7 @@ impl Config { config.server.log_level = normalize_log_level(&config.server.log_level).map_err(ConfigError::InvalidLogLevel)?; + validate_threads(&config.server)?; validate_storage(&config.s3)?; // Create the configured data directories (idempotent). `db_file` is a file, @@ -932,8 +958,8 @@ fn create_dir(path: &Path) -> Result<()> { #[cfg(test)] mod tests { use super::{ - decode_master_key, normalize_base_path, normalize_log_level, validate_storage, Config, - Hdf5Convention, PathBuf, RawConfig, + decode_master_key, normalize_base_path, normalize_log_level, validate_storage, + validate_threads, Config, Hdf5Convention, PathBuf, RawConfig, }; use envconfig::Envconfig; use std::collections::HashMap; @@ -956,6 +982,28 @@ mod tests { Config::from(raw(vars).expect("config should parse")) } + /// Both runtimes have a default size, and neither accepts zero threads: + /// Tokio would panic where the config can name the variable instead. + #[test] + fn thread_counts_default_and_reject_zero() { + let defaults = config(&[]); + assert_eq!(defaults.server.worker_threads, 8); + assert_eq!(defaults.server.api_threads, 4); + assert!(validate_threads(&defaults.server).is_ok()); + + let no_query_threads = config(&[("BEACON_WORKER_THREADS", "0")]); + let error = validate_threads(&no_query_threads.server) + .unwrap_err() + .to_string(); + assert!(error.contains("BEACON_WORKER_THREADS"), "{error}"); + + let no_api_threads = config(&[("BEACON_API_THREADS", "0")]); + let error = validate_threads(&no_api_threads.server) + .unwrap_err() + .to_string(); + assert!(error.contains("BEACON_API_THREADS"), "{error}"); + } + /// Every data path derives from `BEACON_DATA_DIR`. This is a regression guard: /// the paths used to be `lazy_static`s pinned to `./data`, so setting the /// variable relocated only `indexes` and `cache` while the datasets, tables, diff --git a/beacon-server/beacon-server/src/main.rs b/beacon-server/beacon-server/src/main.rs index 88a36638..2b7b5efb 100644 --- a/beacon-server/beacon-server/src/main.rs +++ b/beacon-server/beacon-server/src/main.rs @@ -7,7 +7,7 @@ use std::{net::IpAddr, str::FromStr, sync::Arc}; use anyhow::Context; -use tokio::runtime::Builder; +use tokio::runtime::{Builder, Handle}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use beacon_server::{axum::setup_router, flight_sql, Server}; @@ -21,7 +21,12 @@ static GLOBAL: Jemalloc = Jemalloc; const BEACON_VERSION: &str = env!("CARGO_PKG_VERSION"); -/// Builds the Tokio runtime and hands control to the async entrypoint. +/// Builds the two Tokio runtimes and hands control to the async entrypoint. +/// +/// The API runtime serves HTTP and Flight SQL. The query runtime plans and runs +/// queries, and hosts the crawler and file statistics timers. A partition decode +/// holds a worker until it yields, and Tokio cannot preempt it, so a query on +/// the API runtime would make every other request wait. fn main() -> anyhow::Result<()> { // Load and validate configuration up front so problems (e.g. a malformed // `BEACON_BASE_PATH`) surface as a clean error here. The config is owned and @@ -29,25 +34,51 @@ fn main() -> anyhow::Result<()> { // a process-global. let config = Arc::new(beacon_server_config::Config::load().context("failed to load configuration")?); - let rt = Builder::new_multi_thread() + let api_runtime = Builder::new_multi_thread() + .worker_threads(config.server.api_threads) + .thread_name("beacon-api") + .enable_all() + .build() + .context("failed to build the API Tokio runtime")?; + let query_runtime = Builder::new_multi_thread() .worker_threads(config.server.worker_threads) + .thread_name("beacon-query") .enable_all() .build() - .context("failed to build Tokio runtime")?; + .context("failed to build the query Tokio runtime")?; - rt.block_on(async_main(config)) + // Both runtimes live until this returns, so a task on one can always reach the other. + api_runtime.block_on(async_main(config, query_runtime.handle().clone())) } /// Initializes shared services and starts all configured API transports. -async fn async_main(config: Arc) -> anyhow::Result<()> { +async fn async_main( + config: Arc, + query_runtime: Handle, +) -> anyhow::Result<()> { let log_filter = setup_tracing(&config); install_panic_hook(); tracing::info!("Beacon v{}", BEACON_VERSION); // This line only prints when DEBUG is on, so it confirms the level took effect. tracing::debug!(filter = %log_filter, "debug logging is on"); + tracing::info!( + api_threads = config.server.api_threads, + query_threads = config.server.worker_threads, + "runtime threads" + ); // The server owns the datasets store and hosts the runtime that queries it. - let server = Arc::new(Server::open(config.clone()).await?); + // It opens on the query runtime: the timers it starts spawn onto the ambient + // runtime, and they belong with the queries, not with the API. + let server = { + let config = config.clone(); + let handle = query_runtime.clone(); + query_runtime + .spawn(async move { Server::open(config, handle).await }) + .await + .context("the server did not finish opening")?? + }; + let server = Arc::new(server); // Keep both transports on the same server so metadata and access rules stay aligned. let router = setup_router(server.clone(), config.clone())?; diff --git a/beacon-server/beacon-server/src/server/mod.rs b/beacon-server/beacon-server/src/server/mod.rs index e9fa8655..48d053c7 100644 --- a/beacon-server/beacon-server/src/server/mod.rs +++ b/beacon-server/beacon-server/src/server/mod.rs @@ -30,6 +30,7 @@ use beacon_datafusion_ext::listing_factory::RootStore; use beacon_server_config::Config; use datafusion::execution::object_store::ObjectStoreUrl; use object_store::ObjectStore; +use tokio::runtime::Handle; /// The URL scheme bare dataset paths resolve against inside the runtime. pub const DATASETS_STORE_URL: &str = "datasets://"; @@ -52,15 +53,19 @@ impl Server { /// Open the server described by `config`: build the datasets store, then /// start a runtime over it. /// + /// `query_runtime` is the Tokio runtime queries run on. The binary gives it a + /// runtime of its own, apart from the one that serves the API, so a long scan + /// never holds an API worker. A test passes the runtime it runs on. + /// /// A server is always persistent — the tables store is the single redb file at /// `config.data.db_file`. To get throwaway state, point `config` at a /// temporary directory; there is no separate in-memory mode. - pub async fn open(config: Arc) -> anyhow::Result { + pub async fn open(config: Arc, query_runtime: Handle) -> anyhow::Result { let (store, root) = build_datasets_store(&config)?; let store_url = ObjectStoreUrl::parse(DATASETS_STORE_URL).context("invalid datasets store URL")?; - let runtime = build_runtime(&config, store_url, root, store.clone()) + let runtime = build_runtime(&config, store_url, root, store.clone(), query_runtime) .await .context("failed to start the beacon runtime")?; @@ -258,9 +263,10 @@ async fn build_runtime( store_url: ObjectStoreUrl, root: RootStore, store: Arc, + query_runtime: Handle, ) -> anyhow::Result { let mut builder = RuntimeBuilder::new() - .with_runtime_handle(tokio::runtime::Handle::current()) + .with_runtime_handle(query_runtime) // The store the server owns; the root is what native readers (netCDF-c) // translate object paths against. .with_default_store(store_url, root) diff --git a/beacon-server/beacon-server/tests/common/mod.rs b/beacon-server/beacon-server/tests/common/mod.rs index d823b595..d9559eb7 100644 --- a/beacon-server/beacon-server/tests/common/mod.rs +++ b/beacon-server/beacon-server/tests/common/mod.rs @@ -61,7 +61,8 @@ pub async fn server_with(mut config: beacon_server_config::Config) -> TestServer std::fs::create_dir_all(dir).expect("create temp data dir"); } - let server = Server::open(Arc::new(config)) + // Tests run queries on the runtime they run on; only the binary has two. + let server = Server::open(Arc::new(config), tokio::runtime::Handle::current()) .await .expect("server should open"); TestServer { diff --git a/docs/docs/2.0.0-rc5/server/configuration.md b/docs/docs/2.0.0-rc5/server/configuration.md index 1320cb9f..eca7c555 100644 --- a/docs/docs/2.0.0-rc5/server/configuration.md +++ b/docs/docs/2.0.0-rc5/server/configuration.md @@ -20,7 +20,8 @@ See [S3 object storage](#s3-object-storage). | --- | --- | --- | | `BEACON_HOST` | `0.0.0.0` | IP address the HTTP API listens on. | | `BEACON_PORT` | `5001` | Port the HTTP API listens on. | -| `BEACON_WORKER_THREADS` | `8` | Number of worker threads for the async runtime. | +| `BEACON_WORKER_THREADS` | `8` | Number of threads of the query runtime. This runtime plans and runs queries, and hosts the crawlers and the file statistics. | +| `BEACON_API_THREADS` | `4` | Number of threads of the API runtime. This runtime serves HTTP and Flight SQL. A long query does not block it. | | `BEACON_LOG_LEVEL` | `info` | Log level: `trace`, `debug`, `info`, `warn`, `error`, or `off`. Case does not matter. The level applies to all Beacon crates. At `debug` and `trace`, loud dependencies such as DataFusion, Arrow, `object_store`, and hyper stay at `info`. An unknown value stops the server at startup. | | `RUST_LOG` | _(unset)_ | Full log filter, in [`tracing-subscriber` EnvFilter](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html) syntax (e.g. `debug,datafusion=trace`). It replaces `BEACON_LOG_LEVEL`. Use it to see the dependency logs that `BEACON_LOG_LEVEL` holds back. An invalid value prints a warning, and Beacon uses `BEACON_LOG_LEVEL`. | | `BEACON_BASE_PATH` | _(empty)_ | Optional URL path prefix for the HTTP API, OpenAPI document, and Swagger UI (e.g. `/beacon`). Useful behind a reverse proxy. Normalized to exactly one leading slash and no trailing slash, so `beacon`, `/beacon`, and `/beacon/` are equivalent. Only URL-safe characters are allowed (letters, digits, `-`, `_`, `.`, `~`, and `/` as a separator); any other character causes Beacon to exit at startup with a descriptive error. | diff --git a/docs/docs/2.0.0-rc5/server/performance-tuning.md b/docs/docs/2.0.0-rc5/server/performance-tuning.md index 436d1f3b..558e280d 100644 --- a/docs/docs/2.0.0-rc5/server/performance-tuning.md +++ b/docs/docs/2.0.0-rc5/server/performance-tuning.md @@ -16,9 +16,13 @@ Every setting below is an environment variable. [configuration.md](configuration ### CPU and concurrency +Beacon runs two Tokio runtimes. The API runtime serves HTTP and Flight SQL. The query runtime +plans and runs queries. A long scan holds a query thread, but the API runtime keeps its own +threads, so a login, a health check, or the admin UI answers at once. + #### `BEACON_WORKER_THREADS` -This value sizes the Tokio runtime of Beacon. That runtime runs the API requests and the query work. +This value sizes the query runtime. The crawlers and the file statistics run there too. - On a dedicated machine, set `BEACON_WORKER_THREADS` to the number of physical cores. - On a shared host, set a lower value. Other services then keep enough CPU. @@ -26,6 +30,12 @@ This value sizes the Tokio runtime of Beacon. That runtime runs the API requests More threads help an I/O-heavy workload, such as a read from object storage or a NetCDF read over HTTP. A CPU-heavy workload, such as an aggregate or a join, does not scale past the CPU count. +#### `BEACON_API_THREADS` + +This value sizes the API runtime. The default of `4` is enough for most deployments. The API +runtime encodes and compresses the result stream for each client. A deployment with many +concurrent large downloads can raise the value. + ### Memory and disk spilling #### `BEACON_VM_MEMORY_SIZE`