From d30979d04e0af3522395ee8e6620cfba63accfdc Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Thu, 27 Feb 2025 11:29:14 +0100
Subject: [PATCH] hub script fetch retry (#5379)

* hub script fetch retry

* use backon

* oups

* Update backend/windmill-common/src/scripts.rs

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* retry whole logic

* nits

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
---
 backend/Cargo.lock                     |   1 +
 backend/windmill-common/Cargo.toml     |   1 +
 backend/windmill-common/src/scripts.rs | 104 +++++++++++++++----------
 3 files changed, 66 insertions(+), 40 deletions(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index 926930dcd5..65141b9eb1 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -13682,6 +13682,7 @@ dependencies = [
  "aws-config",
  "aws-sdk-sts",
  "axum",
+ "backon",
  "bytes",
  "chrono",
  "chrono-tz 0.10.1",
diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml
index 8f1bc299ea..9facf76ce4 100644
--- a/backend/windmill-common/Cargo.toml
+++ b/backend/windmill-common/Cargo.toml
@@ -61,6 +61,7 @@ async-stream.workspace = true
 const_format.workspace = true
 crc.workspace = true
 windmill-macros.workspace = true
+backon.workspace = true
 
 semver.workspace = true
 croner = "2.0.6"
diff --git a/backend/windmill-common/src/scripts.rs b/backend/windmill-common/src/scripts.rs
index 9d896e7f5b..5658d505b6 100644
--- a/backend/windmill-common/src/scripts.rs
+++ b/backend/windmill-common/src/scripts.rs
@@ -19,6 +19,8 @@ use crate::{
 
 use crate::worker::HUB_CACHE_DIR;
 use anyhow::Context;
+use backon::ConstantBuilder;
+use backon::{BackoffBuilder, Retryable};
 use serde::de::Error as _;
 use serde::{ser::SerializeSeq, Deserialize, Deserializer, Serialize};
 
@@ -415,6 +417,8 @@ pub async fn get_hub_script_by_path(
         Some(db),
     )
     .await?
+    .error_for_status()
+    .map_err(to_anyhow)?
     .text()
     .await
     .map_err(to_anyhow);
@@ -440,6 +444,8 @@ pub async fn get_hub_script_by_path(
                     Some(db),
                 )
                 .await?
+                .error_for_status()
+                .map_err(to_anyhow)?
                 .text()
                 .await
                 .map_err(to_anyhow)?;
@@ -494,49 +500,67 @@ async fn get_full_hub_script_by_path_inner(
 ) -> crate::error::Result<HubScript> {
     let hub_base_url = HUB_BASE_URL.read().await.clone();
 
-    let result = http_get_from_hub(
-        http_client,
-        &format!("{}/raw2/{}", hub_base_url, path),
-        true,
-        None,
-        db,
-    )
-    .await?
-    .json::<HubScript>()
-    .await
-    .context("Decoding hub response to script");
+    let response = (|| async {
+        let response = http_get_from_hub(
+            http_client,
+            &format!("{}/raw2/{}", hub_base_url, path),
+            true,
+            None,
+            db,
+        )
+        .await
+        .and_then(|r| r.error_for_status().map_err(|e| to_anyhow(e).into()));
 
-    match result {
-        Ok(result) => Ok(result),
-        Err(e) => {
-            if hub_base_url != DEFAULT_HUB_BASE_URL
-                && path
-                    .split("/")
-                    .next()
-                    .is_some_and(|x| x.parse::<i32>().is_ok_and(|x| x < 10_000_000))
-            {
-                tracing::info!(
-                    "Not found on private hub, fallback to default hub for {}",
-                    path
-                );
-                let value = http_get_from_hub(
-                    http_client,
-                    &format!("{}/raw2/{}", DEFAULT_HUB_BASE_URL, path),
-                    true,
-                    None,
-                    db,
-                )
-                .await?
-                .json::<HubScript>()
-                .await
-                .context("Decoding hub response to script")?;
-
-                Ok(value)
-            } else {
-                Err(e)?
+        match response {
+            Ok(response) => Ok(response),
+            Err(e) => {
+                if hub_base_url != DEFAULT_HUB_BASE_URL
+                    && path
+                        .split("/")
+                        .next()
+                        .is_some_and(|x| x.parse::<i32>().is_ok_and(|x| x < 10_000_000))
+                {
+                    // TODO: should only fallback to default hub if status is 404 (hub returns 500 currently)
+                    tracing::info!(
+                        "Not found on private hub, fallback to default hub for {}",
+                        path
+                    );
+                    http_get_from_hub(
+                        http_client,
+                        &format!("{}/raw2/{}", DEFAULT_HUB_BASE_URL, path),
+                        true,
+                        None,
+                        db,
+                    )
+                    .await?
+                    .error_for_status()
+                    .map_err(|e| to_anyhow(e).into())
+                } else {
+                    Err(e)
+                }
             }
         }
-    }
+    })
+    .retry(
+        ConstantBuilder::default()
+            .with_delay(std::time::Duration::from_secs(5))
+            .with_max_times(2)
+            .build(),
+    )
+    .notify(|err, dur| {
+        tracing::warn!(
+            "Could not get hub script at path {path}, retrying in {dur:#?}, err: {err:#?}"
+        );
+    })
+    .sleep(tokio::time::sleep)
+    .await?;
+
+    let script = response
+        .json::<HubScript>()
+        .await
+        .context(format!("Decoding hub response for script at path {path}"))?;
+
+    Ok(script)
 }
 
 #[derive(Deserialize, Serialize)]