engine: Service timeout shutdown (#1919)

* timeout shutdown

* timeout shutdown, clippy & fmt

* timeout shutdown proper name a default value

* timeout shutdown interval && idle timeout

Author: DanCodedThis

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<Arc<UserSession>>;
@@ -145,6 +146,8 @@ pub trait ExecutionService: Send + Sync {
         file_name: &str,
         format: Format,
     ) -> Result<usize>;
+
+    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<std::time::Instant> = 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

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<u64>,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]

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<Db>) {
+async fn shutdown_signal(db: Arc<Db>, execution_svc: Arc<dyn ExecutionService>, timeout: u64) {
     let ctrl_c = async {
         signal::ctrl_c()
             .await
@@ -487,6 +494,11 @@ async fn shutdown_signal(db: Arc<Db>) {
     #[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>) {
             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<openapi::OpenApi> {
     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");
+            }
+        }
+    }
+}