Cleanup DB handling a bit in key-gen/attempt-manager

Cargo.lock

@@ -8603,9 +8603,11 @@ dependencies = [
 name = "serai-processor-frost-attempt-manager"
 version = "0.1.0"
 dependencies = [
+ "borsh",
  "hex",
  "log",
  "modular-frost",
+ "parity-scale-codec",
  "rand_core",
  "serai-db",
  "serai-processor-messages",

processor/frost-attempt-manager/Cargo.toml

@@ -25,5 +25,9 @@ serai-validator-sets-primitives = { path = "../../substrate/validator-sets/primi
 
 hex = { version = "0.4", default-features = false, features = ["std"] }
 log = { version = "0.4", default-features = false, features = ["std"] }
+
+scale = { package = "parity-scale-codec", version = "3", default-features = false, features = ["std"] }
+borsh = { version = "1", default-features = false, features = ["std", "derive", "de_strict_order"] }
 serai-db = { path = "../../common/db" }
+
 messages = { package = "serai-processor-messages", path = "../messages" }

processor/frost-attempt-manager/src/individual.rs

@@ -9,11 +9,19 @@ use frost::{
 
 use serai_validator_sets_primitives::Session;
 
+use serai_db::{Get, DbTxn, Db, create_db};
 use messages::sign::{SignId, ProcessorMessage};
 
+create_db!(
+  FrostAttemptManager {
+    Attempted: (id: [u8; 32]) -> u32,
+  }
+);
+
 /// An instance of a signing protocol with re-attempts handled internally.
 #[allow(clippy::type_complexity)]
-pub(crate) struct SigningProtocol<M: Clone + PreprocessMachine> {
+pub(crate) struct SigningProtocol<D: Db, M: Clone + PreprocessMachine> {
+  db: D,
   // The session this signing protocol is being conducted by.
   session: Session,
   // The `i` of our first, or starting, set of key shares we will be signing with.
@@ -34,12 +42,19 @@ pub(crate) struct SigningProtocol<M: Clone + PreprocessMachine> {
   >,
 }
 
-impl<M: Clone + PreprocessMachine> SigningProtocol<M> {
+impl<D: Db, M: Clone + PreprocessMachine> SigningProtocol<D, M> {
   /// Create a new signing protocol.
-  pub(crate) fn new(session: Session, start_i: Participant, id: [u8; 32], root: Vec<M>) -> Self {
+  pub(crate) fn new(
+    db: D,
+    session: Session,
+    start_i: Participant,
+    id: [u8; 32],
+    root: Vec<M>,
+  ) -> Self {
     log::info!("starting signing protocol {}", hex::encode(id));
 
     Self {
+      db,
       session,
       start_i,
       id,
@@ -70,7 +85,15 @@ impl<M: Clone + PreprocessMachine> SigningProtocol<M> {
       We also won't send the share we were supposed to, unfortunately, yet caching/reloading the
       preprocess has enough safety issues it isn't worth the headache.
     */
-    // TODO
+    {
+      let mut txn = self.db.txn();
+      let prior_attempted = Attempted::get(&txn, self.id);
+      if Some(attempt) <= prior_attempted {
+        return vec![];
+      }
+      Attempted::set(&mut txn, self.id, &attempt);
+      txn.commit();
+    }
 
     log::debug!("attemting a new instance of signing protocol {}", hex::encode(self.id));
 
@@ -248,4 +271,11 @@ impl<M: Clone + PreprocessMachine> SigningProtocol<M> {
 
     Ok(signature)
   }
+
+  /// Cleanup the database entries for a specified signing protocol.
+  pub(crate) fn cleanup(db: &mut D, id: [u8; 32]) {
+    let mut txn = db.txn();
+    Attempted::del(&mut txn, id);
+    txn.commit();
+  }
 }

processor/frost-attempt-manager/src/lib.rs

@@ -8,6 +8,7 @@ use frost::{Participant, sign::PreprocessMachine};
 
 use serai_validator_sets_primitives::Session;
 
+use serai_db::Db;
 use messages::sign::{ProcessorMessage, CoordinatorMessage};
 
 mod individual;
@@ -22,21 +23,28 @@ pub enum Response<M: PreprocessMachine> {
 }
 
 /// A manager of attempts for a variety of signing protocols.
-pub struct AttemptManager<M: Clone + PreprocessMachine> {
+pub struct AttemptManager<D: Db, M: Clone + PreprocessMachine> {
+  db: D,
   session: Session,
   start_i: Participant,
-  active: HashMap<[u8; 32], SigningProtocol<M>>,
+  active: HashMap<[u8; 32], SigningProtocol<D, M>>,
 }
 
-impl<M: Clone + PreprocessMachine> AttemptManager<M> {
+impl<D: Db, M: Clone + PreprocessMachine> AttemptManager<D, M> {
   /// Create a new attempt manager.
-  pub fn new(session: Session, start_i: Participant) -> Self {
-    AttemptManager { session, start_i, active: HashMap::new() }
+  pub fn new(db: D, session: Session, start_i: Participant) -> Self {
+    AttemptManager { db, session, start_i, active: HashMap::new() }
   }
 
   /// Register a signing protocol to attempt.
-  pub fn register(&mut self, id: [u8; 32], machines: Vec<M>) {
-    self.active.insert(id, SigningProtocol::new(self.session, self.start_i, id, machines));
+  ///
+  /// This ID must be unique across all sessions, attempt managers, protocols, etc.
+  pub fn register(&mut self, id: [u8; 32], machines: Vec<M>) -> Vec<ProcessorMessage> {
+    let mut protocol =
+      SigningProtocol::new(self.db.clone(), self.session, self.start_i, id, machines);
+    let messages = protocol.attempt(0);
+    self.active.insert(id, protocol);
+    messages
   }
 
   /// Retire a signing protocol.
@@ -45,10 +53,13 @@ impl<M: Clone + PreprocessMachine> AttemptManager<M> {
   /// This does not stop the protocol from being re-registered and further worked on (with
   /// undefined behavior) then. The higher-level context must never call `register` again with this
   /// ID.
-  // TODO: Also have the DB for this SigningProtocol cleaned up here.
   pub fn retire(&mut self, id: [u8; 32]) {
-    log::info!("retiring signing protocol {}", hex::encode(id));
-    self.active.remove(&id);
+    if self.active.remove(&id).is_none() {
+      log::info!("retiring protocol {}, which we didn't register/already retired", hex::encode(id));
+    } else {
+      log::info!("retired signing protocol {}", hex::encode(id));
+    }
+    SigningProtocol::<D, M>::cleanup(&mut self.db, id);
   }
 
   /// Handle a message for a signing protocol.

processor/key-gen/src/db.rs

@@ -36,10 +36,10 @@ pub(crate) struct Participations {
 }
 
 create_db!(
-  KeyGenDb {
-    ParamsDb: (session: &Session) -> RawParams,
-    ParticipationsDb: (session: &Session) -> Participations,
-    KeySharesDb: (session: &Session) -> Vec<u8>,
+  KeyGen {
+    Params: (session: &Session) -> RawParams,
+    Participations: (session: &Session) -> Participations,
+    KeyShares: (session: &Session) -> Vec<u8>,
   }
 );
 
@@ -48,7 +48,7 @@ impl<P: KeyGenParams> KeyGenDb<P> {
   pub(crate) fn set_params(txn: &mut impl DbTxn, session: Session, params: Params<P>) {
     assert_eq!(params.substrate_evrf_public_keys.len(), params.network_evrf_public_keys.len());
 
-    ParamsDb::set(
+    Params::set(
       txn,
       &session,
       &RawParams {
@@ -68,7 +68,7 @@ impl<P: KeyGenParams> KeyGenDb<P> {
   }
 
   pub(crate) fn params(getter: &impl Get, session: Session) -> Option<Params<P>> {
-    ParamsDb::get(getter, &session).map(|params| Params {
+    Params::get(getter, &session).map(|params| Params {
       t: params.t,
       n: params
         .network_evrf_public_keys
@@ -101,12 +101,13 @@ impl<P: KeyGenParams> KeyGenDb<P> {
     session: Session,
     participations: &Participations,
   ) {
-    ParticipationsDb::set(txn, &session, participations)
+    Participations::set(txn, &session, participations)
   }
   pub(crate) fn participations(getter: &impl Get, session: Session) -> Option<Participations> {
-    ParticipationsDb::get(getter, &session)
+    Participations::get(getter, &session)
   }
 
+  // Set the key shares for a session.
   pub(crate) fn set_key_shares(
     txn: &mut impl DbTxn,
     session: Session,
@@ -120,7 +121,7 @@ impl<P: KeyGenParams> KeyGenDb<P> {
       keys.extend(substrate_keys.serialize().as_slice());
       keys.extend(network_keys.serialize().as_slice());
     }
-    KeySharesDb::set(txn, &session, &keys);
+    KeyShares::set(txn, &session, &keys);
   }
 
   #[allow(clippy::type_complexity)]
@@ -128,7 +129,7 @@ impl<P: KeyGenParams> KeyGenDb<P> {
     getter: &impl Get,
     session: Session,
   ) -> Option<(Vec<ThresholdKeys<Ristretto>>, Vec<ThresholdKeys<P::ExternalNetworkCurve>>)> {
-    let keys = KeySharesDb::get(getter, &session)?;
+    let keys = KeyShares::get(getter, &session)?;
     let mut keys: &[u8] = keys.as_ref();
 
     let mut substrate_keys = vec![];

processor/key-gen/src/lib.rs

@@ -182,7 +182,7 @@ impl<P: KeyGenParams, D: Db> KeyGen<P, D> {
 
     match msg {
       CoordinatorMessage::GenerateKey { session, threshold, evrf_public_keys } => {
-        log::info!("Generating new key. Session: {session:?}");
+        log::info!("generating new key, session: {session:?}");
 
         // Unzip the vector of eVRF keys
         let substrate_evrf_public_keys =
@@ -258,7 +258,7 @@ impl<P: KeyGenParams, D: Db> KeyGen<P, D> {
       }
 
       CoordinatorMessage::Participation { session, participant, participation } => {
-        log::info!("received participation from {:?} for {:?}", participant, session);
+        log::debug!("received participation from {:?} for {:?}", participant, session);
 
         let Params { t: threshold, n, substrate_evrf_public_keys, network_evrf_public_keys } =
           KeyGenDb::<P>::params(txn, session).unwrap();
@@ -293,7 +293,7 @@ impl<P: KeyGenParams, D: Db> KeyGen<P, D> {
             // participations and continue. We solely have to verify them, as to identify malicious
             // participants and prevent DoSs, before returning
             if self.key_shares(session).is_some() {
-              log::info!("already finished generating a key for {:?}", session);
+              log::debug!("already finished generating a key for {:?}", session);
 
               match EvrfDkg::<Ristretto>::verify(
                 &mut OsRng,
@@ -511,6 +511,8 @@ impl<P: KeyGenParams, D: Db> KeyGen<P, D> {
         }
         KeyGenDb::<P>::set_key_shares(txn, session, &substrate_keys, &network_keys);
 
+        log::info!("generated key, session: {session:?}");
+
         // Since no one we verified was invalid, and we had the threshold, yield the new keys
         vec![ProcessorMessage::GeneratedKeyPair {
           session,