Fix race condition in shutdown of background task (#81)

src/db.rs

@@ -4,6 +4,17 @@ use tokio::time::{self, Duration, Instant};
 use bytes::Bytes;
 use std::collections::{BTreeMap, HashMap};
 use std::sync::{Arc, Mutex};
+use tracing::debug;
+
+/// A wrapper around a `Db` instance. This exists to allow orderly cleanup
+/// of the `Db` by signalling the background purge task to shut down when
+/// this struct is dropped.
+#[derive(Debug)]
+pub(crate) struct DbDropGuard {
+    /// The `Db` instance that will be shut down when this `DbHolder` struct
+    /// is dropped.
+    db: Db,
+}
 
 /// Server state shared across all connections.
 ///
@@ -92,6 +103,27 @@ struct Entry {
     expires_at: Option<Instant>,
 }
 
+impl DbDropGuard {
+    /// Create a new `DbHolder`, wrapping a `Db` instance. When this is dropped
+    /// the `Db`'s purge task will be shut down.
+    pub(crate) fn new() -> DbDropGuard {
+        DbDropGuard { db: Db::new() }
+    }
+
+    /// Get the shared database. Internally, this is an
+    /// `Arc`, so a clone only increments the ref count.
+    pub(crate) fn db(&self) -> Db {
+        self.db.clone()
+    }
+}
+
+impl Drop for DbDropGuard {
+    fn drop(&mut self) {
+        // Signal the 'Db' instance to shut down the task that purges expired keys
+        self.db.shutdown_purge_task();
+    }
+}
+
 impl Db {
     /// Create a new, empty, `Db` instance. Allocates shared state and spawns a
     /// background task to manage key expiration.
@@ -244,28 +276,20 @@ impl Db {
             // subscribers. In this case, return `0`.
             .unwrap_or(0)
     }
-}
 
-impl Drop for Db {
-    fn drop(&mut self) {
-        // If this is the last active `Db` instance, the background task must be
-        // notified to shut down.
-        //
-        // First, determine if this is the last `Db` instance. This is done by
-        // checking `strong_count`. The count will be 2. One for this `Db`
-        // instance and one for the handle held by the background task.
-        if Arc::strong_count(&self.shared) == 2 {
-            // The background task must be signaled to shutdown. This is done by
-            // setting `State::shutdown` to `true` and signalling the task.
-            let mut state = self.shared.state.lock().unwrap();
-            state.shutdown = true;
+    /// Signals the purge background task to shut down. This is called by the
+    /// `DbShutdown`s `Drop` implementation.
+    fn shutdown_purge_task(&self) {
+        // The background task must be signaled to shut down. This is done by
+        // setting `State::shutdown` to `true` and signalling the task.
+        let mut state = self.shared.state.lock().unwrap();
+        state.shutdown = true;
 
-            // Drop the lock before signalling the background task. This helps
-            // reduce lock contention by ensuring the background task doesn't
-            // wake up only to be unable to acquire the mutex.
-            drop(state);
-            self.shared.background_task.notify_one();
-        }
+        // Drop the lock before signalling the background task. This helps
+        // reduce lock contention by ensuring the background task doesn't
+        // wake up only to be unable to acquire the mutex.
+        drop(state);
+        self.shared.background_task.notify_one();
     }
 }
 
@@ -349,4 +373,6 @@ async fn purge_expired_tasks(shared: Arc<Shared>) {
             shared.background_task.notified().await;
         }
     }
+
+    debug!("Purge background task shut down")
 }