Add input aggregation in the transaction-chaining scheduler

Also handles some other misc in it.

Cargo.lock

@@ -8741,6 +8741,7 @@ dependencies = [
  "serai-primitives",
  "serai-processor-primitives",
  "serai-processor-scanner",
+ "serai-processor-scheduler-primitives",
 ]
 
 [[package]]

processor/scheduler/utxo/primitives/Cargo.toml

@@ -23,3 +23,4 @@ serai-primitives = { path = "../../../../substrate/primitives", default-features
 
 primitives = { package = "serai-processor-primitives", path = "../../../primitives" }
 scanner = { package = "serai-processor-scanner", path = "../../../scanner" }
+scheduler-primitives = { package = "serai-processor-scheduler-primitives", path = "../../primitives" }

processor/scheduler/utxo/primitives/src/lib.rs

@@ -7,11 +7,22 @@ use core::fmt::Debug;
 use serai_primitives::{Coin, Amount};
 
 use primitives::{ReceivedOutput, Payment};
-use scanner::{ScannerFeed, KeyFor, AddressFor, OutputFor};
+use scanner::{ScannerFeed, KeyFor, AddressFor, OutputFor, EventualityFor};
+use scheduler_primitives::*;
+
+/// A planned transaction.
+pub struct PlannedTransaction<S: ScannerFeed, ST: SignableTransaction, A> {
+  /// The signable transaction.
+  pub signable: ST,
+  /// The Eventuality to watch for.
+  pub eventuality: EventualityFor<S>,
+  /// The auxilliary data for this transaction.
+  pub auxilliary: A,
+}
 
 /// An object able to plan a transaction.
 #[async_trait::async_trait]
-pub trait TransactionPlanner<S: ScannerFeed>: 'static + Send + Sync {
+pub trait TransactionPlanner<S: ScannerFeed, A>: 'static + Send + Sync {
   /// An error encountered when determining the fee rate.
   ///
   /// This MUST be an ephemeral error. Retrying fetching data from the blockchain MUST eventually
@@ -21,8 +32,8 @@ pub trait TransactionPlanner<S: ScannerFeed>: 'static + Send + Sync {
   /// The type representing a fee rate to use for transactions.
   type FeeRate: Clone + Copy;
 
-  /// The type representing a planned transaction.
-  type PlannedTransaction;
+  /// The type representing a signable transaction.
+  type SignableTransaction: SignableTransaction;
 
   /// Obtain the fee rate to pay.
   ///
@@ -62,7 +73,7 @@ pub trait TransactionPlanner<S: ScannerFeed>: 'static + Send + Sync {
     inputs: Vec<OutputFor<S>>,
     payments: Vec<Payment<AddressFor<S>>>,
     change: Option<KeyFor<S>>,
-  ) -> Self::PlannedTransaction;
+  ) -> PlannedTransaction<S, Self::SignableTransaction, A>;
 
   /// Obtain a PlannedTransaction via amortizing the fee over the payments.
   ///
@@ -77,7 +88,7 @@ pub trait TransactionPlanner<S: ScannerFeed>: 'static + Send + Sync {
     inputs: Vec<OutputFor<S>>,
     mut payments: Vec<Payment<AddressFor<S>>>,
     mut change: Option<KeyFor<S>>,
