BM-51: Image / Input fetch retry system. (github#95)

This PR adds a configurable retry count with static wait time between
fetches. Will allow for brokers to configure a retry counter so IPFS
urls have a chance of being indexed between http GET's

closes BM-51

---------

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

Author: mothran

broker.toml

@@ -7,6 +7,7 @@ lookback_blocks = 100
 max_stake = "0.5"
 skip_preflight_ids = []
 max_file_size = 50_000_000
+# max_fetch_retries = 2
 # allow_client_addresses = []
 # lockin_priority_gas = 100
 

crates/boundless-market/src/storage.rs

@@ -131,7 +131,7 @@ pub enum PinataStorageProviderError {
 }
 
 const DEFAULT_PINATA_API_URL: &str = "https://api.pinata.cloud";
-const DEFAULT_GATEWAY_URL: &str = "https://dweb.link";
+const DEFAULT_GATEWAY_URL: &str = "https://gateway.pinata.cloud";
 
 impl PinataStorageProvider {
     pub async fn from_env() -> Result<Self, PinataStorageProviderError> {

crates/broker/src/config.rs

@@ -51,8 +51,10 @@ pub struct MarketConf {
     /// for increasing the priority if competing with multiple provers during the
     /// same block
     pub lockin_priority_gas: Option<u64>,
-    /// Max input / image file size
+    /// Max input / image file size allowed for downloading from request URLs
     pub max_file_size: usize,
+    /// Max retries for fetching input / image contents from URLs
+    pub max_fetch_retries: Option<u8>,
 }
 
 impl Default for MarketConf {
@@ -68,6 +70,7 @@ impl Default for MarketConf {
             allow_client_addresses: None,
             lockin_priority_gas: None,
             max_file_size: 50_000_000,
+            max_fetch_retries: Some(2),
         }
     }
 }
@@ -330,6 +333,7 @@ lookback_blocks = 100
 max_stake = "0.1"
 skip_preflight_ids = ["0x0000000000000000000000000000000000000000000000000000000000000001"]
 max_file_size = 50_000_000
+max_fetch_retries = 10
 allow_client_addresses = ["0x0000000000000000000000000000000000000000"]
 lockin_priority_gas = 100
 
@@ -425,6 +429,7 @@ error = ?"#;
             assert_eq!(config.market.lookback_blocks, 100);
             assert_eq!(config.market.allow_client_addresses, Some(vec![Address::ZERO]));
             assert_eq!(config.market.lockin_priority_gas, Some(100));
+            assert_eq!(config.market.max_fetch_retries, Some(10));
             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));

crates/broker/src/lib.rs

@@ -533,11 +533,18 @@ where
     }
 }
 
