From 929b42b7adacc05fa5ef6bc61f37443b495fa927 Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Sat, 20 Jan 2024 08:37:33 -0800
Subject: [PATCH] Add close() method to ReadTransaction

Calling this method is entirely optional. Transactions still close
automatically, but this can be used to ensure that it can be closed
immediately
---
 src/error.rs        | 16 +++++++++++++++-
 src/transactions.rs | 24 +++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index c475f5a3..c9735790 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,5 +1,5 @@
 use crate::tree_store::{FILE_FORMAT_VERSION, MAX_VALUE_LENGTH};
-use crate::TypeName;
+use crate::{ReadTransaction, TypeName};
 use std::fmt::{Display, Formatter};
 use std::sync::PoisonError;
 use std::{io, panic};
@@ -334,12 +334,15 @@ impl std::error::Error for CompactionError {}
 pub enum TransactionError {
     /// Error from underlying storage
     Storage(StorageError),
+    /// The transaction is still referenced by a table or other object
+    ReadTransactionStillInUse(ReadTransaction<'static>),
 }
 
 impl TransactionError {
     pub(crate) fn into_storage_error(self) -> StorageError {
         match self {
             TransactionError::Storage(storage) => storage,
+            _ => unreachable!(),
         }
     }
 }
@@ -348,6 +351,9 @@ impl From<TransactionError> for Error {
     fn from(err: TransactionError) -> Error {
         match err {
             TransactionError::Storage(storage) => storage.into(),
+            TransactionError::ReadTransactionStillInUse(txn) => {
+                Error::ReadTransactionStillInUse(txn)
+            }
         }
     }
 }
@@ -362,6 +368,9 @@ impl Display for TransactionError {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             TransactionError::Storage(storage) => storage.fmt(f),
+            TransactionError::ReadTransactionStillInUse(_) => {
+                write!(f, "Transaction still in use")
+            }
         }
     }
 }
@@ -453,6 +462,8 @@ pub enum Error {
     TableAlreadyOpen(String, &'static panic::Location<'static>),
     Io(io::Error),
     LockPoisoned(&'static panic::Location<'static>),
+    /// The transaction is still referenced by a table or other object
+    ReadTransactionStillInUse(ReadTransaction<'static>),
 }
 
 impl<T> From<PoisonError<T>> for Error {
@@ -543,6 +554,9 @@ impl Display for Error {
             Error::InvalidSavepoint => {
                 write!(f, "Savepoint is invalid or cannot be created.")
             }
+            Error::ReadTransactionStillInUse(_) => {
+                write!(f, "Transaction still in use")
+            }
         }
     }
 }
diff --git a/src/transactions.rs b/src/transactions.rs
index 568f2e0f..29c2aef7 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -21,7 +21,7 @@ use log::{info, warn};
 use std::borrow::Borrow;
 use std::cmp::min;
 use std::collections::{HashMap, HashSet};
-use std::fmt::{Display, Formatter};
+use std::fmt::{Debug, Display, Formatter};
 use std::marker::PhantomData;
 use std::ops::{RangeBounds, RangeFull};
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -1276,6 +1276,28 @@ impl<'a> ReadTransaction<'a> {
     }
 }
 
+impl ReadTransaction<'static> {
+    /// Close the transaction
+    ///
+    /// Transactions are automatically closed when they and all objects referencing them have been dropped.
+    /// This method can be used to ensure that there are no outstanding objects remaining.
+    ///
+    /// Returns `ReadTransactionStillInUse` error if a table or other object retrieved from the transaction still references this transaction
+    pub fn close(self) -> Result<(), TransactionError> {
+        if Arc::strong_count(&self.transaction_guard) > 1 {
+            return Err(TransactionError::ReadTransactionStillInUse(self));
+        }
+        // No-op, just drop ourself
+        Ok(())
+    }
+}
+
+impl Debug for ReadTransaction<'static> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_str("ReadTransaction")
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::{Database, TableDefinition};