Properly flush and exit on ctrl-c

Fixes #60

Author: rtyler

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hotdog"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["R. Tyler Croy <[email protected]>"]
 edition = "2024"
 

src/serve/mod.rs

@@ -5,6 +5,7 @@
 use crate::connection::*;
 use crate::errors;
 use crate::settings::Settings;
+use crate::sink::Message;
 use crate::sink::Sink;
 use crate::sink::kafka::Kafka;
 use crate::sink::parquet::Parquet;
@@ -22,44 +23,34 @@ use std::sync::Arc;
 pub mod plain;
 pub mod tls;
 
+/// State entity which can be passed into connection handlers and callbacks
+#[derive(Clone)]
 pub struct ServerState {
-    /**
-     * A reference to the global Settings object for all configuration information
-     */
+    /// A reference to the global Settings object for all configuration information
     pub settings: Arc<Settings>,
-    /**
-     * A Sender for sending statistics to the status handler
-     */
+    /// A Sender for sending statistics to the status handler
     pub stats: InputQueueScope,
 }
 
-/**
- * The Server trait describes the necessary functionality to implement a new hotdog backend server
- * which can receive syslog messages
- */
+// The Server trait describes the necessary functionality to implement a new hotdog backend server
+// which can receive syslog messages
 #[async_trait]
 pub trait Server {
-    /**
-     * Bootstrap can/should be overridden by implementations which need to perform some work prior
-     * to the creation of the TcpListener and the incoming connection loop
-     */
+    /// Bootstrap can/should be overridden by implementations which need to perform some work prior
+    /// to the creation of the TcpListener and the incoming connection loop
     fn bootstrap(&mut self, _state: &ServerState) -> Result<(), errors::HotdogError> {
         Ok(())
     }
 
-    /**
-     * Shutdown scan/should be overridden by implementations which need to perform some work after
-     * the termination of the connection accept loop
-     */
+    /// Shutdown scan/should be overridden by implementations which need to perform some work after
+    /// the termination of the connection accept loop
     fn shutdown(&self, _state: &ServerState) -> Result<(), errors::HotdogError> {
         Ok(())
     }
 
-    /**
-     * Handle a single connection
-     *
-     * The close_channel parameter must be a clone of our connection-tracking channel Sender
-     */
+    /// Handle a single connection
+    ///
+    /// The close_channel parameter must be a clone of our connection-tracking channel Sender
     fn handle_connection(
         &self,
         stream: TcpStream,
@@ -78,9 +69,7 @@ pub trait Server {
         Ok(())
     }
 
-    /**
-     * Accept connections on the addr
-     */
+    /// Accept connections on the addr
     async fn accept_loop(
         &mut self,
         addr: &str,
@@ -119,11 +108,29 @@ pub trait Server {
             smol::spawn(async move {
                 debug!("Starting Parquet loop");
                 pq.runloop().await;
+                debug!("Ending Parquet loop");
+                std::process::exit(0);
             })
             .detach();
         }
 
         let sender = sender.expect("Failed to configure a sink properly!");
+        use std::sync::atomic::{AtomicBool, Ordering};
+        let should_exit = Arc::new(AtomicBool::new(false));
+        let se = should_exit.clone();
+
+        let ctrlc_tx = sender.clone();
+
+        ctrlc::set_handler(move || {
+            info!("Interrupt has been received! Attempting to flush");
+            se.store(true, Ordering::SeqCst);
+            let tx = ctrlc_tx.clone();
+            smol::block_on(async move {
+                tx.send(Message::Flush { should_exit: true })
+                    .await
+                    .expect("Failed to send flush command");
+            });
+        });
 
         self.bootstrap(&state)?;
 
@@ -152,10 +159,21 @@ pub trait Server {
                 .stats
                 .gauge(status::Stats::ConnectionCount.into())
                 .value(conn_count);
+
+            if should_exit.load(Ordering::Relaxed) {
+                debug!("Serve has been instructed to exit");
+                break;
+            }
         }
 
         self.shutdown(&state)?;
 
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_placholder() {}
+}

src/sink/mod.rs

@@ -59,7 +59,9 @@ pub enum Message {
         destination: String,
         payload: String,
     },
-    Flush,
+    Flush {
+        should_exit: bool,
+    },
 }
 
 impl Message {
@@ -69,6 +71,10 @@ impl Message {
             payload,
         }
     }
+
+    pub fn flush() -> Self {
+        Message::Flush { should_exit: false }
+    }
 }
 
 #[cfg(test)]

src/sink/parquet.rs

@@ -104,7 +104,7 @@ impl Sink for Parquet {
             let mut timer = smol::Timer::interval(interval);
             while timer.next().await.is_some() {
                 debug!("Timer has fired, issuing a flush");
-                if let Err(e) = timer_tx.send(Message::Flush).await {
+                if let Err(e) = timer_tx.send(Message::flush()).await {
                     error!("Failed to trigger the flush timer in the parquet sink: {e:?}");
                 }
             }
@@ -162,11 +162,11 @@ impl Sink for Parquet {
                                     "Reached the threshold to flush bytes for `{}`",
                                     &destination
                                 );
-                                let _ = self.tx.send(Message::Flush).await;
+                                let _ = self.tx.send(Message::flush()).await;
                             }
                         }
                     }
-                    Message::Flush => {
+                    Message::Flush { should_exit } => {
                         info!("Parquet sink has been told to flush");
 
                         for (destination, buf) in buffer.drain() {
@@ -196,6 +196,11 @@ impl Sink for Parquet {
                             flush_to_parquet(self.store.clone(), schema, &destination, &buf);
                             since_last_flush = Instant::now();
                         }
+
+                        if should_exit {
+                            debug!("Supposed to exit from the sink!");
+                            return;
+                        }
                     }
                 }
             }