Remove Serai from the ConfirmDkgTask

coordinator/src/db.rs

@@ -80,6 +80,8 @@ create_db! {
     ErroneousCosigns: () -> Vec<SignedCosign>,
     // The keys to confirm and set on the Serai network
     KeysToConfirm: (set: ValidatorSet) -> KeyPair,
+    // The key was set on the Serai network
+    KeySet: (set: ValidatorSet) -> (),
   }
 }
 

coordinator/src/dkg_confirmation.rs

@@ -1,5 +1,5 @@
 use core::{ops::Deref, future::Future};
-use std::{boxed::Box, sync::Arc, collections::HashMap};
+use std::{boxed::Box, collections::HashMap};
 
 use zeroize::Zeroizing;
 use rand_core::OsRng;
@@ -18,15 +18,14 @@ use serai_db::{DbTxn, Db as DbTrait};
 use serai_client::{
   primitives::SeraiAddress,
   validator_sets::primitives::{ValidatorSet, musig_context, set_keys_message},
-  SeraiError, Serai,
 };
 
-use serai_task::ContinuallyRan;
+use serai_task::{DoesNotError, ContinuallyRan};
 
 use serai_coordinator_substrate::{NewSetInformation, Keys};
 use serai_coordinator_tributary::{Transaction, DkgConfirmationMessages};
 
