Add key management to the scheduler

processor/scanner/src/eventuality/db.rs

@@ -1,12 +1,17 @@
 use core::marker::PhantomData;
 
 use scale::Encode;
+use borsh::{BorshSerialize, BorshDeserialize};
 use serai_db::{Get, DbTxn, create_db};
 
 use primitives::{EncodableG, Eventuality, EventualityTracker};
 
 use crate::{ScannerFeed, KeyFor, EventualityFor};
 
+// The DB macro doesn't support `BorshSerialize + BorshDeserialize` as a bound, hence this.
+trait Borshy: BorshSerialize + BorshDeserialize {}
+impl<T: BorshSerialize + BorshDeserialize> Borshy for T {}
+
 create_db!(
   ScannerEventuality {
     // The next block to check for resolving eventualities
@@ -15,6 +20,8 @@ create_db!(
     LatestHandledNotableBlock: () -> u64,
 
     SerializedEventualities: <K: Encode>(key: K) -> Vec<u8>,
+
+    RetiredKey: <K: Borshy>(block_number: u64) -> K,
   }
 );
 
@@ -51,7 +58,6 @@ impl<S: ScannerFeed> EventualityDb<S> {
     }
     SerializedEventualities::set(txn, EncodableG(key), &serialized);
   }
-
   pub(crate) fn eventualities(
     getter: &impl Get,
     key: KeyFor<S>,
@@ -66,4 +72,19 @@ impl<S: ScannerFeed> EventualityDb<S> {
     }
     res
   }
+
+  pub(crate) fn retire_key(txn: &mut impl DbTxn, block_number: u64, key: KeyFor<S>) {
+    assert!(
+      RetiredKey::get::<EncodableG<KeyFor<S>>>(txn, block_number).is_none(),
+      "retiring multiple keys within the same block"
+    );
+    RetiredKey::set(txn, block_number, &EncodableG(key));
+  }
+  pub(crate) fn take_retired_key(txn: &mut impl DbTxn, block_number: u64) -> Option<KeyFor<S>> {
+    let res = RetiredKey::get::<EncodableG<KeyFor<S>>>(txn, block_number).map(|res| res.0);
+    if res.is_some() {
+      RetiredKey::del::<EncodableG<KeyFor<S>>>(txn, block_number);
+    }
+    res
+  }
 }

processor/scanner/src/eventuality/mod.rs

@@ -248,6 +248,11 @@ impl<D: Db, S: ScannerFeed, Sch: Scheduler<S>> ContinuallyRan for EventualityTas
       let mut outputs = received_external_outputs;
 
       for key in &keys {
+        // If this is the key's activation block, activate it
+        if key.activation_block_number == b {
+          self.scheduler.activate_key(&mut txn, key.key);
+        }
+
         let completed_eventualities = {
           let mut eventualities = EventualityDb::<S>::eventualities(&txn, key.key);
           let completed_eventualities = block.check_for_eventuality_resolutions(&mut eventualities);
@@ -349,11 +354,18 @@ impl<D: Db, S: ScannerFeed, Sch: Scheduler<S>> ContinuallyRan for EventualityTas
 
             // Retire this key `WINDOW_LENGTH` blocks in the future to ensure the scan task never
             // has a malleable view of the keys.
-            ScannerGlobalDb::<S>::retire_key(&mut txn, b + S::WINDOW_LENGTH, key.key);
+            let retire_at = b + S::WINDOW_LENGTH;
+            ScannerGlobalDb::<S>::retire_key(&mut txn, retire_at, key.key);
+            EventualityDb::<S>::retire_key(&mut txn, retire_at, key.key);
           }
         }
       }
 
+      // If we retired any key at this block, retire it within the scheduler
+      if let Some(key) = EventualityDb::<S>::take_retired_key(&mut txn, b) {
+        self.scheduler.retire_key(&mut txn, key);
+      }
+
       // Update the next-to-check block
       EventualityDb::<S>::set_next_to_check_for_eventualities_block(&mut txn, next_to_check);