BM-273: Fixes batch deadline detection by regular polling for finalize conditions (github#87)

This PR moves the finalize checks to a function that can run
incrementally as we poll for new orders such that batch deadlines can be
detected and finalized without new orders coming into the aggregator to
trigger it.

It also adds a config option for how often to the aggregator should
check for new proofs to aggregate / finalize checks.

Additionally it adds a unit test for batch-deadline detection. 

This PR does increase the get_block_number() RPC calls and means we
should probably make a block number service internally to broker sooner
rather than later.

---------

Co-authored-by: Victor Graf <[email protected]>

Author: mothran

broker.toml

@@ -23,3 +23,4 @@ batch_max_time = 1000
 batch_size = 1
 block_deadline_buffer_secs = 120
 txn_timeout = 45
+# batch_poll_time_ms = 500

crates/broker/src/aggregator.rs

@@ -342,58 +342,18 @@ where
         Ok(())
     }
 
-    async fn aggregate_proofs(&mut self) -> Result<()> {
-        let new_proofs = self
-            .db
-            .get_aggregation_proofs()
-            .await
-            .context("Failed to get pending agg proofs from DB")?;
-
-        if new_proofs.is_empty() {
-            return Ok(());
-        }
-
-        let batch_id = self.db.get_current_batch().await.context("Failed to get current batch")?;
-
-        for agg_proof in new_proofs {
-            match self.aggregate_proof(batch_id, agg_proof.order_id, agg_proof.proof_id).await {
-                Ok(_) => {
-                    tracing::info!("Completed aggregation of proof {:x}", agg_proof.order_id);
-                    self.db
-                        .update_batch(
-                            batch_id,
-                            agg_proof.order_id,
-                            agg_proof.expire_block,
-                            agg_proof.fee,
-                        )
-                        .await
-                        .context("Failed to update batch with new order details")?;
-                }
-                Err(err) => {
-                    tracing::error!(
-                        "Failed to complete aggregation of proof {:x} {err:?}",
-                        agg_proof.order_id
-                    );
-                    if let Err(db_err) =
-                        self.db.set_order_failure(agg_proof.order_id, format!("{err:?}")).await
-                    {
-                        tracing::error!("Failed to mark order failure in db: {db_err}");
-                    }
-                }
-            }
-        }
-
-        // Check if we should finalize the batch
-        // - check current min-deadline and batch timer
-        // - need to fetch current block, might be good to make that a long polling service with a
-        //   Atomic everyone reads
-        // if so:
-        // - finalize
-        // - snark proof
-        // - insert batch data in to DB for finalizer
-        // - mark all orders in batch as Aggregated
-        // - reset batch timer
-
+    /// Check if we should finalize the batch
+    ///
+    /// - check current min-deadline and batch timer
+    /// - need to fetch current block, might be good to make that a long polling service with a
+    ///   Atomic everyone reads
+    ///
+    /// if so:
+    /// - finalize
+    /// - snark proof
+    /// - insert batch data in to DB for finalizer
+    /// - mark all orders in batch as Aggregated
+    async fn check_finalize(&mut self, batch_id: usize) -> Result<()> {
         let (conf_batch_size, conf_batch_time, conf_batch_fees) = {
             let config = self.config.lock_all().context("Failed to lock config")?;
 
@@ -457,11 +417,19 @@ where
             let block_number =
                 self.provider.get_block_number().await.context("Failed to get current block")?;
 
-            // tracing::info!("{} {} {}", self.current_deadline, block_number, self.block_time);
             let remaining_secs = (batch.block_deadline.expect("batch missing block deadline")
                 - block_number)
                 * self.block_time;
             let buffer_secs = conf_block_deadline_buf;
+            // tracing::info!(
+            //     "{:?} {} {} {} {}",
+            //     batch.block_deadline,
+            //     block_number,
+            //     self.block_time,
+            //     remaining_secs,
+            //     buffer_secs
+            // );
+
             if remaining_secs <= buffer_secs {
                 tracing::info!("Batch getting close to deadline {remaining_secs}, finalizing");
                 finalize = true;
@@ -474,6 +442,54 @@ where
 
         Ok(())
     }
+
+    async fn aggregate_proofs(&mut self) -> Result<()> {
+        let new_proofs = self
+            .db
+            .get_aggregation_proofs()
+            .await
+            .context("Failed to get pending agg proofs from DB")?;
+
+        let batch_id = self.db.get_current_batch().await.context("Failed to get current batch")?;
+
+        self.check_finalize(batch_id).await?;
+
+        if new_proofs.is_empty() {
+            return Ok(());
+        }
+
+        for agg_proof in new_proofs {
+            match self.aggregate_proof(batch_id, agg_proof.order_id, agg_proof.proof_id).await {
+                Ok(_) => {
+                    tracing::info!("Completed aggregation of proof {:x}", agg_proof.order_id);
+                    self.db
+                        .update_batch(
+                            batch_id,
+                            agg_proof.order_id,
+                            agg_proof.expire_block,
+                            agg_proof.fee,
+                        )
+                        .await
+                        .context("Failed to update batch with new order details")?;
+                }
+                Err(err) => {
+                    tracing::error!(
+                        "Failed to complete aggregation of proof {:x} {err:?}",
+                        agg_proof.order_id
+                    );
+                    if let Err(db_err) =
+                        self.db.set_order_failure(agg_proof.order_id, format!("{err:?}")).await
+                    {
+                        tracing::error!("Failed to mark order failure in db: {db_err}");
+                    }
+                }
+            }
+        }
+
+        self.check_finalize(batch_id).await?;
+
+        Ok(())
+    }
 }
 
 impl<T, P> RetryTask for AggregatorService<T, P>
@@ -483,12 +499,21 @@ where
 {
     fn spawn(&self) -> RetryRes {
         let mut self_clone = self.clone();
+
         Box::pin(async move {
             tracing::info!("Starting Aggregator service");
             loop {
+                let conf_poll_time_ms = {
+                    let config = self_clone
+                        .config
+                        .lock_all()
+                        .context("Failed to lock config")
+                        .map_err(SupervisorErr::Fault)?;
+                    config.batcher.batch_poll_time_ms.unwrap_or(1000)
+                };
+
                 self_clone.aggregate_proofs().await.map_err(SupervisorErr::Recover)?;
-                // TODO: configuration
-                tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+                tokio::time::sleep(tokio::time::Duration::from_millis(conf_poll_time_ms)).await;
             }
         })
     }
@@ -507,7 +532,7 @@ mod tests {
         network::EthereumWallet,
         node_bindings::Anvil,
         primitives::{aliases::U96, Keccak256, B256},
-        providers::ProviderBuilder,
+        providers::{ext::AnvilApi, ProviderBuilder},
         signers::local::PrivateKeySigner,
     };
     use boundless_market::contracts::{
@@ -838,4 +863,112 @@ mod tests {
         assert!(!batch.orders.is_empty());
         assert_eq!(batch.status, BatchStatus::PendingSubmission);
     }
+
+    #[tokio::test]
+    #[traced_test]
+    async fn deadline_finalize() {
+        let anvil = Anvil::new().spawn();
+        let signer: PrivateKeySigner = anvil.keys()[0].clone().into();
+        let provider = Arc::new(
+            ProviderBuilder::new()
+                .with_recommended_fillers()
+                .wallet(EthereumWallet::from(signer))
+                .on_http(anvil.endpoint().parse().unwrap()),
+        );
+        let db: DbObj = Arc::new(SqliteDb::new("sqlite::memory:").await.unwrap());
+        let config = ConfigLock::default();
+        {
+            let mut config = config.load_write().unwrap();
+            config.batcher.batch_size = Some(2);
+            config.batcher.block_deadline_buffer_secs = 100;
+        }
+
+        let prover: ProverObj = Arc::new(MockProver::default());
+
+        // Pre-prove the echo aka app guest:
+        let image_id = Digest::from(ECHO_ID);
+        let image_id_str = image_id.to_string();
+        prover.upload_image(&image_id_str, ECHO_ELF.to_vec()).await.unwrap();
+        let input_id = prover
+            .upload_input(encode_input(&vec![0x41, 0x41, 0x41, 0x41]).unwrap())
+            .await
+            .unwrap();
+        let proof_res =
+            prover.prove_and_monitor_stark(&image_id_str, &input_id, vec![]).await.unwrap();
+
+        let mut aggregator = AggregatorService::new(
+            db.clone(),
+            provider.clone(),
+            Digest::from(SET_BUILDER_GUEST_ID),
+            SET_BUILDER_GUEST_ELF.to_vec(),
+            Digest::from(ASSESSOR_GUEST_ID),
+            ASSESSOR_GUEST_ELF.to_vec(),
+            Address::ZERO,
+            config,
+            prover,
+            2,
+        )
+        .await
+        .unwrap();
+
+        let customer_signer: PrivateKeySigner = anvil.keys()[1].clone().into();
+        let chain_id = provider.get_chain_id().await.unwrap();
+
+        let min_price = 200000000000000000u64;
+        let order_request = ProvingRequest::new(
+            0,
+            &customer_signer.address(),
+            Requirements {
+                imageId: B256::from_slice(image_id.as_bytes()),
+                predicate: Predicate {
+                    predicateType: PredicateType::PrefixMatch,
+                    data: Default::default(),
+                },
+            },
+            "http://risczero.com/image".into(),
+            Input { inputType: InputType::Inline, data: Default::default() },
+            Offer {
+                minPrice: U96::from(min_price),
+                maxPrice: U96::from(250000000000000000u64),
+                biddingStart: 0,
+                timeout: 50,
+                rampUpPeriod: 1,
+                lockinStake: U96::from(10),
+            },
+        );
+
+        let client_sig = order_request
+            .sign_request(&customer_signer, Address::ZERO, chain_id)
+            .unwrap()
+            .as_bytes();
+
+        let order = Order {
+            status: OrderStatus::PendingAgg,
+            updated_at: Utc::now(),
+            target_block: None,
+            request: order_request,
+            image_id: Some(image_id_str.clone()),
+            input_id: Some(input_id.clone()),
+            proof_id: Some(proof_res.id),
+            expire_block: Some(100),
+            path: None,
+            client_sig: client_sig.into(),
+            lock_price: Some(U256::from(min_price)),
+            error_msg: None,
+        };
+        let order_id = U256::from(order.request.id);
+        db.add_order(order_id, order.clone()).await.unwrap();
+
+        provider.anvil_mine(Some(U256::from(51)), Some(U256::from(2))).await.unwrap();
+
+        aggregator.aggregate_proofs().await.unwrap();
+
+        let db_order = db.get_order(order_id).await.unwrap().unwrap();
+        assert_eq!(db_order.status, OrderStatus::PendingSubmission);
+
+        let (_batch_id, batch) = db.get_complete_batch().await.unwrap().unwrap();
+        assert!(!batch.orders.is_empty());
+        assert_eq!(batch.status, BatchStatus::PendingSubmission);
+        assert!(logs_contain("Batch getting close to deadline"));
+    }
 }

crates/broker/src/config.rs

@@ -123,7 +123,12 @@ pub struct BatcherConfig {
     pub block_deadline_buffer_secs: u64,
     /// Timeout, in seconds for transaction confirmations
     pub txn_timeout: Option<u64>,
-    /// Use the single TXN submission that batchs submit_merkle / fulfill_batch into
+    /// Polling time, in milliseconds
+    ///
+    /// The time between polls for new orders to aggregate and how often to check for
+    /// batch finalize conditions
+    pub batch_poll_time_ms: Option<u64>,
+    /// Use the single TXN submission that batches submit_merkle / fulfill_batch into
     /// A single transaction. Requires the `submitRootAndFulfillBatch` method
     /// be present on the deployed contract
     #[serde(default)]
@@ -138,6 +143,7 @@ impl Default for BatcherConfig {
             batch_max_fees: None,
             block_deadline_buffer_secs: 120,
             txn_timeout: None,
+            batch_poll_time_ms: Some(1000),
             single_txn_fulfill: false,
         }
     }
@@ -346,6 +352,7 @@ batch_max_time = 300
 batch_size = 2
 block_deadline_buffer_secs = 120
 txn_timeout = 45
+batch_poll_time_ms = 1200
 single_txn_fulfill = true"#;
 
     const BAD_CONFIG: &str = r#"
@@ -389,6 +396,7 @@ error = ?"#;
         assert_eq!(config.batcher.batch_max_fees, Some("0.1".into()));
         assert_eq!(config.batcher.block_deadline_buffer_secs, 120);
         assert_eq!(config.batcher.txn_timeout, None);
+        assert_eq!(config.batcher.batch_poll_time_ms, None);
     }
 
     #[tokio::test]
@@ -433,6 +441,7 @@ error = ?"#;
             assert_eq!(config.prover.status_poll_ms, 1000);
             assert!(config.prover.bonsai_r0_zkvm_ver.is_none());
             assert_eq!(config.batcher.txn_timeout, Some(45));
+            assert_eq!(config.batcher.batch_poll_time_ms, Some(1200));
             assert_eq!(config.batcher.single_txn_fulfill, true);
         }
         tracing::debug!("closing...");