-use crate::{KeysToConfirm, TributaryTransactionsFromDkgConfirmation};
+use crate::{KeysToConfirm, KeySet, TributaryTransactionsFromDkgConfirmation};
 
 fn schnorrkel() -> Schnorrkel {
   Schnorrkel::new(b"substrate") // TODO: Pull the constant for this
@@ -128,8 +127,6 @@ pub(crate) struct ConfirmDkgTask<CD: DbTrait, TD: DbTrait> {
   set: NewSetInformation,
   tributary_db: TD,
 
-  serai: Arc<Serai>,
-
   key: Zeroizing<<Ristretto as Ciphersuite>::F>,
   signer: Option<Signer>,
 }
@@ -139,10 +136,9 @@ impl<CD: DbTrait, TD: DbTrait> ConfirmDkgTask<CD, TD> {
     db: CD,
     set: NewSetInformation,
     tributary_db: TD,
-    serai: Arc<Serai>,
     key: Zeroizing<<Ristretto as Ciphersuite>::F>,
   ) -> Self {
-    Self { db, set, tributary_db, serai, key, signer: None }
+    Self { db, set, tributary_db, key, signer: None }
   }
 
   fn slash(db: &mut CD, set: ValidatorSet, validator: SeraiAddress) {
@@ -192,7 +188,7 @@ impl<CD: DbTrait, TD: DbTrait> ConfirmDkgTask<CD, TD> {
 }
 
 impl<CD: DbTrait, TD: DbTrait> ContinuallyRan for ConfirmDkgTask<CD, TD> {
-  type Error = SeraiError;
+  type Error = DoesNotError;
 
   fn run_iteration(&mut self) -> impl Send + Future<Output = Result<bool, Self::Error>> {
     async move {
@@ -416,24 +412,20 @@ impl<CD: DbTrait, TD: DbTrait> ContinuallyRan for ConfirmDkgTask<CD, TD> {
       }
 
       // Check if the key has been set on Serai
-      if KeysToConfirm::get(&self.db, self.set.set).is_some() {
-        let serai = self.serai.as_of_latest_finalized_block().await?;
-        let serai = serai.validator_sets();
-        let is_historic_set = serai.session(self.set.set.network).await?.map(|session| session.0) >
-          Some(self.set.set.session.0);
-        let key_set_on_serai = is_historic_set || serai.keys(self.set.set).await?.is_some();
-        if key_set_on_serai {
-          // Take the keys to confirm so we never instantiate the signer again
-          let mut txn = self.db.txn();
-          KeysToConfirm::take(&mut txn, self.set.set);
-          txn.commit();
+      if KeysToConfirm::get(&self.db, self.set.set).is_some() &&
+        KeySet::get(&self.db, self.set.set).is_some()
+      {
+        // Take the keys to confirm so we never instantiate the signer again
+        let mut txn = self.db.txn();
+        KeysToConfirm::take(&mut txn, self.set.set);
+        KeySet::take(&mut txn, self.set.set);
+        txn.commit();
 
-          // Drop our own signer
-          // The task won't die until the Tributary does, but now it'll never do anything again
-          self.signer = None;
+        // Drop our own signer
+        // The task won't die until the Tributary does, but now it'll never do anything again
+        self.signer = None;
 
-          made_progress = true;
-        }
+        made_progress = true;
       }
 
       Ok(made_progress)

coordinator/src/main.rs

@@ -368,6 +368,7 @@ async fn main() {
       prune_tributary_db(to_cleanup);
       // Remove the keys to confirm for this network
       KeysToConfirm::take(&mut txn, to_cleanup);
+      KeySet::take(&mut txn, to_cleanup);
       // Drain the cosign intents created for this set
       while !Cosigning::<Db>::intended_cosigns(&mut txn, to_cleanup).is_empty() {}
       // Drain the transactions to publish for this set
@@ -461,7 +462,6 @@ async fn main() {
       p2p.clone(),
       &p2p_add_tributary_send,
       tributary,
-      serai.clone(),
       serai_key.clone(),
     )
     .await;
@@ -476,7 +476,6 @@ async fn main() {
       p2p: p2p.clone(),
       p2p_add_tributary: p2p_add_tributary_send.clone(),
       p2p_retire_tributary: p2p_retire_tributary_send.clone(),
-      serai: serai.clone(),
     })
     .continually_run(substrate_task_def, vec![]),
   );

coordinator/src/substrate.rs

@@ -9,10 +9,7 @@ use tokio::sync::mpsc;
 
 use serai_db::{DbTxn, Db as DbTrait};
 
-use serai_client::{
-  validator_sets::primitives::{Session, ValidatorSet},
-  Serai,
-};
+use serai_client::validator_sets::primitives::{Session, ValidatorSet};
 use message_queue::{Service, Metadata, client::MessageQueue};
 
 use tributary_sdk::Tributary;
@@ -22,7 +19,7 @@ use serai_task::ContinuallyRan;
 use serai_coordinator_tributary::Transaction;
 use serai_coordinator_p2p::P2p;
 
-use crate::Db;
+use crate::{Db, KeySet};
 
 pub(crate) struct SubstrateTask<P: P2p> {
   pub(crate) serai_key: Zeroizing<<Ristretto as Ciphersuite>::F>,
@@ -32,7 +29,6 @@ pub(crate) struct SubstrateTask<P: P2p> {
   pub(crate) p2p_add_tributary:
     mpsc::UnboundedSender<(ValidatorSet, Tributary<Db, Transaction, P>)>,
   pub(crate) p2p_retire_tributary: mpsc::UnboundedSender<ValidatorSet>,
-  pub(crate) serai: Arc<Serai>,
 }
 
 impl<P: P2p> ContinuallyRan for SubstrateTask<P> {
@@ -51,8 +47,9 @@ impl<P: P2p> ContinuallyRan for SubstrateTask<P> {
           };
 
           match msg {
-            // TODO: Stop trying to confirm the DKG
-            messages::substrate::CoordinatorMessage::SetKeys { .. } => todo!("TODO"),
+            messages::substrate::CoordinatorMessage::SetKeys { session, .. } => {
+              KeySet::set(&mut txn, ValidatorSet { network, session }, &());
+            }
             messages::substrate::CoordinatorMessage::SlashesReported { session } => {
               let prior_retired = crate::db::RetiredTributary::get(&txn, network);
               let next_to_be_retired =
@@ -150,7 +147,6 @@ impl<P: P2p> ContinuallyRan for SubstrateTask<P> {
           self.p2p.clone(),
           &self.p2p_add_tributary,
           new_set,
-          self.serai.clone(),
           self.serai_key.clone(),
         )
         .await;

coordinator/src/tributary.rs

@@ -11,7 +11,7 @@ use tokio::sync::mpsc;
 use serai_db::{Get, DbTxn, Db as DbTrait, create_db, db_channel};
 
 use scale::Encode;
-use serai_client::{validator_sets::primitives::ValidatorSet, Serai};
+use serai_client::validator_sets::primitives::ValidatorSet;
 
 use tributary_sdk::{TransactionKind, TransactionError, ProvidedError, TransactionTrait, Tributary};
 
@@ -471,7 +471,6 @@ pub(crate) async fn spawn_tributary<P: P2p>(
   p2p: P,
   p2p_add_tributary: &mpsc::UnboundedSender<(ValidatorSet, Tributary<Db, Transaction, P>)>,
   set: NewSetInformation,
-  serai: Arc<Serai>,
   serai_key: Zeroizing<<Ristretto as Ciphersuite>::F>,
 ) {
   // Don't spawn retired Tributaries
@@ -566,7 +565,7 @@ pub(crate) async fn spawn_tributary<P: P2p>(
   // Spawn the task to confirm the DKG result
   let (confirm_dkg_task_def, confirm_dkg_task) = Task::new();
   tokio::spawn(
-    ConfirmDkgTask::new(db.clone(), set.clone(), tributary_db.clone(), serai, serai_key.clone())
+    ConfirmDkgTask::new(db.clone(), set.clone(), tributary_db.clone(), serai_key.clone())
       .continually_run(confirm_dkg_task_def, vec![add_tributary_transactions_task]),
   );