-  ) -> Option<Self::PlannedTransaction> {
+  ) -> Option<PlannedTransaction<S, Self::SignableTransaction, A>> {
     // Sanity checks
     {
       assert!(!inputs.is_empty());

processor/scheduler/utxo/transaction-chaining/Cargo.toml

@@ -30,6 +30,6 @@ serai-primitives = { path = "../../../../substrate/primitives", default-features
 serai-db = { path = "../../../../common/db" }
 
 primitives = { package = "serai-processor-primitives", path = "../../../primitives" }
+scanner = { package = "serai-processor-scanner", path = "../../../scanner" }
 scheduler-primitives = { package = "serai-processor-scheduler-primitives", path = "../../primitives" }
 utxo-scheduler-primitives = { package = "serai-processor-utxo-scheduler-primitives", path = "../primitives" }
-scanner = { package = "serai-processor-scanner", path = "../../../scanner" }

processor/scheduler/utxo/transaction-chaining/src/db.rs

@@ -6,8 +6,8 @@ use serai_primitives::{Coin, Amount};
 
 use serai_db::{Get, DbTxn, create_db};
 
-use primitives::ReceivedOutput;
-use scanner::{ScannerFeed, KeyFor, OutputFor};
+use primitives::{Payment, ReceivedOutput};
+use scanner::{ScannerFeed, KeyFor, AddressFor, OutputFor};
 
 create_db! {
   TransactionChainingScheduler {
@@ -15,7 +15,7 @@ create_db! {
     SerializedOutputs: (key: &[u8], coin: Coin) -> Vec<u8>,
     // We should be immediately able to schedule the fulfillment of payments, yet this may not be
     // possible if we're in the middle of a multisig rotation (as our output set will be split)
-    SerializedQueuedPayments: (key: &[u8]) > Vec<u8>,
+    SerializedQueuedPayments: (key: &[u8], coin: Coin) -> Vec<u8>,
   }
 }
 
@@ -61,13 +61,19 @@ impl<S: ScannerFeed> Db<S> {
   pub(crate) fn queued_payments(
     getter: &impl Get,
     key: KeyFor<S>,
-  ) -> Option<Vec<Payment<S>>> {
+    coin: Coin,
+  ) -> Option<Vec<Payment<AddressFor<S>>>> {
     todo!("TODO")
   }
-  pub(crate) fn set_queued_payments(txn: &mut impl DbTxn, key: KeyFor<S>, queued: Vec<Payment<S>>) {
+  pub(crate) fn set_queued_payments(
+    txn: &mut impl DbTxn,
+    key: KeyFor<S>,
+    coin: Coin,
+    queued: &Vec<Payment<AddressFor<S>>>,
+  ) {
     todo!("TODO")
   }
-  pub(crate) fn del_outputs(txn: &mut impl DbTxn, key: KeyFor<S>) {
-    SerializedQueuedPayments::del(txn, key.to_bytes().as_ref());
+  pub(crate) fn del_queued_payments(txn: &mut impl DbTxn, key: KeyFor<S>, coin: Coin) {
+    SerializedQueuedPayments::del(txn, key.to_bytes().as_ref(), coin);
   }
 }

processor/scheduler/utxo/transaction-chaining/src/lib.rs

@@ -7,11 +7,11 @@ use std::collections::HashMap;
 
 use group::GroupEncoding;
 
-use serai_primitives::Coin;
+use serai_primitives::{Coin, Amount};
 
 use serai_db::DbTxn;
 
-use primitives::{ReceivedOutput, Payment};
+use primitives::{OutputType, ReceivedOutput, Payment};
 use scanner::{
   LifetimeStage, ScannerFeed, KeyFor, AddressFor, OutputFor, EventualityFor, SchedulerUpdate,
   Scheduler as SchedulerTrait,
@@ -22,65 +22,205 @@ use utxo_scheduler_primitives::*;
 mod db;
 use db::Db;
 
-/// A planned transaction.
-pub struct PlannedTransaction<S: ScannerFeed, T> {
-  /// The signable transaction.
-  signable: T,
-  /// The outputs we'll receive from this.
-  effected_received_outputs: OutputFor<S>,
-  /// The Eventuality to watch for.
-  eventuality: EventualityFor<S>,
-}
+/// The outputs which will be effected by a PlannedTransaction and received by Serai.
+pub struct EffectedReceivedOutputs<S: ScannerFeed>(Vec<OutputFor<S>>);
 
 /// A scheduler of transactions for networks premised on the UTXO model which support
 /// transaction chaining.
-pub struct Scheduler<
-  S: ScannerFeed,
-  T,
-  P: TransactionPlanner<S, PlannedTransaction = PlannedTransaction<S, T>>,
->(PhantomData<S>, PhantomData<T>, PhantomData<P>);
+pub struct Scheduler<S: ScannerFeed, P: TransactionPlanner<S, EffectedReceivedOutputs<S>>>(
+  PhantomData<S>,
+  PhantomData<P>,
+);
 
-impl<S: ScannerFeed, T, P: TransactionPlanner<S, PlannedTransaction = PlannedTransaction<S, T>>>
-  Scheduler<S, T, P>
-{
-  fn accumulate_outputs(txn: &mut impl DbTxn, key: KeyFor<S>, outputs: &[OutputFor<S>]) {
-    // Accumulate them in memory
-    let mut outputs_by_coin = HashMap::with_capacity(1);
-    for output in outputs.iter().filter(|output| output.key() == key) {
-      let coin = output.balance().coin;
-      if let std::collections::hash_map::Entry::Vacant(e) = outputs_by_coin.entry(coin) {
-        e.insert(Db::<S>::outputs(txn, key, coin).unwrap());
+impl<S: ScannerFeed, P: TransactionPlanner<S, EffectedReceivedOutputs<S>>> Scheduler<S, P> {
+  fn handle_queued_payments(
+    &mut self,
+    txn: &mut impl DbTxn,
+    active_keys: &[(KeyFor<S>, LifetimeStage)],
+    key: KeyFor<S>,
+  ) -> Vec<EventualityFor<S>> {
+    let mut eventualities = vec![];
+
+    for coin in S::NETWORK.coins() {
+      // Fetch our operating costs and all our outputs
+      let mut operating_costs = Db::<S>::operating_costs(txn, *coin).0;
+      let mut outputs = Db::<S>::outputs(txn, key, *coin).unwrap();
+
+      // Fetch the queued payments
+      let mut payments = Db::<S>::queued_payments(txn, key, *coin).unwrap();
+      if payments.is_empty() {
+        continue;
       }
-      outputs_by_coin.get_mut(&coin).unwrap().push(output.clone());
+
+      // If this is our only key, our outputs and operating costs should be greater than the
+      // payments' value
+      if active_keys.len() == 1 {
+        // The available amount of fulfill is the amount we have plus the amount we'll reduce by
+        // An alternative formulation would be `outputs >= (payments - operating costs)`, but
+        // that'd risk underflow
+        let available =
+          operating_costs + outputs.iter().map(|output| output.balance().amount.0).sum::<u64>();
+        assert!(available >= payments.iter().map(|payment| payment.balance().amount.0).sum::<u64>());
+      }
+
+      let amount_of_payments_that_can_be_handled =
+        |operating_costs: u64, outputs: &[_], payments: &[_]| {
+          let value_available =
+            operating_costs + outputs.iter().map(|output| output.balance().amount.0).sum::<u64>();
+
+          let mut can_handle = 0;
+          let mut value_used = 0;
+          for payment in payments {
+            value_used += payment.balance().amount.0;
+            if value_available < value_used {
+              break;
+            }
+            can_handle += 1;
+          }
+
+          can_handle
+        };
+
+      // Find the set of payments we should fulfill at this time
+      {
+        // Drop to just the payments we currently have the outputs for
+        {
+          let can_handle =
+            amount_of_payments_that_can_be_handled(operating_costs, &outputs, &payments);
+          let remaining_payments = payments.drain(can_handle ..).collect::<Vec<_>>();
+          // Restore the rest to the database
+          Db::<S>::set_queued_payments(txn, key, *coin, &remaining_payments);
+        }
+        let payments_value = payments.iter().map(|payment| payment.balance().amount.0).sum::<u64>();
+
+        // If these payments are worth less than the operating costs, immediately drop them
+        if payments_value <= operating_costs {
+          operating_costs -= payments_value;
+          Db::<S>::set_operating_costs(txn, *coin, Amount(operating_costs));
+          return vec![];
+        }
+
+        // We explicitly sort AFTER deciding which payments to handle so we always handle the
+        // oldest queued payments first (preventing any from eternally being shuffled to the back
+        // of the line)
+        payments.sort_by(|a, b| a.balance().amount.0.cmp(&b.balance().amount.0));
+      }
+      assert!(!payments.is_empty());
+
+      // Find the smallest set of outputs usable to fulfill these outputs
+      // Size is determined by the largest output, not quantity nor aggregate value
+      {
+        // We start by sorting low to high
+        outputs.sort_by(|a, b| a.balance().amount.0.cmp(&b.balance().amount.0));
+
+        let value_needed =
+          payments.iter().map(|payment| payment.balance().amount.0).sum::<u64>() - operating_costs;
+
+        let mut needed = 0;
+        let mut value_present = 0;
+        for output in &outputs {
+          needed += 1;
+          value_present += output.balance().amount.0;
+          if value_present >= value_needed {
+            break;
+          }
+        }
+
+        // Drain, and save back to the DB, the unnecessary outputs
+        let remaining_outputs = outputs.drain(needed ..).collect::<Vec<_>>();
+        Db::<S>::set_outputs(txn, key, *coin, &remaining_outputs);
+      }
+      assert!(!outputs.is_empty());
+
+      // We now have the current operating costs, the outputs we're using, and the payments
+      // The database has the unused outputs/unfilfillable payments
+      // Actually plan/send off the transactions
+
+      // While our set of outputs exceed the input limit, aggregate them
+      while outputs.len() > MAX_INPUTS {
+        let outputs_chunk = outputs.drain(.. MAX_INPUTS).collect::<Vec<_>>();
+
+        // While we're aggregating these outputs, handle any payments we can
+        let payments_chunk = loop {
+          let can_handle =
+            amount_of_payments_that_can_be_handled(operating_costs, &outputs, &payments);
+          let payments_chunk = payments.drain(.. can_handle.min(MAX_OUTPUTS)).collect::<Vec<_>>();
+
+          let payments_value =
+            payments_chunk.iter().map(|payment| payment.balance().amount.0).sum::<u64>();
+          if payments_value <= operating_costs {
+            operating_costs -= payments_value;
+            continue;
+          }
+          break payments_chunk;
+        };
+
+        let Some(planned) = P::plan_transaction_with_fee_amortization(
+          &mut operating_costs,
+          fee_rates[coin],
+          outputs_chunk,
+          payments_chunk,
+          // We always use our key for the change here since we may need this change output to
+          // finish fulfilling these payments
+          Some(key),
+        ) else {
+          // We amortized all payments, and even when just trying to make the change output, these
+          // inputs couldn't afford their own aggregation and were written off
+          continue;
+        };
+
+        // Send the transactions off for signing
+        TransactionsToSign::<P::SignableTransaction>::send(txn, &key, &planned.signable);
+
+        // Push the Eventualities onto the result
+        eventualities.push(planned.eventuality);
+
+        let mut effected_received_outputs = planned.auxilliary.0;
+        // Only handle Change so if someone burns to an External address, we don't use it here
+        // when the scanner will tell us to return it (without accumulating it)
+        effected_received_outputs.retain(|output| output.kind() == OutputType::Change);
+        outputs.append(&mut effected_received_outputs);
+      }
+
+      // Now that we have an aggregated set of inputs, create the tree for payments
+      todo!("TODO");
     }
 
-    // Flush them to the database
-    for (coin, outputs) in outputs_by_coin {
-      Db::<S>::set_outputs(txn, key, coin, &outputs);
-    }
+    eventualities
   }
 }
 
-impl<
-    S: ScannerFeed,
-    T: 'static + Send + Sync + SignableTransaction,
-    P: TransactionPlanner<S, PlannedTransaction = PlannedTransaction<S, T>>,
-  > SchedulerTrait<S> for Scheduler<S, T, P>
+impl<S: ScannerFeed, P: TransactionPlanner<S, EffectedReceivedOutputs<S>>> SchedulerTrait<S>
+  for Scheduler<S, P>
 {
   fn activate_key(&mut self, txn: &mut impl DbTxn, key: KeyFor<S>) {
     for coin in S::NETWORK.coins() {
+      assert!(Db::<S>::outputs(txn, key, *coin).is_none());
       Db::<S>::set_outputs(txn, key, *coin, &[]);
+      assert!(Db::<S>::queued_payments(txn, key, *coin).is_none());
+      Db::<S>::set_queued_payments(txn, key, *coin, &vec![]);
     }
   }
 
   fn flush_key(&mut self, txn: &mut impl DbTxn, retiring_key: KeyFor<S>, new_key: KeyFor<S>) {
-    todo!("TODO")
+    for coin in S::NETWORK.coins() {
+      let still_queued = Db::<S>::queued_payments(txn, retiring_key, *coin).unwrap();
+      let mut new_queued = Db::<S>::queued_payments(txn, new_key, *coin).unwrap();
+
+      let mut queued = still_queued;
+      queued.append(&mut new_queued);
+
+      Db::<S>::set_queued_payments(txn, retiring_key, *coin, &vec![]);
+      Db::<S>::set_queued_payments(txn, new_key, *coin, &queued);
+    }
   }
 
   fn retire_key(&mut self, txn: &mut impl DbTxn, key: KeyFor<S>) {
     for coin in S::NETWORK.coins() {
-      assert!(Db::<S>::outputs(txn, key, *coin).is_none());
+      assert!(Db::<S>::outputs(txn, key, *coin).unwrap().is_empty());
       Db::<S>::del_outputs(txn, key, *coin);
+      assert!(Db::<S>::queued_payments(txn, key, *coin).unwrap().is_empty());
+      Db::<S>::del_queued_payments(txn, key, *coin);
     }
   }
 
@@ -91,12 +231,41 @@ impl<
     update: SchedulerUpdate<S>,
   ) -> HashMap<Vec<u8>, Vec<EventualityFor<S>>> {
     // Accumulate all the outputs
-    for key in active_keys {
-      Self::accumulate_outputs(txn, key.0, update.outputs());
+    for (key, _) in active_keys {
+      // Accumulate them in memory
+      let mut outputs_by_coin = HashMap::with_capacity(1);
+      for output in update.outputs().iter().filter(|output| output.key() == *key) {
+        match output.kind() {
+          OutputType::External | OutputType::Forwarded => {},
+          // TODO: Only accumulate these if we haven't already, but do accumulate if not
+          OutputType::Branch | OutputType::Change => todo!("TODO"),
+        }
+        let coin = output.balance().coin;
+        if let std::collections::hash_map::Entry::Vacant(e) = outputs_by_coin.entry(coin) {
+          e.insert(Db::<S>::outputs(txn, *key, coin).unwrap());
+        }
+        outputs_by_coin.get_mut(&coin).unwrap().push(output.clone());
+      }
+
+      // Flush them to the database
+      for (coin, outputs) in outputs_by_coin {
+        Db::<S>::set_outputs(txn, *key, coin, &outputs);
+      }
     }
 
     let mut fee_rates: HashMap<Coin, _> = todo!("TODO");
 
+    // Fulfill the payments we prior couldn't
+    let mut eventualities = HashMap::new();
+    for (key, _stage) in active_keys {
+      eventualities.insert(
+        key.to_bytes().as_ref().to_vec(),
+        self.handle_queued_payments(txn, active_keys, *key),
+      );
+    }
+
+    // TODO: If this key has been flushed, forward all outputs
+
     // Create the transactions for the forwards/burns
     {
       let mut planned_txs = vec![];
@@ -137,20 +306,14 @@ impl<
         planned_txs.push((key, plan));
       }
 
-      let mut eventualities = HashMap::new();
       for (key, planned_tx) in planned_txs {
         // Send the transactions off for signing
-        TransactionsToSign::<T>::send(txn, &key, &planned_tx.signable);
+        TransactionsToSign::<P::SignableTransaction>::send(txn, &key, &planned_tx.signable);
 
-        // Insert the eventualities into the result
-        eventualities
-          .entry(key.to_bytes().as_ref().to_vec())
-          .or_insert(Vec::with_capacity(1))
-          .push(planned_tx.eventuality);
+        // Insert the Eventualities into the result
+        eventualities[key.to_bytes().as_ref()].push(planned_tx.eventuality);
       }
 
-      // TODO: Fulfill any payments we prior couldn't
-
       eventualities
     }
   }
@@ -159,13 +322,29 @@ impl<
     &mut self,
     txn: &mut impl DbTxn,
     active_keys: &[(KeyFor<S>, LifetimeStage)],
-    payments: Vec<Payment<AddressFor<S>>>,
+    mut payments: Vec<Payment<AddressFor<S>>>,
   ) -> HashMap<Vec<u8>, Vec<EventualityFor<S>>> {
-    // TODO: Find the key to use for fulfillment
-    // TODO: Sort outputs and payments by amount
-    // TODO: For as long as we don't have sufficiently aggregated inputs to handle all payments,
-    // aggregate
-    // TODO: Create the tree for the payments
-    todo!("TODO")
+    // Find the key to filfill these payments with
+    let fulfillment_key = match active_keys[0].1 {
+      LifetimeStage::ActiveYetNotReporting => {
+        panic!("expected to fulfill payments despite not reporting for the oldest key")
+      }
+      LifetimeStage::Active | LifetimeStage::UsingNewForChange => active_keys[0].0,
+      LifetimeStage::Forwarding | LifetimeStage::Finishing => active_keys[1].0,
+    };
+
+    // Queue the payments for this key
+    for coin in S::NETWORK.coins() {
+      let mut queued_payments = Db::<S>::queued_payments(txn, fulfillment_key, *coin).unwrap();
+      queued_payments
+        .extend(payments.iter().filter(|payment| payment.balance().coin == *coin).cloned());
+      Db::<S>::set_queued_payments(txn, fulfillment_key, *coin, &queued_payments);
+    }
+
+    // Handle the queued payments
+    HashMap::from([(
+      fulfillment_key.to_bytes().as_ref().to_vec(),
+      self.handle_queued_payments(txn, active_keys, fulfillment_key),
+    )])
   }
 }