feat: replace PersistSemaphore

RodrigoVillar · RodrigoVillar · commit d63cb4590c1b · 2026-02-25T10:26:24.000-05:00
diff --git a/firewood/src/persist_worker.rs b/firewood/src/persist_worker.rs
@@ -109,11 +109,11 @@ impl PersistWorker {
             channel: PersistChannel::new(commit_count, persist_threshold),
         });
 
-        let persist_loop = PersistLoop {
-            shared: shared.clone(),
-        };
-
-        let handle = thread::spawn(move || persist_loop.run());
+        let bg_shared = shared.clone();
+        let handle = thread::spawn(move || {
+            let _guard = PanicGuard(bg_shared.clone());
+            PersistLoop { shared: bg_shared }.run()
+        });
 
         Self {
             handle: Mutex::new(Some(handle)),
@@ -124,7 +124,14 @@ impl PersistWorker {
     /// Sends `committed` to the background thread for persistence. This call
     /// blocks if the limit of unpersisted commits has been reached.
     pub(crate) fn persist(&self, committed: CommittedRevision) -> Result<(), PersistError> {
-        self.shared.channel.push(committed)
+        if let Err(e) = self.shared.channel.push(committed) {
+            self.join_handle();
+            self.check_error()?;
+
+            return Err(e);
+        }
+
+        Ok(())
     }
 
     /// Sends `nodestore` to the background thread for reaping if archival mode
@@ -133,9 +140,15 @@ impl PersistWorker {
         &self,
         nodestore: NodeStore<Committed, FileBacked>,
     ) -> Result<(), PersistError> {
-        if self.shared.root_store.is_none() {
-            self.shared.channel.reap(nodestore)?;
+        if self.shared.root_store.is_none()
+            && let Err(e) = self.shared.channel.reap(nodestore)
+        {
+            self.join_handle();
+            self.check_error()?;
+
+            return Err(e);
         }
+
         Ok(())
     }
 
@@ -154,12 +167,13 @@ impl PersistWorker {
 
     /// Close the persist worker and persist `latest_committed_revision`.
     pub(crate) fn close(
-        mut self,
+        self,
         latest_committed_revision: CommittedRevision,
     ) -> Result<(), PersistError> {
         self.shared
             .persist_on_shutdown
-            .set(latest_committed_revision);
+            .set(latest_committed_revision)
+            .expect("should be empty");
 
         self.shared.channel.close();
         self.join_handle();
@@ -190,15 +204,15 @@ impl PersistWorker {
 }
 
 #[derive(Debug)]
-struct PersistChannel<T> {
-    state: Mutex<PersistChannelState<T>>,
+struct PersistChannel {
+    state: Mutex<PersistChannelState>,
     /// Condition variable to wake blocked committers when space is available.
     commit_not_full: Condvar,
     /// Condition variable to wake the persister when persistence is needed.
     persist_ready: Condvar,
 }
 
-impl<T> PersistChannel<T> {
+impl PersistChannel {
     const fn new(max_permits: NonZeroU64, persist_threshold: u64) -> Self {
         Self {
             state: Mutex::new(PersistChannelState {
@@ -229,9 +243,9 @@ impl<T> PersistChannel<T> {
         Ok(())
     }
 
-    fn push(&self, item: T) -> Result<(), PersistError> {
+    fn push(&self, item: CommittedRevision) -> Result<(), PersistError> {
         let mut state = self.state.lock();
-        while state.permits_available == 0 {
+        while state.permits_available == 0 && !state.shutdown {
             self.commit_not_full.wait(&mut state);
         }
 
@@ -249,37 +263,39 @@ impl<T> PersistChannel<T> {
         Ok(())
     }
 
-    fn pop(&self) -> Result<PersistDataWrapper<'_, T>, PersistError> {
-        let mut state = self.state.lock();
-        let (permits_to_release, pending_reaps, data) = loop {
-            // Shutdown requested. Return error.
-            if state.shutdown {
-                return Err(PersistError::ChannelDisconnected);
+    fn pop(&self) -> Result<PersistDataWrapper<'_>, PersistError> {
+        let (permits_to_release, pending_reaps, data) = {
+            let mut state = self.state.lock();
+            loop {
+                // Shutdown requested. Return error.
+                if state.shutdown {
+                    return Err(PersistError::ChannelDisconnected);
+                }
+                // Unblock to persist when permits available <= threshold
+                if state.permits_available <= state.persist_threshold && state.data.is_some() {
+                    break (
+                        state
+                            .max_permits
+                            .get()
+                            .saturating_sub(state.permits_available),
+                        std::mem::take(&mut state.pending_reaps),
+                        state.data.take(),
+                    );
+                }
+                // Unblock even if we haven't met the threshold if there are pending reaps.
+                // Permits to release is set to 0, and committed revision is not taken.
+                if !state.pending_reaps.is_empty() {
+                    break (0, std::mem::take(&mut state.pending_reaps), None);
+                }
+                // Block until it is woken up by the committer thread.
+                self.persist_ready.wait(&mut state);
             }
-            // Unblock to persist when permits available <= threshold
-            if state.permits_available <= state.persist_threshold && state.data.is_some() {
-                break (
-                    state
-                        .max_permits
-                        .get()
-                        .saturating_sub(state.permits_available),
-                    std::mem::take(&mut state.pending_reaps),
-                    state.data.take(),
-                );
-            }
-            // Unblock even if we haven't met the threshold if there are pending reaps.
-            // Permits to release is set to 0, and committed revision is not taken.
-            if !state.pending_reaps.is_empty() {
-                break (0, std::mem::take(&mut state.pending_reaps), None);
-            }
-            // Block until it is woken up by the committer thread.
-            self.persist_ready.wait(&mut state);
         };
         Ok(PersistDataWrapper {
             channel: self,
             permits_to_release,
             pending_reaps,
-            data,
+            latest_committed: data,
         })
     }
 
@@ -311,25 +327,37 @@ impl<T> PersistChannel<T> {
 }
 
 #[derive(Debug)]
-struct PersistChannelState<T> {
+struct PersistChannelState {
     permits_available: u64,
     /// Maximum number of unpersisted commits allowed.
     max_permits: NonZeroU64,
     /// Persist when remaining permits are at or below this threshold.
     persist_threshold: u64,
     shutdown: bool,
     pending_reaps: Vec<NodeStore<Committed, FileBacked>>,
-    data: Option<T>,
+    data: Option<CommittedRevision>,
 }
 
-struct PersistDataWrapper<'a, T> {
-    channel: &'a PersistChannel<T>,
+/// RAII wrapper returned by [`PersistChannel::pop`] that carries the data for
+/// one persistence cycle (a committed revision and/or pending reaps).
+///
+/// On drop, it returns consumed permits back to the channel and notifies
+/// blocked committers via [`commit_not_full`](PersistChannel::commit_not_full),
+/// ensuring backpressure is released even if the persist loop exits early due
+/// to an error.
+struct PersistDataWrapper<'a> {
+    channel: &'a PersistChannel,
+    /// Number of permits consumed by commits since the last persist cycle.
+    /// Released back to the channel on drop.
     permits_to_release: u64,
+    /// Expired node stores whose deleted nodes should be returned to free lists.
     pending_reaps: Vec<NodeStore<Committed, FileBacked>>,
-    data: Option<T>,
+    /// The latest committed revision to persist, if the threshold was reached.
+    /// `None` when this cycle was triggered solely by pending reaps.
+    latest_committed: Option<CommittedRevision>,
 }
 
-impl<T> Drop for PersistDataWrapper<'_, T> {
+impl Drop for PersistDataWrapper<'_> {
     /// Handle permit release on drop of the wrapper
     fn drop(&mut self) {
         if self.permits_to_release == 0 {
@@ -359,7 +387,21 @@ struct SharedState {
     /// Unpersisted revision to persist on shutdown.
     persist_on_shutdown: OnceLock<CommittedRevision>,
     /// Channel for coordinating persist and reap operations.
-    channel: PersistChannel<CommittedRevision>,
+    channel: PersistChannel,
+}
+
+/// Closes the persist channel if the background thread panics.
+///
+/// This ensures blocked committers are woken up and see the shutdown
+/// state rather than blocking indefinitely.
+struct PanicGuard(Arc<SharedState>);
+
+impl Drop for PanicGuard {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            self.0.channel.close();
+        }
+    }
 }
 
 /// The background persistence loop that runs in a separate thread.
@@ -393,9 +435,9 @@ impl PersistLoop {
                 self.reap(nodestore)?;
             }
 
-            if let Some(revision) = persist_data.data.take() {
+            if let Some(revision) = persist_data.latest_committed.take() {
                 self.persist_to_disk(&revision)
-                    .and_then(|()| self.save_to_root_store(&revision))?;
+                    .and_then(|()| self.maybe_save_to_root_store(&revision))?;
             }
         }
 
@@ -404,7 +446,7 @@ impl PersistLoop {
             && !self.shared.channel.empty()
         {
             self.persist_to_disk(&revision)
-                .and_then(|()| self.save_to_root_store(&revision))?;
+                .and_then(|()| self.maybe_save_to_root_store(&revision))?;
         }
 
         Ok(())
@@ -423,7 +465,10 @@ impl PersistLoop {
     fn reap(&self, nodestore: NodeStore<Committed, FileBacked>) -> Result<(), PersistError> {
         nodestore
             .reap_deleted(&mut self.shared.header.lock())
-            .map_err(|e| PersistError::FileIo(Arc::new(e)))
+            .map_err(|e| {
+                error!("Failed to reap deleted nodes: {e}");
+                PersistError::FileIo(Arc::new(e))
+            })
     }
 
     /// Saves the revision's root address to `RootStore` if configured.
@@ -438,9 +483,20 @@ impl PersistLoop {
     }
 }
 
+#[crate::metrics("persist.root_store", "persist revision address to root store")]
+fn save_to_root_store(
+    store: &RootStore,
+    hash: &TrieHash,
+    addr: &LinearAddress,
+) -> Result<(), PersistError> {
+    store.add_root(hash, addr).map_err(|e| {
+        error!("Failed to persist revision address to RootStore: {e}");
+        PersistError::RootStore(e.into())
+    })
+}
+
 impl SharedState {
     #[cfg(test)]
-    #[allow(clippy::arithmetic_side_effects)]
     fn wait_all_released(&self) {
         self.channel.wait_all_released();
     }
