Use a single txn for an entire coordinator message

Removes direct DB accesses whre possible. Documents the safety of the rest.
Does uncover one case of unsafety not previously noted.

processor/src/tests/key_gen.rs

@@ -7,7 +7,7 @@ use rand_core::{RngCore, OsRng};
 use group::GroupEncoding;
 use frost::{Participant, ThresholdParams, tests::clone_without};
 
-use serai_db::MemDb;
+use serai_db::{DbTxn, Db, MemDb};
 
 use serai_client::{
   primitives::{MONERO_NET_ID, BlockHash},
@@ -31,19 +31,24 @@ pub async fn test_key_gen<C: Coin>() {
     let mut entropy = Zeroizing::new([0; 32]);
     OsRng.fill_bytes(entropy.as_mut());
     entropies.insert(i, entropy);
-    dbs.insert(i, MemDb::new());
-    key_gens.insert(i, KeyGen::<C, _>::new(dbs[&i].clone(), entropies[&i].clone()));
+    let db = MemDb::new();
+    dbs.insert(i, db.clone());
+    key_gens.insert(i, KeyGen::<C, MemDb>::new(db, entropies[&i].clone()));
   }
 
   let mut all_commitments = HashMap::new();
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
+    let mut txn = dbs.get_mut(&i).unwrap().txn();
     if let ProcessorMessage::Commitments { id, commitments } = key_gen
-      .handle(CoordinatorMessage::GenerateKey {
-        id: ID,
-        params: ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap())
-          .unwrap(),
-      })
+      .handle(
+        &mut txn,
+        CoordinatorMessage::GenerateKey {
+          id: ID,
+          params: ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap())
+            .unwrap(),
+        },
+      )
       .await
     {
       assert_eq!(id, ID);
@@ -51,27 +56,32 @@ pub async fn test_key_gen<C: Coin>() {
     } else {
       panic!("didn't get commitments back");
     }
+    txn.commit();
   }
 
   // 1 is rebuilt on every step
   // 2 is rebuilt here
   // 3 ... are rebuilt once, one at each of the following steps
-  let rebuild = |key_gens: &mut HashMap<_, _>, i| {
+  let rebuild = |key_gens: &mut HashMap<_, _>, dbs: &HashMap<_, MemDb>, i| {
     key_gens.remove(&i);
     key_gens.insert(i, KeyGen::<C, _>::new(dbs[&i].clone(), entropies[&i].clone()));
   };
-  rebuild(&mut key_gens, 1);
-  rebuild(&mut key_gens, 2);
+  rebuild(&mut key_gens, &dbs, 1);
+  rebuild(&mut key_gens, &dbs, 2);
 
   let mut all_shares = HashMap::new();
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
+    let mut txn = dbs.get_mut(&i).unwrap().txn();
     let i = Participant::new(u16::try_from(i).unwrap()).unwrap();
     if let ProcessorMessage::Shares { id, shares } = key_gen
-      .handle(CoordinatorMessage::Commitments {
-        id: ID,
-        commitments: clone_without(&all_commitments, &i),
-      })
+      .handle(
+        &mut txn,
+        CoordinatorMessage::Commitments {
+          id: ID,
+          commitments: clone_without(&all_commitments, &i),
+        },
+      )
       .await
     {
       assert_eq!(id, ID);
@@ -79,24 +89,29 @@ pub async fn test_key_gen<C: Coin>() {
     } else {
       panic!("didn't get shares back");
     }
+    txn.commit();
   }
 
   // Rebuild 1 and 3
-  rebuild(&mut key_gens, 1);
-  rebuild(&mut key_gens, 3);
+  rebuild(&mut key_gens, &dbs, 1);
+  rebuild(&mut key_gens, &dbs, 3);
 
   let mut res = None;
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
+    let mut txn = dbs.get_mut(&i).unwrap().txn();
     let i = Participant::new(u16::try_from(i).unwrap()).unwrap();
     if let ProcessorMessage::GeneratedKeyPair { id, substrate_key, coin_key } = key_gen
-      .handle(CoordinatorMessage::Shares {
-        id: ID,
-        shares: all_shares
-          .iter()
-          .filter_map(|(l, shares)| if i == *l { None } else { Some((*l, shares[&i].clone())) })
-          .collect(),
-      })
+      .handle(
+        &mut txn,
+        CoordinatorMessage::Shares {
+          id: ID,
+          shares: all_shares
+            .iter()
+            .filter_map(|(l, shares)| if i == *l { None } else { Some((*l, shares[&i].clone())) })
+            .collect(),
+        },
+      )
       .await
     {
       assert_eq!(id, ID);
@@ -107,17 +122,25 @@ pub async fn test_key_gen<C: Coin>() {
     } else {
       panic!("didn't get key back");
     }
+    txn.commit();
   }
 
   // Rebuild 1 and 4
-  rebuild(&mut key_gens, 1);
-  rebuild(&mut key_gens, 4);
+  rebuild(&mut key_gens, &dbs, 1);
+  rebuild(&mut key_gens, &dbs, 4);
 
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
+    let mut txn = dbs.get_mut(&i).unwrap().txn();
     let KeyConfirmed { activation_block, substrate_keys, coin_keys } = key_gen
-      .confirm(SubstrateContext { coin_latest_finalized_block: BlockHash([0x11; 32]) }, ID)
+      .confirm(
+        &mut txn,
+        SubstrateContext { coin_latest_finalized_block: BlockHash([0x11; 32]) },
+        ID,
+      )
       .await;
+    txn.commit();
+
     assert_eq!(activation_block, BlockHash([0x11; 32]));
     let params =
       ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap()).unwrap();