Simplify logic and enrich error

[ci skip-build-wheels]

src/rust/engine/src/downloads.rs

@@ -11,7 +11,7 @@ use futures::stream::StreamExt;
 use humansize::{file_size_opts, FileSize};
 use reqwest::Error;
 use tokio_retry::strategy::{jitter, ExponentialBackoff};
-use tokio_retry::RetryIf;
+use tokio_retry::Retry;
 use url::Url;
 
 use crate::context::Core;
@@ -28,39 +28,44 @@ struct NetDownload {
 
 impl NetDownload {
   async fn start(core: &Arc<Core>, url: Url, file_name: String) -> Result<NetDownload, String> {
+    let retry_count = 4;
+
     let try_download = || async {
       core
       .http_client
       .get(url.clone())
       .send()
       .await
-      .map_err(|err| (format!("Error downloading file: {}", err), true))
+      .map_err(|err| format!("Error downloading file: {}", err))
       .and_then(|res|
         // Handle common HTTP errors.
         if res.status().is_server_error() {
-          Err((format!(
+          Err(format!(
             "Server error ({}) downloading file {} from {}",
             res.status().as_str(),
             file_name,
             url,
-          ), true))
+          ))
         } else if res.status().is_client_error() {
-          Err((format!(
+          Err(format!(
             "Client error ({}) downloading file {} from {}",
             res.status().as_str(),
             file_name,
             url,
-          ), false))
+          ))
         } else {
           Ok(res)
         })
     };
 
     // TODO: Allow the retry strategy to be configurable?
-    let retry_strategy = ExponentialBackoff::from_millis(10).map(jitter).take(3);
-    let response = RetryIf::spawn(retry_strategy, try_download, |err: &(String, bool)| err.1)
+    // For now we retry after 10ms, 100ms, 1s, and 10s.
+    let retry_strategy = ExponentialBackoff::from_millis(10)
+      .map(jitter)
+      .take(retry_count);
+    let response = Retry::spawn(retry_strategy, try_download)
       .await
-      .map_err(|(err, _)| err)?;
+      .map_err(|err| format!("After {retry_count} attempts: {err}"))?;
 
     let byte_stream = Pin::new(Box::new(response.bytes_stream()));
     Ok(NetDownload {