add signal handlers

edgeandnode · Oct 18, 2024 · 318996f · 318996f
1 parent 82b4ca4
commit 318996f
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,11 +17,12 @@ serde = { version = "1.0.210", features = ["derive"] }
 serde_json = "1.0.128"
 snmalloc-rs = "0.3.6"
 tokio = { version = "1.40.0", features = [
-  "sync",
   "macros",
   "parking_lot",
   "rt-multi-thread",
   "time",
+  "signal",
+  "sync",
 ] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = [

diff --git a/src/bin/main.rs b/src/bin/main.rs
@@ -92,9 +92,12 @@ async fn run() -> anyhow::Result<()> {
         .into_iter()
         .map(|e| (format!("{}/{}", e.topic(), e.partition()), 0))
         .collect();
+    let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
     loop {
         tokio::select! {
             biased;
+            _ = tokio::signal::ctrl_c() => anyhow::bail!("exit"),
+            _ = sigterm.recv() => anyhow::bail!("exit"),
             _ = consumer.recv() => anyhow::bail!("message received from split consumer"),
             _ = select_all(&mut partition_consumers) => anyhow::bail!("parition consumer exit"),
             _ = source_msg_rx.recv_many(&mut msg_buffer, msg_buffer_limit) => {