fix(FileStore): Stop keep-alive thread on cancellation token. (#4666)

Signed-off-by: Graham King <grahamk@nvidia.com>

lib/runtime/src/distributed.rs

@@ -111,7 +111,7 @@ impl DistributedRuntime {
                     tracing::error!(%err, "Could not connect to etcd. Pass `--store-kv ..` to use a different backend or start etcd."))?;
                 kv::Manager::etcd(etcd_client)
             }
-            kv::Selector::File(root) => kv::Manager::file(root),
+            kv::Selector::File(root) => kv::Manager::file(runtime.primary_token(), root),
             kv::Selector::Memory => kv::Manager::memory(),
         };
 

lib/runtime/src/storage/kv.rs

@@ -265,8 +265,8 @@ impl Manager {
         Self::new(KeyValueStoreEnum::Etcd(EtcdStore::new(etcd_client)))
     }
 
-    pub fn file<P: Into<PathBuf>>(root: P) -> Self {
-        Self::new(KeyValueStoreEnum::File(FileStore::new(root)))
+    pub fn file<P: Into<PathBuf>>(cancel_token: CancellationToken, root: P) -> Self {
+        Self::new(KeyValueStoreEnum::File(FileStore::new(cancel_token, root)))
     }
 
     fn new(s: KeyValueStoreEnum) -> Manager {

lib/runtime/src/storage/kv/file.rs

@@ -20,6 +20,7 @@ use async_trait::async_trait;
 use futures::StreamExt;
 use notify::{Config, Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher, event};
 use parking_lot::Mutex;
+use tokio_util::sync::CancellationToken;
 
 use super::{Bucket, Key, KeyValue, Store, StoreError, StoreOutcome, WatchEvent};
 
@@ -33,6 +34,7 @@ const MIN_KEEP_ALIVE: Duration = Duration::from_secs(1);
 /// Treat as a singleton
 #[derive(Clone)]
 pub struct FileStore {
+    cancel_token: CancellationToken,
     root: PathBuf,
     connection_id: u64,
     /// Directories we may have created files in, for shutdown cleanup and keep-alive.
@@ -41,8 +43,9 @@ pub struct FileStore {
 }
 
 impl FileStore {
-    pub(super) fn new<P: Into<PathBuf>>(root_dir: P) -> Self {
+    pub(super) fn new<P: Into<PathBuf>>(cancel_token: CancellationToken, root_dir: P) -> Self {
         let fs = FileStore {
+            cancel_token,
             root: root_dir.into(),
             connection_id: rand::random::<u64>(),
             active_dirs: Arc::new(Mutex::new(HashMap::new())),
@@ -52,16 +55,27 @@ impl FileStore {
         fs
     }
 
-    /// Keep our files alive and delete expired keys. Does not return.
-    /// We run this in a real thread so it doesn't get delayed by tokio runtime load.
-    /// It doesn't need any cleanup so we don't use cancellation token.
-    fn expiry_thread(&self) -> ! {
+    /// Keep our files alive and delete expired keys.
+    ///
+    /// Does not return until cancellation token cancelled. On shutdown the process will
+    /// often exit before we detect cancellation. That's fine.
+    /// We run this in a real thread so it doesn't get delayed by tokio runtime under heavy load.
+    fn expiry_thread(&self) {
         loop {
             let ttl = self.shortest_ttl();
             let keep_alive_interval = cmp::max(ttl / 3, MIN_KEEP_ALIVE);
 
+            // Check before and after the sleep
+            if self.cancel_token.is_cancelled() {
+                break;
+            }
+
             thread::sleep(keep_alive_interval);
 
+            if self.cancel_token.is_cancelled() {
+                break;
+            }
+
             self.keep_alive();
             if let Err(err) = self.delete_expired_files() {
                 tracing::error!(error = %err, "FileStore delete_expired_files");
@@ -469,13 +483,16 @@ fn to_fs_err<E: std::error::Error>(err: E) -> StoreError {
 mod tests {
     use std::collections::HashSet;
 
+    use tokio_util::sync::CancellationToken;
+
     use crate::storage::kv::{Bucket as _, FileStore, Key, Store as _};
 
     #[tokio::test]
     async fn test_entries_full_path() {
         let t = tempfile::tempdir().unwrap();
 
-        let m = FileStore::new(t.path());
+        let cancel_token = CancellationToken::new();
+        let m = FileStore::new(cancel_token.clone(), t.path());
         let bucket = m.get_or_create_bucket("v1/tests", None).await.unwrap();
         let _ = bucket
             .insert(&Key::new("key1/multi/part".to_string()), "value1".into(), 0)
@@ -487,6 +504,7 @@ mod tests {
             .unwrap();
         let entries = bucket.entries().await.unwrap();
         let keys: HashSet<Key> = entries.into_keys().collect();
+        cancel_token.cancel(); // stop the background thread
 
         assert!(keys.contains(&Key::new("v1/tests/key1/multi/part".to_string())));
         assert!(keys.contains(&Key::new("v1/tests/key2".to_string())));