Fix persistence workers never terminating (#401)

Hey!

This PR fixes an issue where memory is never reclaimed and keeps rising,
due to persistence workers never shutting down and freeing up the used
memory.

I'm running y-sweet as a long-running service, instead of via Plane or
another way to spin up a process per "room". And while experimenting I
also allowed images to be uploaded to the room state, which made the
fact that there was something up with the memory usage very visible, and
after adding a bit of logging/monitoring I could see that the number of
persistence workers never decreased.

In order to fix this, this PR adds a link between the GC Worker and the
persistence worker. When the GC detects that there's no more references
to the document it also notifies the persistence worker to shut down,
and when it does the memory is also freed.

Co-authored-by: Per-Victor Persson <pepe@netlight.com>

Author: Dracyr

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);