Introduce initial-idle-shutdown-timeout-s (#528)

* hrana: Kick IdleKicker only after actual request is received

Previously we Hrana was kicking IdleKicker at startup but that's
no longer acceptable because we want to be able to have different
idle shutdown timeouts for startup that does not receive any requests
and startup that receives some requests.

Signed-off-by: Piotr Jastrzebski <[email protected]>

* Restructure IdleShutdownLayer:new

This change only restructures the IdleShutdownLayer::new.
It keeps the previous behaviour.
New structure of the code will allow introduction of the actual
changes in the next commit in much less invasive way.

Signed-off-by: Piotr Jastrzebski <[email protected]>

* Add initial-idle-shutdown-timeout-s

This new parameter makes it possible to define different
idle shutdown timeout for instance that has been started
but never received any requests.

It is useful for maintainance operations like periodic syncing
replica with the primary even when the replica does not serve
any requests at the moment. The assumption is that such maintainance
operation could wake up the replica for much shorter time than
idle-shutdown-timeout used after receiving actual user request.

Signed-off-by: Piotr Jastrzebski <[email protected]>

---------

Signed-off-by: Piotr Jastrzebski <[email protected]>

Author: haaawk

sqld/src/hrana/ws/conn.rs

@@ -82,10 +82,6 @@ async fn handle_ws<D: Database>(
     };
 
     loop {
-        if let Some(kicker) = conn.server.idle_kicker.as_ref() {
-            kicker.kick();
-        }
-
         tokio::select! {
             Some(client_msg_res) = conn.ws.recv() => {
                 let client_msg = client_msg_res
@@ -122,6 +118,10 @@ async fn handle_ws<D: Database>(
             },
             else => break,
         }
+
+        if let Some(kicker) = conn.server.idle_kicker.as_ref() {
+            kicker.kick();
+        }
     }
 
     close(

sqld/src/hrana/ws/mod.rs

@@ -50,10 +50,6 @@ pub async fn serve(
 
     let mut join_set = tokio::task::JoinSet::new();
     loop {
-        if let Some(kicker) = server.idle_kicker.as_ref() {
-            kicker.kick();
-        }
-
         tokio::select! {
             Some(accept) = accept_rx.recv() => {
                 let conn_id = server.next_conn_id.fetch_add(1, Ordering::AcqRel);
@@ -85,6 +81,10 @@ pub async fn serve(
                 return Ok(())
             }
         }
+
+        if let Some(kicker) = server.idle_kicker.as_ref() {
+            kicker.kick();
+        }
     }
 }
 

sqld/src/lib.rs

@@ -92,6 +92,7 @@ pub struct Config {
     pub rpc_server_ca_cert: Option<PathBuf>,
     pub bottomless_replication: Option<bottomless::replicator::Options>,
     pub idle_shutdown_timeout: Option<Duration>,
+    pub initial_idle_shutdown_timeout: Option<Duration>,
     pub load_from_dump: Option<PathBuf>,
     pub max_log_size: u64,
     pub max_log_duration: Option<f32>,
@@ -130,6 +131,7 @@ impl Default for Config {
             rpc_server_ca_cert: None,
             bottomless_replication: None,
             idle_shutdown_timeout: None,
+            initial_idle_shutdown_timeout: None,
             load_from_dump: None,
             max_log_size: 200,
             max_log_duration: None,
@@ -655,9 +657,13 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
             Ok(())
         });
 
-        let idle_shutdown_layer = config
-            .idle_shutdown_timeout
-            .map(|d| IdleShutdownLayer::new(d, shutdown_sender.clone()));
+        let idle_shutdown_layer = config.idle_shutdown_timeout.map(|d| {
+            IdleShutdownLayer::new(
+                d,
+                config.initial_idle_shutdown_timeout,
+                shutdown_sender.clone(),
+            )
+        });
 
         let stats = Stats::new(&config.db_path)?;
 

sqld/src/main.rs

@@ -122,6 +122,12 @@ struct Cli {
     #[clap(long, env = "SQLD_IDLE_SHUTDOWN_TIMEOUT_S")]
     idle_shutdown_timeout_s: Option<u64>,
 
+    /// Like idle_shutdown_timeout_s but used only once after the server is started.
+    /// After that server either is shut down because it does not receive any requests
+    /// or idle_shutdown_timeout_s is used moving forward.
+    #[clap(long, env = "SQLD_INITIAL_IDLE_SHUTDOWN_TIMEOUT_S")]
+    initial_idle_shutdown_timeout_s: Option<u64>,
+
     /// Load the dump at the provided path.
     /// Requires that the node is not in replica mode
     #[clap(long, env = "SQLD_LOAD_DUMP_PATH", conflicts_with = "primary_grpc_url")]
@@ -272,6 +278,9 @@ fn config_from_args(args: Cli) -> Result<Config> {
             None
         },
         idle_shutdown_timeout: args.idle_shutdown_timeout_s.map(Duration::from_secs),
+        initial_idle_shutdown_timeout: args
+            .initial_idle_shutdown_timeout_s
+            .map(Duration::from_secs),
         load_from_dump: args.load_from_dump,
         max_log_size: args.max_log_size,
         max_log_duration: args.max_log_duration,

sqld/src/utils/services/idle_shutdown.rs

@@ -14,31 +14,33 @@ pub struct IdleShutdownLayer {
 }
 
 impl IdleShutdownLayer {
-    pub fn new(idle_timeout: Duration, shutdown_notifier: mpsc::Sender<()>) -> Self {
+    pub fn new(
+        idle_timeout: Duration,
+        initial_idle_timeout: Option<Duration>,
+        shutdown_notifier: mpsc::Sender<()>,
+    ) -> Self {
         let (sender, mut receiver) = watch::channel(());
         let connected_replicas = Arc::new(AtomicUsize::new(0));
         let connected_replicas_clone = connected_replicas.clone();
+        let mut sleep_time = initial_idle_timeout.unwrap_or(idle_timeout);
         tokio::spawn(async move {
             loop {
                 // FIXME: if we measure that this is causing performance issues, we may want to
                 // implement some debouncing.
-                let timeout_fut = timeout(idle_timeout, receiver.changed());
-                match timeout_fut.await {
-                    Ok(Ok(_)) => continue,
-                    Ok(Err(_)) => break,
-                    Err(_) => {
-                        if connected_replicas_clone.load(Ordering::SeqCst) > 0 {
-                            continue;
-                        }
-                        tracing::info!(
-                            "Idle timeout, no new connection in {idle_timeout:.0?}. Shutting down.",
-                        );
-                        shutdown_notifier
-                            .send(())
-                            .await
-                            .expect("failed to shutdown gracefully");
-                    }
+                let timeout_res = timeout(sleep_time, receiver.changed()).await;
+                if let Ok(Err(_)) = timeout_res {
+                    break;
                 }
+                if timeout_res.is_err() && connected_replicas_clone.load(Ordering::SeqCst) == 0 {
+                    tracing::info!(
+                        "Idle timeout, no new connection in {sleep_time:.0?}. Shutting down.",
+                    );
+                    shutdown_notifier
+                        .send(())
+                        .await
+                        .expect("failed to shutdown gracefully");
+                }
+                sleep_time = idle_timeout;
             }
 
             tracing::debug!("idle shutdown loop exited");