Make transactions undroppable

coordinator/cosign/src/delay.rs literally demonstrates how we'd need to rewrite
our handling of transactions with this change. It can be cleaned up a bit but
already identifies ergonomic issues. It also doesn't model passing an &mut txn
to an async function, which would also require using the droppable wrapper
struct.

To locally see this build, run

RUSTFLAGS="-Zpanic_abort_tests -C panic=abort" cargo +nightly build -p serai-cosign --all-targets

To locally see this fail to build, run

cargo build -p serai-cosign --all-targets

While it doesn't say which line causes it fail to build, the only distinction
is panic=unwind.

For more context, please see #578.

common/db/src/lib.rs

@@ -30,13 +30,53 @@ pub trait Get {
 /// is undefined. The transaction may block, deadlock, panic, overwrite one of the two values
 /// randomly, or any other action, at time of write or at time of commit.
 #[must_use]
-pub trait DbTxn: Send + Get {
+pub trait DbTxn: Sized + Send + Get {
   /// Write a value to this key.
   fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>);
   /// Delete the value from this key.
   fn del(&mut self, key: impl AsRef<[u8]>);
   /// Commit this transaction.
   fn commit(self);
+  /// Close this transaction.
+  ///
+  /// This is equivalent to `Drop` on transactions which can be dropped. This is explicit and works
+  /// with transactions which can't be dropped.
+  fn close(self) {
+    drop(self);
+  }
+}
+
+// Credit for the idea goes to https://jack.wrenn.fyi/blog/undroppable
+pub struct Undroppable<T>(Option<T>);
+impl<T> Drop for Undroppable<T> {
+  fn drop(&mut self) {
+    // Use an assertion at compile time to prevent this code from compiling if generated
+    #[allow(clippy::assertions_on_constants)]
+    const {
+      assert!(false, "Undroppable DbTxn was dropped. Ensure all code paths call commit or close");
+    }
+  }
+}
+impl<T: DbTxn> Get for Undroppable<T> {
+  fn get(&self, key: impl AsRef<[u8]>) -> Option<Vec<u8>> {
+    self.0.as_ref().unwrap().get(key)
+  }
+}
+impl<T: DbTxn> DbTxn for Undroppable<T> {
+  fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) {
+    self.0.as_mut().unwrap().put(key, value);
+  }
+  fn del(&mut self, key: impl AsRef<[u8]>) {
+    self.0.as_mut().unwrap().del(key);
+  }
+  fn commit(mut self) {
+    self.0.take().unwrap().commit();
+    let _ = core::mem::ManuallyDrop::new(self);
+  }
+  fn close(mut self) {
+    drop(self.0.take().unwrap());
+    let _ = core::mem::ManuallyDrop::new(self);
+  }
 }
 
 /// A database supporting atomic transaction.
@@ -51,6 +91,10 @@ pub trait Db: 'static + Send + Sync + Clone + Get {
     let dst_len = u8::try_from(item_dst.len()).unwrap();
     [[db_len].as_ref(), db_dst, [dst_len].as_ref(), item_dst, key.as_ref()].concat()
   }
-  /// Open a new transaction.
-  fn txn(&mut self) -> Self::Transaction<'_>;
+  /// Open a new transaction which may be dropped.
+  fn unsafe_txn(&mut self) -> Self::Transaction<'_>;
+  /// Open a new transaction which must be committed or closed.
+  fn txn(&mut self) -> Undroppable<Self::Transaction<'_>> {
+    Undroppable(Some(self.unsafe_txn()))
+  }
 }

common/db/src/mem.rs

@@ -74,7 +74,7 @@ impl Get for MemDb {
 }
 impl Db for MemDb {
   type Transaction<'a> = MemDbTxn<'a>;
-  fn txn(&mut self) -> MemDbTxn<'_> {
+  fn unsafe_txn(&mut self) -> MemDbTxn<'_> {
     MemDbTxn(self, HashMap::new(), HashSet::new())
   }
 }

common/db/src/parity_db.rs

@@ -37,7 +37,7 @@ impl Get for Arc<ParityDb> {
 }
 impl Db for Arc<ParityDb> {
   type Transaction<'a> = Transaction<'a>;
-  fn txn(&mut self) -> Self::Transaction<'_> {
+  fn unsafe_txn(&mut self) -> Self::Transaction<'_> {
     Transaction(self, vec![])
   }
 }

common/db/src/rocks.rs

@@ -39,7 +39,7 @@ impl<T: ThreadMode> Get for Arc<OptimisticTransactionDB<T>> {
 }
 impl<T: Send + ThreadMode + 'static> Db for Arc<OptimisticTransactionDB<T>> {
   type Transaction<'a> = Transaction<'a, T>;
-  fn txn(&mut self) -> Self::Transaction<'_> {
+  fn unsafe_txn(&mut self) -> Self::Transaction<'_> {
     let mut opts = WriteOptions::default();
     opts.set_sync(true);
     Transaction(self.transaction_opt(&opts, &Default::default()), &**self)