From f1f9de0b671c78c3446c60d47501a588de085d0c Mon Sep 17 00:00:00 2001 From: DanCodedThis Date: Tue, 11 Nov 2025 19:43:17 +0200 Subject: [PATCH 1/4] timeout shutdown --- crates/core-executor/src/service.rs | 17 +++++++- crates/embucketd/src/cli.rs | 8 ++++ crates/embucketd/src/main.rs | 64 +++++++++++++++++++++++++++-- 3 files changed, 85 insertions(+), 4 deletions(-) diff --git a/crates/core-executor/src/service.rs b/crates/core-executor/src/service.rs index 487228c66..08af8a015 100644 --- a/crates/core-executor/src/service.rs +++ b/crates/core-executor/src/service.rs @@ -19,7 +19,7 @@ use std::sync::atomic::Ordering; use std::vec; use std::{collections::HashMap, sync::Arc}; use time::{Duration as DateTimeDuration, OffsetDateTime}; - +use tokio::io::AsyncWriteExt; use super::error::{self as ex_error, Result}; use super::models::{AsyncQueryHandle, QueryContext, QueryResult, QueryResultStatus}; use super::running_queries::{RunningQueries, RunningQueriesRegistry, RunningQuery}; @@ -145,6 +145,8 @@ pub trait ExecutionService: Send + Sync { file_name: &str, format: Format, ) -> Result; + + async fn timeout_signal(&self, period: Duration) -> (); } pub struct CoreExecutionService { @@ -891,6 +893,19 @@ impl ExecutionService for CoreExecutionService { Ok(rows_loaded) } + + async fn timeout_signal(&self, period: Duration) -> () { + let mut interval = tokio::time::interval(period); + interval.tick().await; // The first tick completes immediately; skip. + loop { + interval.tick().await; + let sessions_in_use = self.df_sessions.read().await.len() > 0; + let queries_in_use = self.queries.count() > 0; + if !sessions_in_use && !queries_in_use { + return (); + } + } + } } //Test environment diff --git a/crates/embucketd/src/cli.rs b/crates/embucketd/src/cli.rs index 9db5256cb..dd1358a44 100644 --- a/crates/embucketd/src/cli.rs +++ b/crates/embucketd/src/cli.rs @@ -284,6 +284,14 @@ pub struct CliOpts { help = "Tracing span processor" )] pub tracing_span_processor: TracingSpanProcessor, + + #[arg( + long, + env = "TIMEOUT", + default_value = "300", + help = "Service timeout in seconds" + )] + pub timeout: Option, } #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] diff --git a/crates/embucketd/src/main.rs b/crates/embucketd/src/main.rs index 1e3a3b420..e7c80f8a8 100644 --- a/crates/embucketd/src/main.rs +++ b/crates/embucketd/src/main.rs @@ -34,7 +34,7 @@ use axum::{ routing::{get, post}, }; use clap::Parser; -use core_executor::service::CoreExecutionService; +use core_executor::service::{CoreExecutionService, ExecutionService}; use core_executor::utils::Config as ExecutionConfig; use core_history::SlateDBHistoryStore; use core_metastore::SlateDBMetastore; @@ -345,8 +345,9 @@ async fn async_main( .expect("Failed to bind to address"); let addr = listener.local_addr().expect("Failed to get local address"); tracing::info!(%addr, "Listening on http"); + let timeout = opts.timeout.unwrap(); axum::serve(listener, router) - .with_graceful_shutdown(shutdown_signal(Arc::new(db.clone()))) + .with_graceful_shutdown(shutdown_signal(Arc::new(db.clone()), execution_svc.clone(), timeout)) .await .expect("Failed to start server"); @@ -469,7 +470,7 @@ fn setup_tracing(opts: &cli::CliOpts) -> SdkTracerProvider { clippy::redundant_pub_crate, clippy::cognitive_complexity )] -async fn shutdown_signal(db: Arc) { +async fn shutdown_signal(db: Arc, execution_svc: Arc, timeout: u64) { let ctrl_c = async { signal::ctrl_c() .await @@ -487,6 +488,8 @@ async fn shutdown_signal(db: Arc) { #[cfg(not(unix))] let terminate = std::future::pending::<()>(); + let timeout = execution_svc.timeout_signal(tokio::time::Duration::from_secs(timeout)); + tokio::select! { () = ctrl_c => { db.close().await.expect("Failed to close database"); @@ -496,6 +499,9 @@ async fn shutdown_signal(db: Arc) { db.close().await.expect("Failed to close database"); tracing::warn!("SIGTERM received, starting graceful shutdown"); }, + () = timeout => { + tracing::warn!("No sessions in use & no running queries - timeout, starting graceful shutdown"); + } } tracing::warn!("signal received, starting graceful shutdown"); @@ -513,3 +519,55 @@ fn load_openapi_spec() -> Option { original_spec.paths = openapi::Paths::new(); Some(original_spec) } + +#[cfg(test)] +mod tests { + use api_sessions::session::SessionStore; + use core_executor::models::QueryContext; + use core_executor::service::ExecutionService; + use core_executor::service::make_test_execution_svc; + use core_executor::session::to_unix; + use std::sync::atomic::Ordering; + use std::time::Duration; + use time::OffsetDateTime; + use tokio::time::sleep; + + #[tokio::test] + #[allow(clippy::expect_used, clippy::too_many_lines)] + async fn test_timeout_signal() { + let execution_svc = make_test_execution_svc().await; + + let df_session_id = "fasfsafsfasafsass".to_string(); + let user_session = execution_svc + .create_session(&df_session_id) + .await + .expect("Failed to create a session"); + + execution_svc + .query(&df_session_id, "SELECT SLEEP(5)", QueryContext::default()) + .await + .expect("Failed to execute query (session deleted)"); + + user_session + .expiry + .store(to_unix(OffsetDateTime::now_utc()), Ordering::Relaxed); + + let session_store = SessionStore::new(execution_svc.clone()); + + tokio::task::spawn({ + let session_store = session_store.clone(); + async move { + session_store + .continuously_delete_expired(Duration::from_secs(1)) + .await; + } + }); + + let timeout = execution_svc.timeout_signal(Duration::from_secs(1)); + tokio::select! { + () = timeout => { + tracing::warn!("No sessions in use & no running queries - timeout, starting graceful shutdown"); + } + } + } +} From 8969812b645e2f238d1977ecc6781cffdaa70d83 Mon Sep 17 00:00:00 2001 From: DanCodedThis Date: Tue, 11 Nov 2025 19:51:10 +0200 Subject: [PATCH 2/4] timeout shutdown, clippy & fmt --- crates/core-executor/src/service.rs | 37 ++++++++++++++--------------- crates/embucketd/src/main.rs | 7 ++++-- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/crates/core-executor/src/service.rs b/crates/core-executor/src/service.rs index 08af8a015..a04740dce 100644 --- a/crates/core-executor/src/service.rs +++ b/crates/core-executor/src/service.rs @@ -1,4 +1,19 @@ +use super::error::{self as ex_error, Result}; +use super::models::{AsyncQueryHandle, QueryContext, QueryResult, QueryResultStatus}; +use super::running_queries::{RunningQueries, RunningQueriesRegistry, RunningQuery}; +use super::session::UserSession; +use crate::running_queries::RunningQueryId; +use crate::session::{SESSION_INACTIVITY_EXPIRATION_SECONDS, to_unix}; +use crate::tracing::SpanTracer; +use crate::utils::{Config, MemPoolType}; use bytes::{Buf, Bytes}; +use core_history::HistoryStore; +use core_history::SlateDBHistoryStore; +use core_history::{QueryRecordId, QueryResultError, QueryStatus}; +use core_metastore::{ + Database, Metastore, Schema, SchemaIdent, SlateDBMetastore, TableIdent as MetastoreTableIdent, + Volume, VolumeType, +}; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::csv::ReaderBuilder; use datafusion::arrow::csv::reader::Format; @@ -13,29 +28,13 @@ use datafusion::execution::memory_pool::{ }; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion_common::TableReference; +use df_catalog::catalog_list::{DEFAULT_CATALOG, EmbucketCatalogList}; use snafu::ResultExt; use std::num::NonZeroUsize; use std::sync::atomic::Ordering; use std::vec; use std::{collections::HashMap, sync::Arc}; use time::{Duration as DateTimeDuration, OffsetDateTime}; -use tokio::io::AsyncWriteExt; -use super::error::{self as ex_error, Result}; -use super::models::{AsyncQueryHandle, QueryContext, QueryResult, QueryResultStatus}; -use super::running_queries::{RunningQueries, RunningQueriesRegistry, RunningQuery}; -use super::session::UserSession; -use crate::running_queries::RunningQueryId; -use crate::session::{SESSION_INACTIVITY_EXPIRATION_SECONDS, to_unix}; -use crate::tracing::SpanTracer; -use crate::utils::{Config, MemPoolType}; -use core_history::HistoryStore; -use core_history::SlateDBHistoryStore; -use core_history::{QueryRecordId, QueryResultError, QueryStatus}; -use core_metastore::{ - Database, Metastore, Schema, SchemaIdent, SlateDBMetastore, TableIdent as MetastoreTableIdent, - Volume, VolumeType, -}; -use df_catalog::catalog_list::{DEFAULT_CATALOG, EmbucketCatalogList}; use tokio::sync::RwLock; use tokio::sync::oneshot; use tokio::time::{Duration, timeout}; @@ -899,10 +898,10 @@ impl ExecutionService for CoreExecutionService { interval.tick().await; // The first tick completes immediately; skip. loop { interval.tick().await; - let sessions_in_use = self.df_sessions.read().await.len() > 0; + let sessions_in_use = !self.df_sessions.read().await.is_empty(); let queries_in_use = self.queries.count() > 0; if !sessions_in_use && !queries_in_use { - return (); + return; } } } diff --git a/crates/embucketd/src/main.rs b/crates/embucketd/src/main.rs index e7c80f8a8..b5711e8e7 100644 --- a/crates/embucketd/src/main.rs +++ b/crates/embucketd/src/main.rs @@ -347,7 +347,11 @@ async fn async_main( tracing::info!(%addr, "Listening on http"); let timeout = opts.timeout.unwrap(); axum::serve(listener, router) - .with_graceful_shutdown(shutdown_signal(Arc::new(db.clone()), execution_svc.clone(), timeout)) + .with_graceful_shutdown(shutdown_signal( + Arc::new(db.clone()), + execution_svc.clone(), + timeout, + )) .await .expect("Failed to start server"); @@ -530,7 +534,6 @@ mod tests { use std::sync::atomic::Ordering; use std::time::Duration; use time::OffsetDateTime; - use tokio::time::sleep; #[tokio::test] #[allow(clippy::expect_used, clippy::too_many_lines)] From 4a2348daeab908ac0f43163bb5e9709c25876540 Mon Sep 17 00:00:00 2001 From: DanCodedThis Date: Tue, 11 Nov 2025 20:05:41 +0200 Subject: [PATCH 3/4] timeout shutdown proper name a default value --- crates/embucketd/src/cli.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/embucketd/src/cli.rs b/crates/embucketd/src/cli.rs index dd1358a44..f42d22d5e 100644 --- a/crates/embucketd/src/cli.rs +++ b/crates/embucketd/src/cli.rs @@ -287,9 +287,9 @@ pub struct CliOpts { #[arg( long, - env = "TIMEOUT", - default_value = "300", - help = "Service timeout in seconds" + env = "IDLE_TIMEOUT_SECONDS", + default_value = "18000", + help = "Service idle timeout in seconds" )] pub timeout: Option, } From b6e78aaa3ed9ca26d0c92ac0c2b368f5e3fa89b7 Mon Sep 17 00:00:00 2001 From: DanCodedThis Date: Wed, 12 Nov 2025 15:49:54 +0200 Subject: [PATCH 4/4] timeout shutdown interval && idle timeout --- crates/core-executor/src/service.rs | 34 +++++++++++++++++++++++------ crates/embucketd/src/main.rs | 11 +++++++--- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/crates/core-executor/src/service.rs b/crates/core-executor/src/service.rs index a04740dce..dc994a812 100644 --- a/crates/core-executor/src/service.rs +++ b/crates/core-executor/src/service.rs @@ -43,6 +43,8 @@ use uuid::Uuid; const DEFAULT_SCHEMA: &str = "public"; +pub const TIMEOUT_SIGNAL_INTERVAL_SECONDS: u64 = 60; + #[async_trait::async_trait] pub trait ExecutionService: Send + Sync { async fn create_session(&self, session_id: &str) -> Result>; @@ -145,7 +147,7 @@ pub trait ExecutionService: Send + Sync { format: Format, ) -> Result; - async fn timeout_signal(&self, period: Duration) -> (); + async fn timeout_signal(&self, interval: Duration, idle_timeout: Duration) -> (); } pub struct CoreExecutionService { @@ -893,15 +895,33 @@ impl ExecutionService for CoreExecutionService { Ok(rows_loaded) } - async fn timeout_signal(&self, period: Duration) -> () { - let mut interval = tokio::time::interval(period); + async fn timeout_signal(&self, interval: Duration, idle_timeout: Duration) -> () { + let mut interval = tokio::time::interval(interval); interval.tick().await; // The first tick completes immediately; skip. + let mut idle_since: Option = None; loop { interval.tick().await; - let sessions_in_use = !self.df_sessions.read().await.is_empty(); - let queries_in_use = self.queries.count() > 0; - if !sessions_in_use && !queries_in_use { - return; + let sessions_empty = { + let sessions = self.df_sessions.read().await; + sessions.is_empty() + }; + let queries_empty = self.queries.count() == 0; + let idle_now = sessions_empty && queries_empty; + match (idle_now, idle_since) { + (true, None) => { + // just entered idle + idle_since = Some(std::time::Instant::now()); + } + (true, Some(since)) => { + if since.elapsed() >= idle_timeout { + // stayed idle long enough + return; + } + } + (false, _) => { + // became active again, reset the idle window + idle_since = None; + } } } } diff --git a/crates/embucketd/src/main.rs b/crates/embucketd/src/main.rs index b5711e8e7..5492f69db 100644 --- a/crates/embucketd/src/main.rs +++ b/crates/embucketd/src/main.rs @@ -34,7 +34,9 @@ use axum::{ routing::{get, post}, }; use clap::Parser; -use core_executor::service::{CoreExecutionService, ExecutionService}; +use core_executor::service::{ + CoreExecutionService, ExecutionService, TIMEOUT_SIGNAL_INTERVAL_SECONDS, +}; use core_executor::utils::Config as ExecutionConfig; use core_history::SlateDBHistoryStore; use core_metastore::SlateDBMetastore; @@ -492,7 +494,10 @@ async fn shutdown_signal(db: Arc, execution_svc: Arc, #[cfg(not(unix))] let terminate = std::future::pending::<()>(); - let timeout = execution_svc.timeout_signal(tokio::time::Duration::from_secs(timeout)); + let timeout = execution_svc.timeout_signal( + tokio::time::Duration::from_secs(TIMEOUT_SIGNAL_INTERVAL_SECONDS), + tokio::time::Duration::from_secs(timeout), + ); tokio::select! { () = ctrl_c => { @@ -566,7 +571,7 @@ mod tests { } }); - let timeout = execution_svc.timeout_signal(Duration::from_secs(1)); + let timeout = execution_svc.timeout_signal(Duration::from_secs(1), Duration::from_secs(3)); tokio::select! { () = timeout => { tracing::warn!("No sessions in use & no running queries - timeout, starting graceful shutdown");