Get rid of the weak reference

Author: rklaehn

Cargo.lock

@@ -45,6 +45,7 @@ reflink-copy = "0.1.24"
 irpc = { version = "0.5.0", features = ["rpc", "quinn_endpoint_setup", "message_spans", "stream", "derive"], default-features = false }
 iroh-metrics = { version = "0.35" }
 entity-manager = { path = "../entity-manager" }
+atomic_refcell = "0.1.13"
 
 [dev-dependencies]
 clap = { version = "4.5.31", features = ["derive"] }

Cargo.toml

@@ -64,7 +64,6 @@
 //! safely shut down as well. Any store refs you are holding will be inoperable
 //! after this.
 use std::{
-    collections::{HashMap, HashSet},
     fmt, fs,
     future::Future,
     io::Write,
@@ -93,7 +92,7 @@ use meta::{list_blobs, Snapshot};
 use n0_future::{future::yield_now, io};
 use nested_enum_utils::enum_conversions;
 use range_collections::range_set::RangeSetRange;
-use tokio::task::{Id, JoinError, JoinSet};
+use tokio::task::{JoinError, JoinSet};
 use tracing::{error, instrument, trace};
 
 use crate::{
@@ -117,7 +116,7 @@ use crate::{
     },
 };
 mod bao_file;
-use bao_file::{BaoFileHandle, BaoFileHandleWeak};
+use bao_file::BaoFileHandle;
 mod delete_set;
 mod entry_state;
 mod import;
@@ -259,7 +258,7 @@ impl HashContext {
         &self.ctx.options
     }
 
-    pub async fn lock(&self) -> tokio::sync::MutexGuard<'_, Option<BaoFileHandleWeak>> {
+    pub async fn lock(&self) -> tokio::sync::MutexGuard<'_, Option<BaoFileHandle>> {
         self.slot.0.lock().await
     }
 
@@ -377,14 +376,9 @@ async fn open_bao_file(
 /// An entry for each hash, containing a weak reference to a BaoFileHandle
 /// wrapped in a tokio mutex so handle creation is sequential.
 #[derive(Debug, Clone, Default)]
-pub(crate) struct Slot(Arc<tokio::sync::Mutex<Option<BaoFileHandleWeak>>>);
+pub(crate) struct Slot(Arc<tokio::sync::Mutex<Option<BaoFileHandle>>>);
 
 impl Slot {
-    pub async fn is_live(&self) -> bool {
-        let slot = self.0.lock().await;
-        slot.as_ref().map(|weak| !weak.is_dead()).unwrap_or(false)
-    }
-
     /// Get the handle if it exists and is still alive, otherwise load it from the database.
     /// If there is nothing in the database, create a new in-memory handle.
     ///
@@ -396,14 +390,12 @@ impl Slot {
         T: Default,
     {
         let mut slot = self.0.lock().await;
-        if let Some(weak) = &*slot {
-            if let Some(handle) = weak.upgrade() {
-                return Ok((handle, Default::default()));
-            }
+        if let Some(handle) = &*slot {
+            return Ok((handle.clone(), Default::default()));
         }
         let handle = make().await;
         if let Ok((handle, _)) = &handle {
-            *slot = Some(handle.downgrade());
+            *slot = Some(handle.clone());
         }
         handle
     }
@@ -845,7 +837,7 @@ async fn finish_import_impl(import_data: ImportEntry, ctx: HashContext) -> io::R
         }
     }
     let guard = ctx.lock().await;
-    let handle = guard.as_ref().and_then(|x| x.upgrade());
+    let handle = guard.as_ref().map(|x| x.clone());
     // if I do have an existing handle, I have to possibly deal with observers.
     // if I don't have an existing handle, there are 2 cases:
     //   the entry exists in the db, but we don't have a handle

src/store/fs.rs

@@ -4,7 +4,7 @@ use std::{
     io,
     ops::Deref,
     path::Path,
-    sync::{Arc, Weak},
+    sync::Arc,
 };
 
 use bao_tree::{
@@ -503,27 +503,6 @@ impl BaoFileStorage {
     }
 }
 
-/// A weak reference to a bao file handle.
-#[derive(Debug, Clone)]
-pub struct BaoFileHandleWeak(Weak<BaoFileHandleInner>);
-
-impl BaoFileHandleWeak {
-    /// Upgrade to a strong reference if possible.
-    pub fn upgrade(&self) -> Option<BaoFileHandle> {
-        let inner = self.0.upgrade()?;
-        if let &BaoFileStorage::Poisoned = inner.storage.borrow().deref() {
-            trace!("poisoned storage, cannot upgrade");
-            return None;
-        };
-        Some(BaoFileHandle(inner))
-    }
-
-    /// True if the handle is definitely dead.
-    pub fn is_dead(&self) -> bool {
-        self.0.strong_count() == 0
-    }
-}
-
 /// The inner part of a bao file handle.
 pub struct BaoFileHandleInner {
     pub(crate) storage: watch::Sender<BaoFileStorage>,
@@ -546,8 +525,8 @@ impl fmt::Debug for BaoFileHandleInner {
 #[derive(Debug, Clone, derive_more::Deref)]
 pub struct BaoFileHandle(Arc<BaoFileHandleInner>);
 
-impl Drop for BaoFileHandle {
-    fn drop(&mut self) {
+impl BaoFileHandle {
+    pub fn persist(&mut self) {
         self.0.storage.send_if_modified(|guard| {
             if Arc::strong_count(&self.0) > 1 {
                 return false;
@@ -564,6 +543,9 @@ impl Drop for BaoFileHandle {
             };
             let options = &self.options;
             let path = options.path.bitfield_path(&self.hash);
+            if fs.bitfield.size == 8000000 {
+                println!("PERSISTING THE TEST CASE FILE to {}", path.display());
+            }
             trace!(
                 "writing bitfield for hash {} to {}",
                 self.hash,
@@ -582,6 +564,12 @@ impl Drop for BaoFileHandle {
     }
 }
 
+impl Drop for BaoFileHandle {
+    fn drop(&mut self) {
+        self.persist();
+    }
+}
+
 /// A reader for a bao file, reading just the data.
 #[derive(Debug)]
 pub struct DataReader(BaoFileHandle);
@@ -753,11 +741,6 @@ impl BaoFileHandle {
         self.hash
     }
 
-    /// Downgrade to a weak reference.
-    pub fn downgrade(&self) -> BaoFileHandleWeak {
-        BaoFileHandleWeak(Arc::downgrade(&self.0))
-    }
-
     /// Write a batch and notify the db
     pub(super) async fn write_batch(
         &self,

src/store/fs/bao_file.rs

@@ -192,6 +192,7 @@ mod tests {
     use std::{
         io::{self},
         path::Path,
+        time::Duration,
     };
 
     use bao_tree::{io::EncodeError, ChunkNum};
@@ -299,10 +300,15 @@ mod tests {
             let outboard_path = options.outboard_path(&bh);
             let sizes_path = options.sizes_path(&bh);
             let bitfield_path = options.bitfield_path(&bh);
+            tokio::time::sleep(Duration::from_millis(100)).await; // allow for some time for the file to be written
             assert!(data_path.exists());
             assert!(outboard_path.exists());
             assert!(sizes_path.exists());
-            assert!(bitfield_path.exists());
+            assert!(
+                bitfield_path.exists(),
+                "Bitfield file {} should exist",
+                bitfield_path.display()
+            );
             gc_run_once(store, &mut live).await?;
             assert!(!data_path.exists());
             assert!(!outboard_path.exists());