From 089305b57d162e6d90a311635783fcf5e5b78b93 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Mon, 19 Feb 2024 10:38:19 +1300
Subject: [PATCH] Gracefully shutdown

---
 src/globals.rs | 16 +++++++++++----
 src/main.rs    | 53 +++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/src/globals.rs b/src/globals.rs
index fda6936..77bcb68 100644
--- a/src/globals.rs
+++ b/src/globals.rs
@@ -2,8 +2,10 @@ use crate::config::{Config, FriendlyConfig};
 use crate::store::Store;
 use hyper::server::conn::Http;
 use lazy_static::lazy_static;
+use std::sync::atomic::AtomicUsize;
 use std::sync::OnceLock;
-use tokio::sync::broadcast::Sender;
+use tokio::sync::broadcast::Sender as BroadcastSender;
+use tokio::sync::watch::Sender as WatchSender;
 use tokio::sync::RwLock;
 
 pub struct Globals {
@@ -16,7 +18,10 @@ pub struct Globals {
     /// Every handler needs to listen to it and check if the incoming event matches any
     /// subscribed fitlers for their client, and if so, send the event to their client under
     /// that subscription.
-    pub new_events: Sender<usize>,
+    pub new_events: BroadcastSender<usize>,
+
+    pub num_clients: AtomicUsize,
+    pub shutting_down: WatchSender<bool>,
 }
 
 lazy_static! {
@@ -25,14 +30,17 @@ lazy_static! {
         http_server.http1_only(true);
         http_server.http1_keep_alive(true);
 
-        let (sender, _) = tokio::sync::broadcast::channel(512);
+        let (new_events, _) = tokio::sync::broadcast::channel(512);
+        let (shutting_down, _) = tokio::sync::watch::channel(false);
 
         Globals {
             config: RwLock::new(FriendlyConfig::default().into_config().unwrap()),
             store: OnceLock::new(),
             http_server,
             rid: OnceLock::new(),
-            new_events: sender,
+            new_events,
+            num_clients: AtomicUsize::new(0),
+            shutting_down,
         }
     };
 }
diff --git a/src/main.rs b/src/main.rs
index 6f2782e..984afc8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -30,7 +30,9 @@ use std::future::Future;
 use std::io::Read;
 use std::net::SocketAddr;
 use std::pin::Pin;
+use std::sync::atomic::Ordering;
 use std::task::{Context, Poll};
+use std::time::Duration;
 use textnonce::TextNonce;
 use tokio::net::{TcpListener, TcpStream};
 use tokio::signal::unix::{signal, SignalKind};
@@ -120,6 +122,41 @@ async fn main() -> Result<(), Error> {
         };
     }
 
+    // Pre-sync in case something below hangs up
+    let _ = GLOBALS.store.get().unwrap().sync();
+
+    // Set the shutting down signal
+    let _ = GLOBALS.shutting_down.send(true);
+
+    // Wait for active websockets to shutdown gracefully
+    let mut num_clients = GLOBALS.num_clients.load(Ordering::Relaxed);
+    if num_clients != 0 {
+        log::info!("Waiting for {num_clients} websockets to shutdown...");
+
+        // We will check if all clients have shutdown every 25ms
+        let sleep = tokio::time::sleep(Duration::from_millis(25));
+        tokio::pin!(sleep);
+
+        while num_clients != 0 {
+            // If we get another shutdown signal, stop waiting for websockets
+            tokio::select! {
+                v = interrupt_signal.recv() => if v.is_some() {
+                    break;
+                },
+                v = quit_signal.recv() => if v.is_some() {
+                    break;
+                },
+                v = terminate_signal.recv() => if v.is_some() {
+                    break;
+                },
+                () = &mut sleep => {
+                    num_clients = GLOBALS.num_clients.load(Ordering::Relaxed);
+                    continue;
+                }
+            }
+        }
+    }
+
     log::info!("Syncing and shutting down.");
     let _ = GLOBALS.store.get().unwrap().sync();
 
@@ -209,6 +246,9 @@ async fn handle_http_request(
                         user: None,
                     };
 
+                    // Increment count of active websockets
+                    let _ = GLOBALS.num_clients.fetch_add(1, Ordering::SeqCst);
+
                     // Handle the websocket
                     if let Err(e) = ws_service.handle_websocket_stream().await {
                         if matches!(
@@ -223,6 +263,9 @@ async fn handle_http_request(
                         }
                     }
 
+                    // DecrementIncrement count of active websockets
+                    let _ = GLOBALS.num_clients.fetch_sub(1, Ordering::SeqCst);
+
                     log::info!("{}: websocket ended", peer);
                 }
                 Err(e) => {
@@ -263,6 +306,9 @@ impl WebSocketService {
     }
 
     async fn handle_websocket_stream(&mut self) -> Result<(), Error> {
+        // Subscribe to the shutting down channel
+        let mut shutting_down = GLOBALS.shutting_down.subscribe();
+
         // Subscribe to the new_events broadcast channel
         let mut new_events = GLOBALS.new_events.subscribe();
 
@@ -278,13 +324,18 @@ impl WebSocketService {
                             let message = message?;
                             self.handle_websocket_message(message).await?;
                         },
-                        None => break, // websocket must be closed
+                        None => break, // the websocket is closed
                     }
                 },
                 offset_result = new_events.recv() => {
                     let offset = offset_result?;
                     self.handle_new_event(offset).await?;
                 },
+                _r = shutting_down.changed() => {
+                    // Shutdown the websocket gracefully
+                    self.websocket.send(Message::Close(None)).await?;
+                    break;
+                },
             }
         }