-async fn upload_image_uri(prover: &ProverObj, order: &Order, max_size: usize) -> Result<String> {
-    let uri = UriHandlerBuilder::new(&order.request.imageUrl)
-        .set_max_size(max_size)
-        .build()
-        .context("Uri parse failure")?;
+async fn upload_image_uri(
+    prover: &ProverObj,
+    order: &Order,
+    max_size: usize,
+    retries: Option<u8>,
+) -> Result<String> {
+    let mut uri = UriHandlerBuilder::new(&order.request.imageUrl).set_max_size(max_size);
+
+    if let Some(retry) = retries {
+        uri = uri.set_retries(retry);
+    }
+    let uri = uri.build().context("Uri parse failure")?;
 
     if !uri.exists() {
         let image_data = uri
@@ -566,7 +573,12 @@ async fn upload_image_uri(prover: &ProverObj, order: &Order, max_size: usize) ->
         Ok(uri.id().context("Invalid image URI type")?)
     }
 }
-async fn upload_input_uri(prover: &ProverObj, order: &Order, max_size: usize) -> Result<String> {
+async fn upload_input_uri(
+    prover: &ProverObj,
+    order: &Order,
+    max_size: usize,
+    retries: Option<u8>,
+) -> Result<String> {
     Ok(match order.request.input.inputType {
         InputType::Inline => prover
             .upload_input(order.request.input.data.to_vec())
@@ -577,10 +589,12 @@ async fn upload_input_uri(prover: &ProverObj, order: &Order, max_size: usize) ->
             let input_uri_str =
                 std::str::from_utf8(&order.request.input.data).context("input url is not utf8")?;
             tracing::debug!("Input URI string: {input_uri_str}");
-            let input_uri = UriHandlerBuilder::new(input_uri_str)
-                .set_max_size(max_size)
-                .build()
-                .context("Failed to parse input uri")?;
+            let mut input_uri = UriHandlerBuilder::new(input_uri_str).set_max_size(max_size);
+
+            if let Some(retry) = retries {
+                input_uri = input_uri.set_retries(retry);
+            }
+            let input_uri = input_uri.build().context("Failed to parse input uri")?;
 
             if !input_uri.exists() {
                 let input_data = input_uri

crates/broker/src/order_picker.rs

@@ -125,7 +125,7 @@ where
             return Ok(());
         }
 
-        let (skip_preflight, max_size, peak_prove_khz) = {
+        let (skip_preflight, max_size, peak_prove_khz, fetch_retries) = {
             let config = self.config.lock_all().context("Failed to read config")?;
             let skip_preflight =
                 if let Some(skip_preflights) = config.market.skip_preflight_ids.as_ref() {
@@ -134,7 +134,12 @@ where
                     false
                 };
 
-            (skip_preflight, config.market.max_file_size, config.market.peak_prove_khz)
+            (
+                skip_preflight,
+                config.market.max_file_size,
+                config.market.peak_prove_khz,
+                config.market.max_fetch_retries,
+            )
         };
 
         if skip_preflight {
@@ -147,13 +152,13 @@ where
         }
 
         // TODO: Move URI handling like this into the prover impls
-        let image_id = crate::upload_image_uri(&self.prover, order, max_size)
+        let image_id = crate::upload_image_uri(&self.prover, order, max_size, fetch_retries)
             .await
-            .context("Failed to upload image_id")?;
+            .context("Failed to fetch and upload image_id")?;
 
-        let input_id = crate::upload_input_uri(&self.prover, order, max_size)
+        let input_id = crate::upload_input_uri(&self.prover, order, max_size, fetch_retries)
             .await
-            .context("Failed to upload input_id")?;
+            .context("Failed to fetch and upload input_id")?;
 
         // Record the image/input IDs for proving stage
         self.db

crates/broker/src/proving.rs

@@ -39,22 +39,22 @@ impl ProvingService {
     }
 
     pub async fn prove_order(&self, order_id: U256, order: Order) -> Result<()> {
-        let max_file_size = {
+        let (max_file_size, fetch_retries) = {
             let config = self.config.lock_all().context("Failed to read config")?;
-            config.market.max_file_size
+            (config.market.max_file_size, config.market.max_fetch_retries)
         };
 
         // If the ID's are not present then upload them now
         // Mostly hit by skipping pre-flight
         let image_id = match order.image_id.as_ref() {
             Some(val) => val.clone(),
-            None => crate::upload_image_uri(&self.prover, &order, max_file_size)
+            None => crate::upload_image_uri(&self.prover, &order, max_file_size, fetch_retries)
                 .await
                 .context("Failed to upload image")?,
         };
         let input_id = match order.input_id.as_ref() {
             Some(val) => val.clone(),
-            None => crate::upload_input_uri(&self.prover, &order, max_file_size)
+            None => crate::upload_input_uri(&self.prover, &order, max_file_size, fetch_retries)
                 .await
                 .context("Failed to upload input")?,
         };

crates/broker/src/storage.rs

@@ -31,6 +31,9 @@ pub enum StorageErr {
     #[error("HTTP status error {0}")]
     HttpStatusErr(String),
 
+    #[error("HTTP fetch failed after {0} retries")]
+    FetchRetryMax(u8),
+
     #[error("Authority missing")]
     AuthorityMissing,
 
@@ -45,8 +48,11 @@ pub struct UriHandler {
     uri: url::Url,
     uri_scheme: String,
     max_size: Option<usize>,
+    retries: u8,
 }
 
+const DEFAULT_RETRY_NUMB: u8 = 1;
+
 impl UriHandler {
     fn supported_scheme(scheme: &str) -> bool {
         if risc0_zkvm::is_dev_mode() {
@@ -59,7 +65,11 @@ impl UriHandler {
         matches!(authority, "image" | "input")
     }
 
-    pub fn new(uri_str: &str, max_size: Option<usize>) -> Result<Self, StorageErr> {
+    pub fn new(
+        uri_str: &str,
+        max_size: Option<usize>,
+        retries: Option<u8>,
+    ) -> Result<Self, StorageErr> {
         let uri = url::Url::parse(uri_str)?;
 
         let scheme = uri.scheme().to_string();
@@ -86,7 +96,12 @@ impl UriHandler {
             }
         }
 
-        Ok(Self { uri, uri_scheme: scheme, max_size })
+        Ok(Self {
+            uri,
+            uri_scheme: scheme,
+            max_size,
+            retries: retries.unwrap_or(DEFAULT_RETRY_NUMB),
+        })
     }
 
     pub fn exists(&self) -> bool {
@@ -103,15 +118,29 @@ impl UriHandler {
         match self.uri_scheme.as_ref() {
             "bonsai" => Err(StorageErr::BonsaiFetch),
             "http" | "https" => {
-                let res = reqwest::get(self.uri.to_string()).await?;
-                let status = res.status();
-                if !status.is_success() {
-                    let body = res.text().await?;
-                    return Err(StorageErr::HttpStatusErr(format!(
-                        "HTTP fetch err: {} - {}",
-                        status, body
-                    )));
-                }
+                let mut retry = 0;
+                let res = loop {
+                    // TODO: move these ?'s to captures + retries
+                    // currently only retry on http status code failures
+                    let res = reqwest::get(self.uri.to_string()).await?;
+                    let status = res.status();
+                    if status.is_success() {
+                        break res;
+                    } else {
+                        let body = res.text().await?;
+                        tracing::error!(
+                            "HTTP error fetching contents {retry}/{}: {status} - {body}",
+                            self.retries
+                        );
+                        if retry == self.retries {
+                            return Err(StorageErr::FetchRetryMax(self.retries));
+                        }
+                        retry += 1;
+                        // TODO configurable...
+                        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                        continue;
+                    }
+                };
 
                 let mut buffer = vec![];
                 if let Some(content_length) = res.content_length() {
@@ -163,20 +192,26 @@ impl UriHandler {
 pub struct UriHandlerBuilder {
     uri_str: String,
     max_size: Option<usize>,
+    retries: Option<u8>,
 }
 
 impl UriHandlerBuilder {
     pub fn new(uri_str: &str) -> Self {
-        Self { uri_str: uri_str.into(), max_size: None }
+        Self { uri_str: uri_str.into(), max_size: None, retries: None }
     }
 
     pub fn set_max_size(mut self, max_size: usize) -> Self {
         self.max_size = Some(max_size);
         self
     }
 
+    pub fn set_retries(mut self, retries: u8) -> Self {
+        self.retries = Some(retries);
+        self
+    }
+
     pub fn build(self) -> Result<UriHandler, StorageErr> {
-        UriHandler::new(&self.uri_str, self.max_size)
+        UriHandler::new(&self.uri_str, self.max_size, self.retries)
     }
 }
 
@@ -189,6 +224,8 @@ impl Display for UriHandler {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use httpmock::prelude::*;
+    use tracing_test::traced_test;
 
     #[test]
     fn bonsai_uri_parser() {
@@ -230,22 +267,65 @@ mod tests {
 
     #[tokio::test]
     async fn http_fetch() {
-        let handler = UriHandlerBuilder::new("https://risczero.com/")
-            .set_max_size(1_000_000)
-            .build()
-            .unwrap();
+        let server = MockServer::start();
+        let resp_data = vec![0x41, 0x41, 0x41, 0x41];
+        let get_mock = server.mock(|when, then| {
+            when.method(GET).path("/image");
+            then.status(200).body(&resp_data);
+        });
+
+        let url = format!("http://{}/image", server.address());
+        let handler = UriHandlerBuilder::new(&url).set_max_size(1_000_000).build().unwrap();
+        assert!(!handler.exists());
+
+        let data = handler.fetch().await.unwrap();
+        assert_eq!(data, resp_data);
+        get_mock.assert();
+    }
+
+    #[traced_test]
+    #[tokio::test]
+    async fn http_fetch_retry() {
+        static mut REQ_COUNT: u32 = 0;
+
+        let server = MockServer::start();
+        let get_mock = server.mock(|when, then| {
+            when.method(GET).path("/image").matches(|_req: &HttpMockRequest| {
+                let req = unsafe {
+                    let req = REQ_COUNT;
+                    REQ_COUNT += 1;
+                    req
+                };
+                req >= 1
+            });
+            then.status(200).body("TEST");
+        });
+
+        let url = format!("http://{}/image", server.address());
+        let handler =
+            UriHandlerBuilder::new(&url).set_max_size(1_000_000).set_retries(1).build().unwrap();
         assert!(!handler.exists());
 
         let _data = handler.fetch().await.unwrap();
+        get_mock.assert();
+        assert!(logs_contain("HTTP error fetching contents 0/1"));
     }
 
     #[tokio::test]
     #[should_panic(expected = "TooLarge")]
     async fn max_size_limit() {
-        let handler =
-            UriHandlerBuilder::new("https://risczero.com/").set_max_size(1).build().unwrap();
+        let server = MockServer::start();
+        let resp_data = vec![0x41, 0x41, 0x41, 0x41];
+        let get_mock = server.mock(|when, then| {
+            when.method(GET).path("/image");
+            then.status(200).body(&resp_data);
+        });
+
+        let url = format!("http://{}/image", server.address());
+        let handler = UriHandlerBuilder::new(&url).set_max_size(1).build().unwrap();
         assert!(!handler.exists());
 
         let _data = handler.fetch().await.unwrap();
+        get_mock.assert();
     }
 }