Complete serai-coordinator-p2p

coordinator/p2p/src/heartbeat.rs

@@ -1,7 +1,7 @@
 use core::future::Future;
 use std::time::{Duration, SystemTime};
 
-use serai_client::validator_sets::primitives::ValidatorSet;
+use serai_client::validator_sets::primitives::{MAX_KEY_SHARES_PER_SET, ValidatorSet};
 
 use futures_lite::FutureExt;
 
@@ -15,19 +15,32 @@ use crate::{Heartbeat, Peer, P2p};
 // Amount of blocks in a minute
 const BLOCKS_PER_MINUTE: usize = (60 / (tributary::tendermint::TARGET_BLOCK_TIME / 1000)) as usize;
 
-/// The maximum amount of blocks to include/included within a batch.
-pub const BLOCKS_PER_BATCH: usize = BLOCKS_PER_MINUTE + 1;
+/// The minimum amount of blocks to include/included within a batch, assuming there's blocks to
+/// include in the batch.
+///
+/// This decides the size limit of the Batch (the Block size limit multiplied by the minimum amount
+/// of blocks we'll send). The actual amount of blocks sent will be the amount which fits within
+/// the size limit.
+pub const MIN_BLOCKS_PER_BATCH: usize = BLOCKS_PER_MINUTE + 1;
+
+/// The size limit for a batch of blocks sent in response to a Heartbeat.
+///
+/// This estimates the size of a commit as `32 + (MAX_VALIDATORS * 128)`. At the time of writing, a
+/// commit is `8 + (validators * 32) + (32 + (validators * 32))` (for the time, list of validators,
+/// and aggregate signature). Accordingly, this should be a safe over-estimate.
+pub const BATCH_SIZE_LIMIT: usize = MIN_BLOCKS_PER_BATCH *
+  (tributary::BLOCK_SIZE_LIMIT + 32 + ((MAX_KEY_SHARES_PER_SET as usize) * 128));
 
 /// Sends a heartbeat to other validators on regular intervals informing them of our Tributary's
 /// tip.
 ///
 /// If the other validator has more blocks then we do, they're expected to inform us. This forms
 /// the sync protocol for our Tributaries.
