Fix persistence workers never terminating due to never receiving a shutdown signal

Author: Per-Victor Persson

crates/y-sweet-core/src/sync_kv.rs

@@ -17,6 +17,7 @@ pub struct SyncKv {
     key: String,
     dirty: AtomicBool,
     dirty_callback: Box<dyn Fn() + Send + Sync>,
+    shutdown: AtomicBool,
 }
 
 impl SyncKv {
@@ -44,11 +45,12 @@ impl SyncKv {
             key,
             dirty: AtomicBool::new(false),
             dirty_callback: Box::new(callback),
+            shutdown: AtomicBool::new(false),
         })
     }
 
     fn mark_dirty(&self) {
-        if !self.dirty.load(Ordering::Relaxed) {
+        if !self.dirty.load(Ordering::Relaxed) && !self.shutdown.load(Ordering::SeqCst) {
             self.dirty.store(true, Ordering::Relaxed);
             (self.dirty_callback)();
         }
@@ -88,6 +90,16 @@ impl SyncKv {
     pub fn is_empty(&self) -> bool {
         self.data.lock().unwrap().is_empty()
     }
+
+    pub fn is_shutdown(&self) -> bool {
+        self.shutdown.load(Ordering::SeqCst)
+    }
+
+    pub fn shutdown(&self) {
+        self.shutdown.store(true, Ordering::SeqCst);
+        // Call the callback one last time to wake up the persistence worker
+        (self.dirty_callback)();
+    }
 }
 
 impl<'d> DocOps<'d> for SyncKv {}

crates/y-sweet/src/server.rs

@@ -204,6 +204,10 @@ impl Server {
 
                     if checkpoints_without_refs >= 2 {
                         tracing::info!("GCing doc");
+                        if let Some(doc) = docs.get(&doc_id) {
+                            doc.sync_kv().shutdown();
+                        }
+
                         docs.remove(&doc_id);
                         break;
                     }
@@ -229,6 +233,9 @@ impl Server {
             let is_done = tokio::select! {
                 v = recv.recv() => v.is_none(),
                 _ = cancellation_token.cancelled() => true,
+                _ = tokio::time::sleep(checkpoint_freq) => {
+                    sync_kv.is_shutdown()
+                }
             };
 
             tracing::info!("Received signal. done: {}", is_done);