Support multiple batches per block by the coordinator

Also corrects an assumption block hash == batch ID.

coordinator/src/db.rs

@@ -1,3 +1,6 @@
+use scale::Encode;
+use serai_client::primitives::{NetworkId, BlockHash};
+
 pub use serai_db::*;
 
 use crate::tributary::TributarySpec;
@@ -42,6 +45,42 @@ impl<'a, D: Db> MainDb<'a, D> {
     txn.commit();
   }
 
+  fn batches_in_block_key(network: NetworkId, block: [u8; 32]) -> Vec<u8> {
+    Self::main_key(b"batches_in_block", (network, block).encode())
+  }
+  pub fn batches_in_block<G: Get>(
+    getter: &G,
+    network: NetworkId,
+    block: [u8; 32],
+  ) -> Vec<[u8; 32]> {
+    getter
+      .get(Self::batches_in_block_key(network, block))
+      .expect("asking for batches in block for block without batches")
+      .chunks(32)
+      .map(|id| id.try_into().unwrap())
+      .collect()
+  }
+  pub fn add_batch_to_block(
+    txn: &mut D::Transaction<'_>,
+    network: NetworkId,
+    block: BlockHash,
+    id: [u8; 32],
+  ) {
+    let key = Self::batches_in_block_key(network, block.0);
+    let Some(mut existing) = txn.get(&key) else {
+      txn.put(&key, block.0);
+      return;
+    };
+
+    if existing.chunks(32).any(|existing_id| existing_id == id) {
+      // TODO: Is this an invariant?
+      return;
+    }
+
+    existing.extend(block.0);
+    txn.put(&key, existing);
+  }
+
   fn first_preprocess_key(id: [u8; 32]) -> Vec<u8> {
     Self::main_key(b"first_preprocess", id)
   }

coordinator/src/main.rs

@@ -17,7 +17,7 @@ use ciphersuite::{group::ff::PrimeField, Ciphersuite, Ristretto};
 use serai_db::{DbTxn, Db};
 use serai_env as env;
 
-use serai_client::{Public, Serai};
+use serai_client::{primitives::NetworkId, Public, Serai};
 
 use message_queue::{Service, client::MessageQueue};
 
