From a691be21c84fe3ce7d69f461b651e8df7349fa2b Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Thu, 19 Sep 2024 01:05:36 -0400
Subject: [PATCH] Call tidy_keys upon queue_key

Prevents the potential case of the substrate task and the scan task writing to
the same storage slot at once.
---
 processor/scanner/src/db.rs       | 46 +++++++++++++++++--------------
 processor/scanner/src/scan/mod.rs |  3 --
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/processor/scanner/src/db.rs b/processor/scanner/src/db.rs
index 884e0e2b..a985ba43 100644
--- a/processor/scanner/src/db.rs
+++ b/processor/scanner/src/db.rs
@@ -116,6 +116,28 @@ impl<S: ScannerFeed> ScannerGlobalDb<S> {
     StartBlock::set(txn, &block)
   }
 
+  fn tidy_keys(txn: &mut impl DbTxn) {
+    let mut keys: Vec<SeraiKeyDbEntry<EncodableG<KeyFor<S>>>> =
+      ActiveKeys::get(txn).expect("retiring key yet no active keys");
+    let Some(key) = keys.first() else { return };
+
+    // Get the block we're scanning for next
+    let block_number = next_to_scan_for_outputs_block::<S>(txn).expect(
+      "tidying keys despite never setting the next to scan for block (done on initialization)",
+    );
+    // If this key is scheduled for retiry...
+    if let Some(retire_at) = RetireAt::get(txn, key.key) {
+      // And is retired by/at this block...
+      if retire_at <= block_number {
+        // Remove it from the list of keys
+        let key = keys.remove(0);
+        ActiveKeys::set(txn, &keys);
+        // Also clean up the retiry block
+        RetireAt::del(txn, key.key);
+      }
+    }
+  }
+
   /// Queue a key.
   ///
   /// Keys may be queued whenever, so long as they're scheduled to activate `WINDOW_LENGTH` blocks
@@ -165,6 +187,9 @@ impl<S: ScannerFeed> ScannerGlobalDb<S> {
     // Push and save the next key
     keys.push(SeraiKeyDbEntry { activation_block_number, key: EncodableG(key) });
     ActiveKeys::set(txn, &keys);
+
+    // Now tidy the keys, ensuring this has a maximum length of 2
+    Self::tidy_keys(txn);
   }
   /// Retire a key.
   ///
@@ -181,27 +206,6 @@ impl<S: ScannerFeed> ScannerGlobalDb<S> {
 
     RetireAt::set(txn, EncodableG(key), &at_block);
   }
-  pub(crate) fn tidy_keys(txn: &mut impl DbTxn) {
-    let mut keys: Vec<SeraiKeyDbEntry<EncodableG<KeyFor<S>>>> =
-      ActiveKeys::get(txn).expect("retiring key yet no active keys");
-    let Some(key) = keys.first() else { return };
-
-    // Get the block we're scanning for next
-    let block_number = next_to_scan_for_outputs_block::<S>(txn).expect(
-      "tidying keys despite never setting the next to scan for block (done on initialization)",
-    );
-    // If this key is scheduled for retiry...
-    if let Some(retire_at) = RetireAt::get(txn, key.key) {
-      // And is retired by/at this block...
-      if retire_at <= block_number {
-        // Remove it from the list of keys
-        let key = keys.remove(0);
-        ActiveKeys::set(txn, &keys);
-        // Also clean up the retiry block
-        RetireAt::del(txn, key.key);
-      }
-    }
-  }
   /// Fetch the active keys, as of the next-to-scan-for-outputs Block.
   ///
   /// This means the scan task should scan for all keys returned by this.
diff --git a/processor/scanner/src/scan/mod.rs b/processor/scanner/src/scan/mod.rs
index 7004a4d9..0ebdf992 100644
--- a/processor/scanner/src/scan/mod.rs
+++ b/processor/scanner/src/scan/mod.rs
@@ -116,9 +116,6 @@ impl<D: Db, S: ScannerFeed> ContinuallyRan for ScanTask<D, S> {
 
         assert_eq!(ScanDb::<S>::next_to_scan_for_outputs_block(&txn).unwrap(), b);
 
-        // Tidy the keys, then fetch them
-        // We don't have to tidy them here, we just have to somewhere, so why not here?
-        ScannerGlobalDb::<S>::tidy_keys(&mut txn);
         let keys = ScannerGlobalDb::<S>::active_keys_as_of_next_to_scan_for_outputs_block(&txn)
           .expect("scanning for a blockchain without any keys set");