engine: Timeout signal (#10)

* fix sessions

* fix sessions fmt

* timeout signal

Author: DanCodedThis

crates/embucketd/src/cli.rs

@@ -152,6 +152,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(Debug, Clone, ValueEnum)]

crates/embucketd/src/main.rs

@@ -21,7 +21,7 @@ use axum::{
 use catalog_metastore::InMemoryMetastore;
 use clap::Parser;
 use dotenv::dotenv;
-use executor::service::CoreExecutionService;
+use executor::service::{CoreExecutionService, ExecutionService, TIMEOUT_SIGNAL_INTERVAL_SECONDS};
 use executor::utils::Config as ExecutionConfig;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry_sdk::Resource;
@@ -218,8 +218,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())
+        .with_graceful_shutdown(shutdown_signal(execution_svc.clone(), timeout))
         .await
         .expect("Failed to start server");
 
@@ -337,7 +338,7 @@ fn setup_tracing(opts: &cli::CliOpts) -> SdkTracerProvider {
     clippy::redundant_pub_crate,
     clippy::cognitive_complexity
 )]
-async fn shutdown_signal() {
+async fn shutdown_signal(execution_svc: Arc<dyn ExecutionService>, timeout: u64) {
     let ctrl_c = async {
         signal::ctrl_c()
             .await
@@ -355,13 +356,21 @@ async fn shutdown_signal() {
     #[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 => {
             tracing::warn!("Ctrl+C received, starting graceful shutdown");
         },
         () = terminate => {
             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");
@@ -371,3 +380,54 @@ async fn shutdown_signal() {
 #[derive(OpenApi)]
 #[openapi()]
 pub struct ApiDoc;
+
+#[cfg(test)]
+mod tests {
+    use api_snowflake_rest_sessions::session::SessionStore;
+    use executor::models::QueryContext;
+    use executor::service::ExecutionService;
+    use executor::service::make_test_execution_svc;
+    use 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");
+            }
+        }
+    }
+}

crates/executor/src/service.rs

@@ -42,6 +42,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>>;
@@ -127,6 +129,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 {
@@ -767,6 +771,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