From 9adefa4c2c047ce09d6783be0393d8fca99629c0 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sat, 26 Aug 2023 21:09:57 -0400
Subject: [PATCH] Add code to handle a race condition around first_preprocess

---
 coordinator/src/db.rs   |  4 ++--
 coordinator/src/main.rs | 18 ++++++++++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/coordinator/src/db.rs b/coordinator/src/db.rs
index aac6a513..a7edd3cc 100644
--- a/coordinator/src/db.rs
+++ b/coordinator/src/db.rs
@@ -92,7 +92,7 @@ impl<'a, D: Db> MainDb<'a, D> {
     }
     txn.put(key, preprocess);
   }
-  pub fn first_preprocess<G: Get>(getter: &G, id: [u8; 32]) -> Vec<u8> {
-    getter.get(Self::first_preprocess_key(id)).expect("asked for first preprocess we never saved")
+  pub fn first_preprocess<G: Get>(getter: &G, id: [u8; 32]) -> Option<Vec<u8>> {
+    getter.get(Self::first_preprocess_key(id))
   }
 }
diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 221baf7f..e9a89e33 100644
--- a/coordinator/src/main.rs
+++ b/coordinator/src/main.rs
@@ -694,6 +694,20 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
       let key = key.clone();
       let tributaries = tributaries.clone();
       async move {
+        // SubstrateBlockAck is fired before Preprocess, creating a race between Tributary ack
+        // of the SubstrateBlock and the sending of all Preprocesses
+        // A similar race condition exists when multiple Batches are present in a block
+        // This waits until the necessary preprocess is available
+        let get_preprocess = |raw_db, id| async move {
+          loop {
+            let Some(preprocess) = MainDb::<D>::first_preprocess(raw_db, id) else {
+              sleep(Duration::from_millis(100)).await;
+              continue;
+            };
+            return preprocess;
+          }
+        };
+
         let (ids, txs) = match id_type {
           RecognizedIdType::Block => {
             let block = id;
@@ -704,7 +718,7 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
               txs.push(Transaction::BatchPreprocess(SignData {
                 plan: *id,
                 attempt: 0,
-                data: MainDb::<D>::first_preprocess(&raw_db, *id),
+                data: get_preprocess(&raw_db, *id).await,
                 signed: Transaction::empty_signed(),
               }));
             }
@@ -716,7 +730,7 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
             vec![Transaction::SignPreprocess(SignData {
               plan: id,
               attempt: 0,
-              data: MainDb::<D>::first_preprocess(&raw_db, id),
+              data: get_preprocess(&raw_db, id).await,
               signed: Transaction::empty_signed(),
             })],
           ),