StorageAutoPtr

voltrevo · Nov 1, 2023 · 350fa65 · 350fa65
1 parent e454708
commit 350fa65
Show file tree

Hide file tree

Showing 9 changed files with 148 additions and 43 deletions.
diff --git a/storage/src/demo_val.rs b/storage/src/demo_val.rs
@@ -25,7 +25,10 @@ impl DemoVal {
     &self,
     storage: &mut Storage<SB>,
   ) -> Result<Vec<u64>, Box<dyn Error>> {
-    storage.sb.transaction(|sb| self.numbers_impl(sb))
+    storage
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&storage.sb), |sb| self.numbers_impl(sb))
   }
 
   fn write_to_entry<'a, SB: StorageBackend, Tx: StorageTx<'a, SB>>(

diff --git a/storage/src/lib.rs b/storage/src/lib.rs
@@ -6,6 +6,7 @@ mod demo_val;
 
 mod rc_key;
 mod sled_backend;
+mod storage_auto_ptr;
 mod storage_backend;
 mod storage_entity;
 mod storage_entry;
@@ -19,6 +20,7 @@ pub use self::storage_tx::StorageTx;
 pub use memory_backend::MemoryBackend;
 pub use rc_key::RcKey;
 pub use sled_backend::SledBackend;
+pub use storage_auto_ptr::StorageAutoPtr;
 pub use storage_entity::StorageEntity;
 pub use storage_entry::{StorageEntry, StorageEntryReader, StorageEntryWriter};
 pub use storage_ptr::{storage_head_ptr, StorageEntryPtr, StorageHeadPtr, StoragePtr};
diff --git a/storage/src/memory_backend.rs b/storage/src/memory_backend.rs
@@ -1,8 +1,8 @@
-use std::{collections::HashMap, error::Error};
+use std::{cell::RefCell, collections::HashMap, error::Error, rc::Weak};
 
 use crate::{
   rc_key::RcKey, storage_backend::StorageError, storage_ptr::StorageEntryPtr,
-  storage_tx::StorageTx, StorageBackend, StoragePtr,
+  storage_tx::StorageTx, StorageAutoPtr, StorageBackend, StorageEntity, StoragePtr,
 };
 
 #[derive(Default)]
@@ -22,11 +22,12 @@ impl StorageBackend for MemoryBackend {
   type CustomError = Box<dyn Error>;
   type Tx<'a> = MemoryTx<'a>;
 
-  fn transaction<F, T>(&mut self, f: F) -> Result<T, Box<dyn Error>>
+  fn transaction<F, T>(&mut self, self_weak: Weak<RefCell<Self>>, f: F) -> Result<T, Box<dyn Error>>
   where
     F: Fn(&mut Self::Tx<'_>) -> Result<T, StorageError<Self>>,
   {
     let mut handle = MemoryTx {
+      backend: self_weak,
       ref_deltas: Default::default(),
       cache: Default::default(),
       storage: self,
@@ -56,6 +57,7 @@ impl StorageBackend for MemoryBackend {
 }
 
 pub struct MemoryTx<'a> {
+  backend: Weak<RefCell<MemoryBackend>>,
   ref_deltas: HashMap<(u64, u64, u64), i64>,
   cache: HashMap<RcKey, StorageEntryPtr>,
   storage: &'a mut MemoryBackend,
@@ -77,6 +79,17 @@ impl<'a> StorageTx<'a, MemoryBackend> for MemoryTx<'a> {
     Ok(self.storage.data.get(&ptr.data).cloned())
   }
 
+  fn get_auto_ptr<SE: for<'b> StorageEntity<'b, MemoryBackend, MemoryTx<'b>>>(
+    &mut self,
+    ptr: StorageEntryPtr,
+  ) -> StorageAutoPtr<MemoryBackend, SE> {
+    StorageAutoPtr {
+      _marker: std::marker::PhantomData,
+      sb: self.backend.clone(),
+      ptr,
+    }
+  }
+
   fn write_bytes<T>(
     &mut self,
     ptr: StoragePtr<T>,

diff --git a/storage/src/sled_backend.rs b/storage/src/sled_backend.rs
@@ -1,8 +1,8 @@
-use std::{collections::HashMap, error::Error};
+use std::{cell::RefCell, collections::HashMap, error::Error, rc::Weak};
 
 use crate::{
   rc_key::RcKey, storage_backend::StorageError, storage_ptr::StorageEntryPtr,
-  storage_tx::StorageTx, StorageBackend, StoragePtr,
+  storage_tx::StorageTx, StorageAutoPtr, StorageBackend, StoragePtr,
 };
 
 pub struct SledBackend {
@@ -30,14 +30,15 @@ impl StorageBackend for SledBackend {
   type CustomError = sled::transaction::ConflictableTransactionError<Box<dyn Error>>;
   type Tx<'a> = SledTx<'a>;
 
-  fn transaction<F, T>(&mut self, f: F) -> Result<T, Box<dyn Error>>
+  fn transaction<F, T>(&mut self, self_weak: Weak<RefCell<Self>>, f: F) -> Result<T, Box<dyn Error>>
   where
     F: Fn(&mut Self::Tx<'_>) -> Result<T, StorageError<Self>>,
   {
     self
       .db
       .transaction(|tx| {
         let mut handle = SledTx {
+          backend: self_weak.clone(),
           ref_deltas: Default::default(),
           cache: Default::default(),
           tx,
@@ -72,6 +73,7 @@ impl StorageBackend for SledBackend {
 }
 
 pub struct SledTx<'a> {
+  backend: Weak<RefCell<SledBackend>>,
   ref_deltas: HashMap<(u64, u64, u64), i64>,
   cache: HashMap<RcKey, StorageEntryPtr>,
   tx: &'a sled::transaction::TransactionalTree,
@@ -99,6 +101,19 @@ impl<'a> StorageTx<'a, SledBackend> for SledTx<'a> {
     Ok(value)
   }
 
+  fn get_auto_ptr<
+    SE: for<'b> crate::StorageEntity<'b, SledBackend, <SledBackend as StorageBackend>::Tx<'b>>,
+  >(
+    &mut self,
+    ptr: StorageEntryPtr,
+  ) -> crate::StorageAutoPtr<SledBackend, SE> {
+    StorageAutoPtr {
+      _marker: std::marker::PhantomData,
+      sb: self.backend.clone(),
+      ptr,
+    }
+  }
+
   fn write_bytes<T>(
     &mut self,
     ptr: StoragePtr<T>,

diff --git a/storage/src/storage.rs b/storage/src/storage.rs
@@ -1,85 +1,118 @@
+use std::cell::RefCell;
 use std::error::Error;
+use std::rc::Rc;
 
+use crate::storage_auto_ptr::StorageAutoPtr;
 use crate::storage_entity::StorageEntity;
 use crate::storage_ptr::{tmp_at_ptr, tmp_count_ptr, StorageEntryPtr, StorageHeadPtr};
-use crate::storage_tx::StorageTx;
-use crate::{StorageBackend, StorageError};
+use crate::{StorageBackend, StorageError, StorageTx};
 
 pub struct Storage<SB: StorageBackend> {
-  pub(crate) sb: SB,
+  pub(crate) sb: Rc<RefCell<SB>>,
 }
 
 impl<SB: StorageBackend> Storage<SB> {
   pub fn new(sb: SB) -> Self {
-    Self { sb }
+    Self {
+      sb: Rc::new(RefCell::new(sb)),
+    }
   }
 
   pub fn get_head<SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>>(
     &mut self,
     ptr: StorageHeadPtr,
   ) -> Result<Option<SE>, Box<dyn Error>> {
-    self.sb.transaction(|sb| sb.get_head(ptr))
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| sb.get_head(ptr))
   }
 
   pub fn get<SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>>(
     &mut self,
     ptr: StorageEntryPtr,
   ) -> Result<SE, Box<dyn Error>> {
     // TODO: Avoid going through a transaction when read-only
-    self.sb.transaction(|sb| {
-      let entry = sb
-        .read(ptr)?
-        .ok_or(StorageError::Error("Ptr not found".into()))?;
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| {
+        let entry = sb
+          .read(ptr)?
+          .ok_or(StorageError::Error("Ptr not found".into()))?;
+
+        SE::from_storage_entry(sb, entry)
+      })
+  }
 
-      SE::from_storage_entry(sb, entry)
-    })
+  pub fn get_auto_ptr<SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>>(
+    &mut self,
+    ptr: StorageEntryPtr,
+  ) -> StorageAutoPtr<SB, SE> {
+    StorageAutoPtr {
+      _marker: std::marker::PhantomData,
+      sb: Rc::downgrade(&self.sb),
+      ptr,
+    }
   }
 
   pub fn set_head<SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>>(
     &mut self,
     ptr: StorageHeadPtr,
     value: &SE,
   ) -> Result<(), Box<dyn Error>> {
-    self.sb.transaction(|sb| sb.set_head(ptr, value))
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| sb.set_head(ptr, value))
   }
 
   pub fn remove_head(&mut self, ptr: StorageHeadPtr) -> Result<(), Box<dyn Error>> {
-    self.sb.transaction(|sb| sb.remove_head(ptr))
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| sb.remove_head(ptr))
   }
 
   pub fn store_tmp<SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>>(
     &mut self,
     value: &SE,
   ) -> Result<StorageEntryPtr, Box<dyn Error>> {
-    self.sb.transaction(|sb| {
-      let tmp_count = sb.read(tmp_count_ptr())?.unwrap_or(0);
-      let tmp_ptr = tmp_at_ptr(tmp_count);
-      sb.set_head(tmp_ptr, value)?;
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| {
+        let tmp_count = sb.read(tmp_count_ptr())?.unwrap_or(0);
+        let tmp_ptr = tmp_at_ptr(tmp_count);
+        sb.set_head(tmp_ptr, value)?;
 
-      sb.write(tmp_count_ptr(), Some(&(tmp_count + 1)))?;
+        sb.write(tmp_count_ptr(), Some(&(tmp_count + 1)))?;
 
-      let ptr = sb.read(tmp_ptr)?.unwrap_or_else(|| panic!("Ptr not found"));
+        let ptr = sb.read(tmp_ptr)?.unwrap_or_else(|| panic!("Ptr not found"));
 
-      Ok(ptr)
-    })
+        Ok(ptr)
+      })
   }
 
   pub fn clear_tmp(&mut self) -> Result<(), Box<dyn Error>> {
-    self.sb.transaction(|sb| {
-      let tmp_count = sb.read(tmp_count_ptr())?.unwrap_or(0);
+    self
+      .sb
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| {
+        let tmp_count = sb.read(tmp_count_ptr())?.unwrap_or(0);
 
-      for i in 0..tmp_count {
-        sb.remove_head(tmp_at_ptr(i))?;
-      }
+        for i in 0..tmp_count {
+          sb.remove_head(tmp_at_ptr(i))?;
+        }
 
-      sb.write(tmp_count_ptr(), None)?;
+        sb.write(tmp_count_ptr(), None)?;
 
-      Ok(())
-    })
+        Ok(())
+      })
   }
 
   pub fn is_empty(&self) -> bool {
-    self.sb.is_empty()
+    self.sb.borrow().is_empty()
   }
 
   #[cfg(test)]
@@ -89,6 +122,9 @@ impl<SB: StorageBackend> Storage<SB> {
   ) -> Result<Option<u64>, Box<dyn Error>> {
     self
       .sb
-      .transaction(|sb| Ok(sb.read(ptr)?.map(|entry| entry.ref_count)))
+      .borrow_mut()
+      .transaction(Rc::downgrade(&self.sb), |sb| {
+        Ok(sb.read(ptr)?.map(|entry| entry.ref_count))
+      })
   }
 }
diff --git a/storage/src/storage_auto_ptr.rs b/storage/src/storage_auto_ptr.rs
@@ -0,0 +1,27 @@
+use std::{cell::RefCell, error::Error, rc::Weak};
+
+use crate::{StorageBackend, StorageEntity, StorageEntryPtr, StorageTx};
+
+pub struct StorageAutoPtr<SB: StorageBackend, SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>> {
+  pub(crate) _marker: std::marker::PhantomData<SE>,
+  pub(crate) sb: Weak<RefCell<SB>>, // TODO: Does this need to be weak?
+  pub(crate) ptr: StorageEntryPtr,
+}
+
+impl<SB: StorageBackend, SE: for<'a> StorageEntity<'a, SB, SB::Tx<'a>>> StorageAutoPtr<SB, SE> {
+  pub fn resolve(&self) -> Result<Option<SE>, Box<dyn Error>> {
+    let sb = match self.sb.upgrade() {
+      Some(sb) => sb,
+      None => return Err("Storage backend dropped".into()),
+    };
+
+    let res = sb.borrow_mut().transaction(self.sb.clone(), |sb| {
+      Ok(match sb.read(self.ptr)? {
+        Some(entry) => Some(SE::from_storage_entry(sb, entry)?),
+        None => None,
+      })
+    });
+
+    res
+  }
+}
diff --git a/storage/src/storage_backend.rs b/storage/src/storage_backend.rs
@@ -1,12 +1,16 @@
-use std::error::Error;
+use std::{cell::RefCell, error::Error, rc::Weak};
 
 use crate::storage_tx::StorageTx;
 
 pub trait StorageBackend: Sized {
   type CustomError;
   type Tx<'a>: StorageTx<'a, Self>;
 
-  fn transaction<F, T>(&mut self, f: F) -> Result<T, Box<dyn Error>>
+  fn transaction<F, T>(
+    &mut self,
+    self_weak: Weak<RefCell<Self>>,
+    f: F,
+  ) -> Result<T, Box<dyn Error>>
   where
     F: Fn(&mut Self::Tx<'_>) -> Result<T, StorageError<Self>>;
 

diff --git a/storage/src/storage_tx.rs b/storage/src/storage_tx.rs
@@ -4,8 +4,8 @@ use rand::thread_rng;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-  storage_backend::StorageError, RcKey, StorageBackend, StorageEntity, StorageEntryPtr,
-  StorageHeadPtr, StoragePtr,
+  storage_backend::StorageError, RcKey, StorageAutoPtr, StorageBackend, StorageEntity,
+  StorageEntryPtr, StorageHeadPtr, StoragePtr,
 };
 
 pub trait StorageTx<'a, SB: StorageBackend>: Sized {
@@ -32,6 +32,11 @@ pub trait StorageTx<'a, SB: StorageBackend>: Sized {
       .map_err(StorageError::from)
   }
 
+  fn get_auto_ptr<SE: for<'b> StorageEntity<'b, SB, SB::Tx<'b>>>(
+    &mut self,
+    ptr: StorageEntryPtr,
+  ) -> StorageAutoPtr<SB, SE>;
+
   fn read_or_err<T: for<'de> Deserialize<'de>>(
     &mut self,
     ptr: StoragePtr<T>,

diff --git a/storage/src/tests.rs b/storage/src/tests.rs
@@ -120,7 +120,7 @@ mod tests_ {
 
       storage.remove_head(storage_head_ptr(b"test")).unwrap();
 
-      assert_eq!(storage.sb.len(), 0);
+      assert_eq!(storage.sb.borrow().len(), 0);
       assert!(storage.is_empty());
     }