From 550bcffc915b66fbb1995a544cc8a4583c6307de Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Mon, 8 May 2023 21:09:50 -0700
Subject: [PATCH] Fix crash during I/O error handling

If an I/O occurring during commit(), the transaction would then try to
abort() which might panic since the data structures are left in an
inconsistent state
---
 src/transactions.rs                       | 6 ++++--
 src/tree_store/page_store/page_manager.rs | 1 +
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/transactions.rs b/src/transactions.rs
index 75348aaf..6d93146d 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -456,6 +456,8 @@ impl<'db> WriteTransaction<'db> {
     /// All writes performed in this transaction will be visible to future transactions, and are
     /// durable as consistent with the [`Durability`] level set by [`Self::set_durability`]
     pub fn commit(mut self) -> Result {
+        // Set completed flag first, so that we don't go through the abort() path on drop, if this fails
+        self.completed = true;
         self.table_tree
             .write()
             .unwrap()
@@ -476,7 +478,6 @@ impl<'db> WriteTransaction<'db> {
             Durability::Paranoid => self.durable_commit(false, true)?,
         }
 
-        self.completed = true;
         #[cfg(feature = "logging")]
         info!(
             "Finished commit of transaction id={:?}",
@@ -490,6 +491,8 @@ impl<'db> WriteTransaction<'db> {
     ///
     /// All writes performed in this transaction will be rolled back
     pub fn abort(mut self) -> Result {
+        // Set completed flag first, so that we don't go through the abort() path on drop, if this fails
+        self.completed = true;
         self.abort_inner()
     }
 
@@ -498,7 +501,6 @@ impl<'db> WriteTransaction<'db> {
         info!("Aborting transaction id={:?}", self.transaction_id);
         self.table_tree.write().unwrap().clear_table_root_updates();
         self.mem.rollback_uncommitted_writes()?;
-        self.completed = true;
         #[cfg(feature = "logging")]
         info!("Finished abort of transaction id={:?}", self.transaction_id);
         Ok(())
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index b2c5fb25..cffd5d1f 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -887,6 +887,7 @@ impl TransactionalMemory {
                 "Dirty pages outstanding: {dirty_pages:?}"
             );
         }
+        assert!(!self.needs_recovery.load(Ordering::Acquire));
         let mut state = self.state.lock().unwrap();
         // The layout to restore
         let (restore, restore_tracker_page) = if self.read_from_secondary.load(Ordering::Acquire) {