From 1debd08c026c6b9d66fed48f5d454fad425919e2 Mon Sep 17 00:00:00 2001
From: Amanda Graven <amanda@amandag.net>
Date: Tue, 22 Jun 2021 17:03:32 +0200
Subject: [PATCH] Add graceful shutdown on SIGTERM

---
 Cargo.toml  |  2 +-
 src/main.rs | 30 ++++++++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2d21b0c..61a0fc9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,7 +17,7 @@ humantime = "2.1"
 rpassword = "5.0"
 serde_json = "1.0"
 time = "0.2"
-tokio = { version = "1.0", features = ["macros", "time", "rt-multi-thread"]}
+tokio = { version = "1.6", features = ["macros", "signal", "time", "rt-multi-thread"]}
 
 [dependencies.matrix-sdk]
 git = "https://github.com/matrix-org/matrix-rust-sdk"
diff --git a/src/main.rs b/src/main.rs
index ac87051..507a867 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,4 @@
-use std::{convert::TryFrom, fs::File, time::Duration};
+use std::{convert::TryFrom, fs::File, sync::Arc, time::Duration};
 
 use matrix_sdk::{
 	events::{
@@ -13,6 +13,7 @@ use matrix_sdk::{
 	identifiers::{RoomId, UserId},
 	Client, ClientConfig, LoopCtrl, Session, SyncSettings,
 };
+use tokio::signal::unix::SignalKind;
 
 const HELP_MESSAGE_PLAIN: &str = r#"Help:
 !remindme help - show this help message
@@ -99,9 +100,28 @@ async fn restore_login() -> Result<Client, Box<dyn std::error::Error>> {
 }
 
 async fn sync(client: Client) {
+	let stop = Arc::new(tokio::sync::RwLock::new(false));
+	// install signal handler
+	{
+		let stop = stop.clone();
+		tokio::task::spawn(async move {
+			let mut stream = match tokio::signal::unix::signal(SignalKind::terminate()) {
+				Ok(stream) => stream,
+				Err(e) => {
+					eprintln!("Attaching signal handler failed: {}", e);
+					return;
+				}
+			};
+			loop {
+				stream.recv().await;
+				*stop.write().await = true;
+			}
+		});
+	}
 	client
 		.sync_with_callback(SyncSettings::new(), |response| {
 			let client = client.clone();
+			let stop = stop.clone();
 			async move {
 				for (room_id, room) in response.rooms.invite {
 					for event in room.invite_state.events {
@@ -131,7 +151,13 @@ async fn sync(client: Client) {
 						};
 					}
 				}
-				LoopCtrl::Continue
+				match *stop.read().await {
+					false => LoopCtrl::Continue,
+					true => {
+						eprintln!("Shutting down gracefully");
+						LoopCtrl::Break
+					}
+				}
 			}
 		})
 		.await;