From c6a74b900078bae43f87b543ab948b3c65e553ea Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Mon, 19 Feb 2024 09:59:46 +1300
Subject: [PATCH] Trap signals and sync database on shutdown

---
 src/main.rs      | 60 +++++++++++++++++++++++++++++++++++-------------
 src/store/mod.rs |  7 ++++++
 2 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 1035741..6f2782e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -33,6 +33,7 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 use textnonce::TextNonce;
 use tokio::net::{TcpListener, TcpStream};
+use tokio::signal::unix::{signal, SignalKind};
 use tungstenite::protocol::WebSocketConfig;
 use tungstenite::Message;
 
@@ -76,26 +77,53 @@ async fn main() -> Result<(), Error> {
     // Store config into GLOBALS
     *GLOBALS.config.write().await = config;
 
-    // Accepts network connections and spawn a task to serve each one
-    loop {
-        let (tcp_stream, peer_addr) = listener.accept().await?;
+    let mut interrupt_signal = signal(SignalKind::interrupt())?;
+    let mut quit_signal = signal(SignalKind::quit())?;
+    let mut terminate_signal = signal(SignalKind::terminate())?;
 
-        if let Some(tls_acceptor) = &maybe_tls_acceptor {
-            let tls_acceptor_clone = tls_acceptor.clone();
-            tokio::spawn(async move {
-                match tls_acceptor_clone.accept(tcp_stream).await {
-                    Err(e) => log::error!("{}", e),
-                    Ok(tls_stream) => {
-                        if let Err(e) = serve(MaybeTlsStream::Rustls(tls_stream), peer_addr).await {
-                            log::error!("{}", e);
+    loop {
+        tokio::select! {
+            // Exits gracefully upon exit-type signals
+            v = interrupt_signal.recv() => if v.is_some() {
+                log::info!("SIGINT");
+                break;
+            },
+            v = quit_signal.recv() => if v.is_some() {
+                log::info!("SIGQUIT");
+                break;
+            },
+            v = terminate_signal.recv() => if v.is_some() {
+                log::info!("SIGTERM");
+                break;
+            },
+
+            // Accepts network connections and spawn a task to serve each one
+            v = listener.accept() => {
+                let (tcp_stream, peer_addr) = v?;
+
+                if let Some(tls_acceptor) = &maybe_tls_acceptor {
+                    let tls_acceptor_clone = tls_acceptor.clone();
+                    tokio::spawn(async move {
+                        match tls_acceptor_clone.accept(tcp_stream).await {
+                            Err(e) => log::error!("{}", e),
+                            Ok(tls_stream) => {
+                                if let Err(e) = serve(MaybeTlsStream::Rustls(tls_stream), peer_addr).await {
+                                    log::error!("{}", e);
+                                }
+                            }
                         }
-                    }
+                    });
+                } else {
+                    serve(MaybeTlsStream::Plain(tcp_stream), peer_addr).await?;
                 }
-            });
-        } else {
-            serve(MaybeTlsStream::Plain(tcp_stream), peer_addr).await?;
-        }
+            }
+        };
     }
+
+    log::info!("Syncing and shutting down.");
+    let _ = GLOBALS.store.get().unwrap().sync();
+
+    Ok(())
 }
 
 // Serve a single network connection
diff --git a/src/store/mod.rs b/src/store/mod.rs
index e96d880..ff3c86c 100644
--- a/src/store/mod.rs
+++ b/src/store/mod.rs
@@ -87,6 +87,13 @@ impl Store {
         })
     }
 
+    /// Sync the data to disk. This happens periodically, but sometimes it's useful to force
+    /// it.
+    pub fn sync(&self) -> Result<(), Error> {
+        self.env.force_sync()?;
+        Ok(())
+    }
+
     /// Store an event.
     ///
     /// Returns the offset where the event is stored at, which can be used to fetch