-pub struct HeartbeatTask<TD: Db, Tx: TransactionTrait, P: P2p> {
-  set: ValidatorSet,
-  tributary: Tributary<TD, Tx, P>,
-  reader: TributaryReader<TD, Tx>,
-  p2p: P,
+pub(crate) struct HeartbeatTask<TD: Db, Tx: TransactionTrait, P: P2p> {
+  pub(crate) set: ValidatorSet,
+  pub(crate) tributary: Tributary<TD, Tx, P>,
+  pub(crate) reader: TributaryReader<TD, Tx>,
+  pub(crate) p2p: P,
 }
 
 impl<TD: Db, Tx: TransactionTrait, P: P2p> ContinuallyRan for HeartbeatTask<TD, Tx, P> {
@@ -80,7 +93,7 @@ impl<TD: Db, Tx: TransactionTrait, P: P2p> ContinuallyRan for HeartbeatTask<TD,
 
             // This is the final batch if it has less than the maximum amount of blocks
             // (signifying there weren't more blocks after this to fill the batch with)
-            let final_batch = blocks.len() < BLOCKS_PER_BATCH;
+            let final_batch = blocks.len() < MIN_BLOCKS_PER_BATCH;
 
             // Sync each block
             for block_with_commit in blocks {

coordinator/p2p/src/lib.rs

@@ -3,18 +3,23 @@
 #![deny(missing_docs)]
 
 use core::future::Future;
+use std::collections::HashMap;
 
 use borsh::{BorshSerialize, BorshDeserialize};
 
 use serai_client::{primitives::NetworkId, validator_sets::primitives::ValidatorSet};
 
-use serai_cosign::SignedCosign;
+use serai_db::Db;
+use tributary::{ReadWrite, TransactionTrait, Tributary, TributaryReader};
+use serai_cosign::{SignedCosign, Cosigning};
 
-/// A oneshot channel.
-pub mod oneshot;
+use tokio::sync::{mpsc, oneshot};
+
+use serai_task::{Task, ContinuallyRan};
 
 /// The heartbeat task, effecting sync of Tributaries
 pub mod heartbeat;
+use crate::heartbeat::HeartbeatTask;
 
 /// A heartbeat for a Tributary.
 #[derive(Clone, Copy, BorshSerialize, BorshDeserialize, Debug)]
@@ -74,3 +79,116 @@ pub trait P2p: Send + Sync + Clone + tributary::P2p + serai_cosign::RequestNotab
   /// A cancel-safe future for the next cosign received.
   fn cosign(&self) -> impl Send + Future<Output = SignedCosign>;
 }
+
+fn handle_notable_cosigns_request<D: Db>(
+  db: &D,
+  global_session: [u8; 32],
+  channel: oneshot::Sender<Vec<SignedCosign>>,
+) {
+  let cosigns = Cosigning::<D>::notable_cosigns(db, global_session);
+  channel.send(cosigns).expect("channel listening for cosign oneshot response was dropped?");
+}
+
+fn handle_heartbeat<D: Db, T: TransactionTrait>(
+  reader: &TributaryReader<D, T>,
+  mut latest_block_hash: [u8; 32],
+  channel: oneshot::Sender<Vec<TributaryBlockWithCommit>>,
+) {
+  let mut res_size = 8;
+  let mut res = vec![];
+  // This former case should be covered by this latter case
+  while (res.len() < heartbeat::MIN_BLOCKS_PER_BATCH) || (res_size < heartbeat::BATCH_SIZE_LIMIT) {
+    let Some(block_after) = reader.block_after(&latest_block_hash) else { break };
+
+    let block = reader.block(&block_after).unwrap().serialize();
+    let commit = reader.commit(&block_after).unwrap();
+    res_size += 8 + block.len() + 8 + commit.len();
+    res.push(TributaryBlockWithCommit { block, commit });
+
+    latest_block_hash = block_after;
+  }
+  channel
+    .send(res)
+    .map_err(|_| ())
+    .expect("channel listening for heartbeat oneshot response was dropped?");
+}
+
+/// Run the P2P instance.
+///
+/// `add_tributary`'s and `retire_tributary's senders, along with `send_cosigns`'s receiver, must
+/// never be dropped. `retire_tributary` is not required to only be instructed with added
+/// Tributaries.
+pub async fn run<TD: Db, Tx: TransactionTrait, P: P2p>(
+  db: impl Db,
+  p2p: P,
+  mut add_tributary: mpsc::UnboundedReceiver<(ValidatorSet, Tributary<TD, Tx, P>)>,
+  mut retire_tributary: mpsc::UnboundedReceiver<ValidatorSet>,
+  send_cosigns: mpsc::UnboundedSender<SignedCosign>,
+) {
+  let mut readers = HashMap::<ValidatorSet, TributaryReader<TD, Tx>>::new();
+  let mut tributaries = HashMap::<[u8; 32], mpsc::UnboundedSender<Vec<u8>>>::new();
+  let mut heartbeat_tasks = HashMap::<ValidatorSet, _>::new();
+
+  loop {
+    tokio::select! {
+      tributary = add_tributary.recv() => {
+        let (set, tributary) = tributary.expect("add_tributary send was dropped?");
+        let reader = tributary.reader();
+        readers.insert(set, reader.clone());
+
+        let (heartbeat_task_def, heartbeat_task) = Task::new();
+        tokio::spawn(
+          (HeartbeatTask {
+            set,
+            tributary: tributary.clone(),
+            reader: reader.clone(),
+            p2p: p2p.clone(),
+          }).continually_run(heartbeat_task_def, vec![])
+        );
+        heartbeat_tasks.insert(set, heartbeat_task);
+
+        let (tributary_message_send, mut tributary_message_recv) = mpsc::unbounded_channel();
+        tributaries.insert(tributary.genesis(), tributary_message_send);
+        // For as long as this sender exists, handle the messages from it on a dedicated task
+        tokio::spawn(async move {
+          while let Some(message) = tributary_message_recv.recv().await {
+            tributary.handle_message(&message).await;
+          }
+        });
+      }
+      set = retire_tributary.recv() => {
+        let set = set.expect("retire_tributary send was dropped?");
+        let Some(reader) = readers.remove(&set) else { continue };
+        tributaries.remove(&reader.genesis()).expect("tributary reader but no tributary");
+        heartbeat_tasks.remove(&set).expect("tributary but no heartbeat task");
+      }
+
+      (heartbeat, channel) = p2p.heartbeat() => {
+        if let Some(reader) = readers.get(&heartbeat.set) {
+          let reader = reader.clone(); // This is a cheap clone
+          // We spawn this on a task due to the DB reads needed
+          tokio::spawn(async move {
+            handle_heartbeat(&reader, heartbeat.latest_block_hash, channel)
+          });
+        }
+      }
+      (global_session, channel) = p2p.notable_cosigns_request() => {
+        tokio::spawn({
+          let db = db.clone();
+          async move { handle_notable_cosigns_request(&db, global_session, channel) }
+        });
+      }
+      (tributary, message) = p2p.tributary_message() => {
+        if let Some(tributary) = tributaries.get(&tributary) {
+          tributary.send(message).expect("tributary message recv was dropped?");
+        }
+      }
+      cosign = p2p.cosign() => {
+        // We don't call `Cosigning::intake_cosign` here as that can only be called from a single
+        // location. We also need to intake the cosigns we produce, which means we need to merge
+        // these streams (signing, network) somehow. That's done with this mpsc channel
+        send_cosigns.send(cosign).expect("channel receiving cosigns was dropped?");
+      }
+    }
+  }
+}

coordinator/p2p/src/oneshot.rs

@@ -1,35 +0,0 @@
-use core::{
-  pin::Pin,
-  task::{Poll, Context},
-  future::Future,
-};
-
-pub use async_channel::{SendError, RecvError};
-
-/// The sender for a oneshot channel.
-pub struct Sender<T: Send>(async_channel::Sender<T>);
-impl<T: Send> Sender<T> {
-  /// Send a value down the channel.
-  ///
-  /// Returns an error if the channel's receiver was dropped.
-  pub fn send(self, msg: T) -> Result<(), SendError<T>> {
-    self.0.send_blocking(msg)
-  }
-}
-
-/// The receiver for a oneshot channel.
-pub struct Receiver<T: Send>(async_channel::Receiver<T>);
-impl<T: Send> Future for Receiver<T> {
-  type Output = Result<T, RecvError>;
-  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-    let recv = self.0.recv();
-    futures_lite::pin!(recv);
-    recv.poll(cx)
-  }
-}
-
-/// Create a new oneshot channel.
-pub fn channel<T: Send>() -> (Sender<T>, Receiver<T>) {
-  let (send, recv) = async_channel::bounded(1);
-  (Sender(send), Receiver(recv))
-}