@@ -146,7 +146,7 @@ pub async fn scan_substrate<D: Db, Pro: Processors>(
 pub async fn scan_tributaries<D: Db, Pro: Processors, P: P2p>(
   raw_db: D,
   key: Zeroizing<<Ristretto as Ciphersuite>::F>,
-  recognized_id_send: UnboundedSender<([u8; 32], RecognizedIdType, [u8; 32])>,
+  recognized_id_send: UnboundedSender<(NetworkId, [u8; 32], RecognizedIdType, [u8; 32])>,
   p2p: P,
   processors: Pro,
   serai: Arc<Serai>,
@@ -544,6 +544,7 @@ pub async fn handle_processors<D: Db, Pro: Processors, P: P2p>(
             // before this message finishes it handling (or with this message's finished handling)
             let mut txn = db.txn();
             MainDb::<D>::save_first_preprocess(&mut txn, id.id, preprocess);
+            MainDb::<D>::add_batch_to_block(&mut txn, msg.network, block, id.id);
             txn.commit();
 
             // TODO: This will publish one ExternalBlock per Batch. We should only publish one per
@@ -701,21 +702,27 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
     let tributaries = tributaries.clone();
     async move {
       loop {
-        if let Some((genesis, id_type, id)) = recognized_id_recv.recv().await {
-          let mut tx = match id_type {
-            RecognizedIdType::Block => Transaction::BatchPreprocess(SignData {
-              plan: id,
-              attempt: 0,
-              data: MainDb::<D>::first_preprocess(&raw_db, id),
-              signed: Transaction::empty_signed(),
-            }),
+        if let Some((network, genesis, id_type, id)) = recognized_id_recv.recv().await {
+          let txs = match id_type {
+            RecognizedIdType::Block => {
+              let mut txs = vec![];
+              for id in MainDb::<D>::batches_in_block(&raw_db, network, id) {
+                txs.push(Transaction::BatchPreprocess(SignData {
+                  plan: id,
+                  attempt: 0,
+                  data: MainDb::<D>::first_preprocess(&raw_db, id),
+                  signed: Transaction::empty_signed(),
+                }));
+              }
+              txs
+            }
 
-            RecognizedIdType::Plan => Transaction::SignPreprocess(SignData {
+            RecognizedIdType::Plan => vec![Transaction::SignPreprocess(SignData {
               plan: id,
               attempt: 0,
               data: MainDb::<D>::first_preprocess(&raw_db, id),
               signed: Transaction::empty_signed(),
-            }),
+            })],
           };
 
           let tributaries = tributaries.read().await;
@@ -724,15 +731,17 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
           };
           let tributary = tributary.tributary.read().await;
 
-          // TODO: Same note as prior nonce acquisition
-          log::trace!("getting next nonce for Tributary TX containing Batch signing data");
-          let nonce = tributary
-            .next_nonce(Ristretto::generator() * key.deref())
-            .await
-            .expect("publishing a TX to a tributary we aren't in");
-          tx.sign(&mut OsRng, genesis, &key, nonce);
+          for mut tx in txs {
+            // TODO: Same note as prior nonce acquisition
+            log::trace!("getting next nonce for Tributary TX containing Batch signing data");
+            let nonce = tributary
+              .next_nonce(Ristretto::generator() * key.deref())
+              .await
+              .expect("publishing a TX to a tributary we aren't in");
+            tx.sign(&mut OsRng, genesis, &key, nonce);
 
-          publish_transaction(&tributary, tx).await;
+            publish_transaction(&tributary, tx).await;
+          }
         } else {
           log::warn!("recognized_id_send was dropped. are we shutting down?");
           break;

coordinator/src/tributary/handle.rs

@@ -19,6 +19,7 @@ use tokio::sync::mpsc::UnboundedSender;
 
 use serai_client::{
   Signature,
+  primitives::NetworkId,
   validator_sets::primitives::{ValidatorSet, KeyPair, musig_context, set_keys_message},
   subxt::utils::Encoded,
   Serai,
@@ -232,7 +233,7 @@ pub async fn handle_application_tx<
   publish_serai_tx: PST,
   genesis: [u8; 32],
   key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
-  recognized_id: &UnboundedSender<([u8; 32], RecognizedIdType, [u8; 32])>,
+  recognized_id: &UnboundedSender<(NetworkId, [u8; 32], RecognizedIdType, [u8; 32])>,
   txn: &mut <D as Db>::Transaction<'_>,
 ) {
   // Used to determine if an ID is acceptable
@@ -433,7 +434,7 @@ pub async fn handle_application_tx<
       // Because this external block has been finalized, its batch ID should be authorized
       TributaryDb::<D>::recognize_id(txn, Zone::Batch.label(), genesis, block);
       recognized_id
-        .send((genesis, RecognizedIdType::Block, block))
+        .send((spec.set().network, genesis, RecognizedIdType::Block, block))
         .expect("recognized_id_recv was dropped. are we shutting down?");
     }
 
@@ -446,7 +447,7 @@ pub async fn handle_application_tx<
       for id in plan_ids {
         TributaryDb::<D>::recognize_id(txn, Zone::Sign.label(), genesis, id);
         recognized_id
-          .send((genesis, RecognizedIdType::Plan, id))
+          .send((spec.set().network, genesis, RecognizedIdType::Plan, id))
           .expect("recognized_id_recv was dropped. are we shutting down?");
       }
     }

coordinator/src/tributary/scanner.rs

@@ -4,7 +4,9 @@ use zeroize::Zeroizing;
 
 use ciphersuite::{Ciphersuite, Ristretto};
 
-use serai_client::{validator_sets::primitives::ValidatorSet, subxt::utils::Encoded};
+use serai_client::{
+  primitives::NetworkId, validator_sets::primitives::ValidatorSet, subxt::utils::Encoded,
+};
 
 use tokio::sync::mpsc::UnboundedSender;
 
@@ -43,7 +45,7 @@ async fn handle_block<
 >(
   db: &mut TributaryDb<D>,
   key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
-  recognized_id: &UnboundedSender<([u8; 32], RecognizedIdType, [u8; 32])>,
+  recognized_id: &UnboundedSender<(NetworkId, [u8; 32], RecognizedIdType, [u8; 32])>,
   processors: &Pro,
   publish_serai_tx: PST,
   spec: &TributarySpec,
@@ -109,7 +111,7 @@ pub async fn handle_new_blocks<
 >(
   db: &mut TributaryDb<D>,
   key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
-  recognized_id: &UnboundedSender<([u8; 32], RecognizedIdType, [u8; 32])>,
+  recognized_id: &UnboundedSender<(NetworkId, [u8; 32], RecognizedIdType, [u8; 32])>,
   processors: &Pro,
   publish_serai_tx: PST,
   spec: &TributarySpec,