diff --git a/crates/core-executor/src/service.rs b/crates/core-executor/src/service.rs index 487228c66..dc994a812 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 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}; @@ -44,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,6 +146,8 @@ pub trait ExecutionService: Send + Sync { file_name: &str, format: Format, ) -> Result; + + async fn timeout_signal(&self, interval: Duration, idle_timeout: Duration) -> (); } pub struct CoreExecutionService { @@ -891,6 +894,37 @@ impl ExecutionService for CoreExecutionService { Ok(rows_loaded) } + + 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_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; + } + } + } + } } //Test environment diff --git a/crates/embucketd/src/cli.rs b/crates/embucketd/src/cli.rs index 9db5256cb..f42d22d5e 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 = "IDLE_TIMEOUT_SECONDS", + default_value = "18000", + help = "Service idle 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..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; +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; @@ -345,8 +347,13 @@ 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 +476,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 +494,11 @@ 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_SIGNAL_INTERVAL_SECONDS), + tokio::time::Duration::from_secs(timeout), + ); + tokio::select! { () = ctrl_c => { db.close().await.expect("Failed to close database"); @@ -496,6 +508,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 +528,54 @@ 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; + + #[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), Duration::from_secs(3)); + tokio::select! { + () = timeout => { + tracing::warn!("No sessions in use & no running queries - timeout, starting graceful shutdown"); + } + } + } +}