From 2c4384625a9042fc00438dcb83a951e5917f1640 Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Thu, 12 Aug 2021 13:37:47 +0100
Subject: [PATCH] feat: shutdown Database and Server on drop (#2241)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
---
 server/src/database.rs | 16 +++++++++++++++-
 server/src/lib.rs      | 13 +++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/server/src/database.rs b/server/src/database.rs
index fcb6df16be..0ddd5da739 100644
--- a/server/src/database.rs
+++ b/server/src/database.rs
@@ -10,7 +10,7 @@ use futures::{FutureExt, TryFutureExt};
 use generated_types::database_rules::encode_database_rules;
 use internal_types::freezable::Freezable;
 use object_store::path::{ObjectStorePath, Path};
-use observability_deps::tracing::{error, info};
+use observability_deps::tracing::{error, info, warn};
 use parking_lot::RwLock;
 use persistence_windows::checkpoint::ReplayPlan;
 use snafu::{ResultExt, Snafu};
@@ -250,6 +250,20 @@ impl Database {
     }
 }
 
+impl Drop for Database {
+    fn drop(&mut self) {
+        let db_name = &self.shared.config.name;
+        if !self.shared.shutdown.is_cancelled() {
+            warn!(%db_name, "database dropped without calling shutdown()");
+            self.shared.shutdown.cancel();
+        }
+
+        if self.join.clone().now_or_never().is_none() {
+            warn!(%db_name, "database dropped without waiting for worker termination");
+        }
+    }
+}
+
 /// State shared with the `Database` background worker
 #[derive(Debug)]
 struct DatabaseShared {
diff --git a/server/src/lib.rs b/server/src/lib.rs
index 002d7f5e0c..2cc411f198 100644
--- a/server/src/lib.rs
+++ b/server/src/lib.rs
@@ -366,6 +366,19 @@ pub struct Server<M: ConnectionManager> {
     shared: Arc<ServerShared>,
 }
 
+impl<M: ConnectionManager> Drop for Server<M> {
+    fn drop(&mut self) {
+        if !self.shared.shutdown.is_cancelled() {
+            warn!("server dropped without calling shutdown()");
+            self.shared.shutdown.cancel();
+        }
+
+        if self.join.clone().now_or_never().is_none() {
+            warn!("server dropped without waiting for worker termination");
+        }
+    }
+}
+
 #[derive(Debug)]
 struct ServerShared {
     /// A token that is used to trigger shutdown of the background worker