fix(cubestore): Retying download from external locations (#6321)

Author: waralexrom

rust/cubestore/cubestore/src/import/mod.rs

@@ -41,6 +41,7 @@ use crate::util::decimal::Decimal;
 use crate::util::maybe_owned::MaybeOwnedStr;
 use crate::CubeError;
 use datafusion::cube_ext::ordfloat::OrdF64;
+use tokio::time::{sleep, Duration};
 
 pub mod limits;
 
@@ -465,31 +466,13 @@ impl ImportServiceImpl {
         temp_dir: &Path,
     ) -> Result<(File, Option<TempPath>), CubeError> {
         if location.starts_with("http") {
-            let (file, path) = tempfile::Builder::new()
-                .prefix(&table_id.to_string())
-                .tempfile_in(temp_dir)
-                .map_err(|e| {
-                    CubeError::internal(format!(
-                        "Open tempfile in {}: {}",
-                        temp_dir.to_str().unwrap_or("<invalid>"),
-                        e
-                    ))
-                })?
-                .into_parts();
-            let mut file = File::from_std(file);
-            let mut stream = reqwest::get(location).await?.bytes_stream();
-            let mut size = 0;
-            while let Some(bytes) = stream.next().await {
-                let bytes = bytes?;
-                let slice = bytes.as_ref();
-                size += slice.len();
-                file.write_all(slice).await?;
-            }
+            let (file, size, path) = self
+                .download_http_location(location, table_id, temp_dir)
+                .await?;
             log::info!("Import downloaded {} ({} bytes)", location, size);
             self.meta_store
                 .update_location_download_size(table_id, location.to_string(), size as u64)
                 .await?;
-            file.seek(SeekFrom::Start(0)).await?;
             Ok((file, Some(path)))
         } else if location.starts_with("temp://") {
             let temp_file = self.download_temp_file(location).await?;
@@ -509,6 +492,74 @@ impl ImportServiceImpl {
         }
     }
 
+    async fn download_http_location(
+        &self,
+        location: &str,
+        table_id: u64,
+        temp_dir: &Path,
+    ) -> Result<(File, usize, TempPath), CubeError> {
+        let max_retries: i32 = 10;
+        let mut retry_attempts = max_retries;
+        let mut retries_sleep = Duration::from_millis(100);
+        let sleep_multiplier = 2;
+        loop {
+            retry_attempts -= 1;
+            let result = self
+                .try_download_http_location(location, table_id, temp_dir)
+                .await;
+
+            if retry_attempts <= 0 {
+                return result;
+            }
+            match result {
+                Ok(size) => {
+                    return Ok(size);
+                }
+                Err(err) => {
+                    log::error!(
+                        "Import {} download error: {}. Retrying {}/{}...",
+                        location,
+                        err,
+                        retry_attempts,
+                        max_retries
+                    );
+                    sleep(retries_sleep).await;
+                    retries_sleep *= sleep_multiplier;
+                }
+            }
+        }
+    }
+
+    async fn try_download_http_location(
+        &self,
+        location: &str,
+        table_id: u64,
+        temp_dir: &Path,
+    ) -> Result<(File, usize, TempPath), CubeError> {
+        let (file, path) = tempfile::Builder::new()
+            .prefix(&table_id.to_string())
+            .tempfile_in(temp_dir)
+            .map_err(|e| {
+                CubeError::internal(format!(
+                    "Open tempfile in {}: {}",
+                    temp_dir.to_str().unwrap_or("<invalid>"),
+                    e
+                ))
+            })?
+            .into_parts();
+        let mut file = File::from_std(file);
+        let mut stream = reqwest::get(location).await?.bytes_stream();
+        let mut size = 0;
+        while let Some(bytes) = stream.next().await {
+            let bytes = bytes?;
+            let slice = bytes.as_ref();
+            size += slice.len();
+            file.write_all(slice).await?;
+        }
+        file.seek(SeekFrom::Start(0)).await?;
+        Ok((file, size, path))
+    }
+
     async fn download_temp_file(&self, location: &str) -> Result<File, CubeError> {
         let to_download = ImportServiceImpl::temp_uploads_path(location);
         // TODO check file size