introduce TxnGuard trait to pass either TxnGuardOwned/TxnGuardShared

Author: MarinPostma

libsql-wal/src/registry.rs

@@ -27,7 +27,7 @@ use crate::segment::Segment;
 use crate::segment::{current::CurrentSegment, sealed::SealedSegment};
 use crate::shared_wal::{SharedWal, SwapLog};
 use crate::storage::{OnStoreCallback, Storage};
-use crate::transaction::{Transaction, TxGuard};
+use crate::transaction::TxGuard;
 use crate::{LibsqlFooter, LIBSQL_PAGE_SIZE};
 use libsql_sys::name::NamespaceName;
 
@@ -100,48 +100,13 @@ where
     S: Storage<Segment = SealedSegment<IO::File>>,
 {
     #[tracing::instrument(skip_all)]
-    fn swap_current(&self, shared: &SharedWal<IO>, tx: &TxGuard<<IO as Io>::File>) -> Result<()> {
+    fn swap_current(
+        &self,
+        shared: &SharedWal<IO>,
+        tx: &dyn TxGuard<<IO as Io>::File>,
+    ) -> Result<()> {
         assert!(tx.is_commited());
-        // at this point we must hold a lock to a commited transaction.
-
-        let current = shared.current.load();
-        if current.is_empty() {
-            return Ok(());
-        }
-        let start_frame_no = current.next_frame_no();
-        let path = self
-            .path
-            .join(shared.namespace().as_str())
-            .join(format!("{}:{start_frame_no:020}.seg", shared.namespace()));
-
-        let segment_file = self.io.open(true, true, true, &path)?;
-        let salt = self.io.with_rng(|rng| rng.gen());
-        let new = CurrentSegment::create(
-            segment_file,
-            path,
-            start_frame_no,
-            current.db_size(),
-            current.tail().clone(),
-            salt,
-            current.log_id(),
-        )?;
-        // sealing must the last fallible operation, because we don't want to end up in a situation
-        // where the current log is sealed and it wasn't swapped.
-        if let Some(sealed) = current.seal()? {
-            new.tail().push(sealed.clone());
-            maybe_store_segment(
-                self.storage.as_ref(),
-                &self.checkpoint_notifier,
-                &shared.namespace,
-                &shared.durable_frame_no,
-                sealed,
-            );
-        }
-
-        shared.current.swap(Arc::new(new));
-        tracing::debug!("current segment swapped");
-
-        Ok(())
+        self.swap_current_inner(shared)
     }
 }
 
@@ -498,6 +463,48 @@ where
 
         Ok(())
     }
+
+    #[tracing::instrument(skip_all)]
+    fn swap_current_inner(&self, shared: &SharedWal<IO>) -> Result<()> {
+        let current = shared.current.load();
+        if current.is_empty() {
+            return Ok(());
+        }
+        let start_frame_no = current.next_frame_no();
+        let path = self
+            .path
+            .join(shared.namespace().as_str())
+            .join(format!("{}:{start_frame_no:020}.seg", shared.namespace()));
+
+        let segment_file = self.io.open(true, true, true, &path)?;
+        let salt = self.io.with_rng(|rng| rng.gen());
+        let new = CurrentSegment::create(
+            segment_file,
+            path,
+            start_frame_no,
+            current.db_size(),
+            current.tail().clone(),
+            salt,
+            current.log_id(),
+        )?;
+        // sealing must the last fallible operation, because we don't want to end up in a situation
+        // where the current log is sealed and it wasn't swapped.
+        if let Some(sealed) = current.seal()? {
+            new.tail().push(sealed.clone());
+            maybe_store_segment(
+                self.storage.as_ref(),
+                &self.checkpoint_notifier,
+                &shared.namespace,
+                &shared.durable_frame_no,
+                sealed,
+            );
+        }
+
+        shared.current.swap(Arc::new(new));
+        tracing::debug!("current segment swapped");
+
+        Ok(())
+    }
 }
 
 #[tracing::instrument(skip_all, fields(namespace = shared.namespace().as_str()))]

libsql-wal/src/shared_wal.rs

@@ -35,7 +35,7 @@ pub struct WalLock {
 }
 
 pub(crate) trait SwapLog<IO: Io>: Sync + Send + 'static {
-    fn swap_current(&self, shared: &SharedWal<IO>, tx: &TxGuard<IO::File>) -> Result<()>;
+    fn swap_current(&self, shared: &SharedWal<IO>, tx: &dyn TxGuard<IO::File>) -> Result<()>;
 }
 
 pub struct SharedWal<IO: Io> {
@@ -300,7 +300,7 @@ impl<IO: Io> SharedWal<IO> {
     }
 
     /// Swap the current log. A write lock must be held, but the transaction must be must be committed already.
-    pub(crate) fn swap_current(&self, tx: &TxGuard<IO::File>) -> Result<()> {
+    pub(crate) fn swap_current(&self, tx: &impl TxGuard<IO::File>) -> Result<()> {
         self.registry.swap_current(self, tx)?;
         Ok(())
     }

libsql-wal/src/transaction.rs

@@ -189,20 +189,25 @@ impl<F> DerefMut for TxGuardOwned<F> {
     }
 }
 
-pub struct TxGuard<'a, F> {
+pub trait TxGuard<F>: Deref<Target = WriteTransaction<F>> + DerefMut + Send + Sync { }
+
+impl<'a, F: Send + Sync> TxGuard<F> for TxGuardShared<'a, F> { }
+impl<F: Send + Sync> TxGuard<F> for TxGuardOwned<F> { }
+
+pub struct TxGuardShared<'a, F> {
     _lock: async_lock::MutexGuardArc<Option<u64>>,
     inner: &'a mut WriteTransaction<F>,
 }
 
-impl<'a, F> Deref for TxGuard<'a, F> {
+impl<'a, F> Deref for TxGuardShared<'a, F> {
     type Target = WriteTransaction<F>;
 
     fn deref(&self) -> &Self::Target {
         &self.inner
     }
 }
 
-impl<'a, F> DerefMut for TxGuard<'a, F> {
+impl<'a, F> DerefMut for TxGuardShared<'a, F> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.inner
     }
@@ -225,11 +230,11 @@ impl<F> WriteTransaction<F> {
         savepoint_id
     }
 
-    pub fn lock(&mut self) -> TxGuard<F> {
+    pub fn lock(&mut self) -> TxGuardShared<F> {
         let g = self.wal_lock.tx_id.lock_arc_blocking();
         match *g {
             // we still hold the lock, we can proceed
-            Some(id) if self.id == id => TxGuard {
+            Some(id) if self.id == id => TxGuardShared {
                 _lock: g,
                 inner: self,
             },