state: Reduce the panics in the parquet flusher, add retries to StorageProvider, send flusher failure.

arroyo-state/src/parquet.rs

@@ -1253,8 +1253,23 @@ impl ParquetFlusher {
     fn start(mut self) {
         tokio::spawn(async move {
             loop {
-                if !self.flush_iteration().await.unwrap() {
-                    return;
+                match self.flush_iteration().await {
+                    Ok(continue_flushing) => {
+                        if !continue_flushing {
+                            return;
+                        }
+                    }
+                    Err(err) => {
+                        self.control_tx
+                            .send(ControlResp::TaskFailed {
+                                operator_id: self.task_info.operator_id.clone(),
+                                task_index: self.task_info.task_index,
+                                error: err.to_string(),
+                            })
+                            .await
+                            .unwrap();
+                        return;
+                    }
                 }
             }
         });
@@ -1271,8 +1286,8 @@ impl ParquetFlusher {
         let cursor = Vec::new();
         let mut writer = ArrowWriter::try_new(cursor, record_batch.schema(), Some(props)).unwrap();
         writer.write(&record_batch)?;
-        writer.flush().unwrap();
-        let parquet_bytes = writer.into_inner().unwrap();
+        writer.flush()?;
+        let parquet_bytes = writer.into_inner()?;
         let bytes = parquet_bytes.len();
         self.storage.put(key, parquet_bytes).await?;
         Ok(bytes)
@@ -1494,10 +1509,13 @@ impl ParquetFlusher {
                 operator_id: self.task_info.operator_id.clone(),
                 subtask_metadata,
             }))
-            .await
-            .unwrap();
+            .await?;
         if cp.then_stop {
-            self.finish_tx.take().unwrap().send(()).unwrap();
+            self.finish_tx
+                .take()
+                .unwrap()
+                .send(())
+                .map_err(|_| anyhow::anyhow!("can't send finish"))?;
             return Ok(false);
         }
         Ok(true)

arroyo-storage/Cargo.toml

@@ -10,6 +10,7 @@ default = []
 [dependencies]
 arroyo-types = { path = "../arroyo-types" }
 bytes = "1.4.0"
+tracing = "0.1"
 # used only for getting local AWS credentials; can be removed once we have a
 # better way to do this
 rusoto_core = "0.48.0"

arroyo-storage/src/lib.rs

@@ -119,6 +119,31 @@ fn matchers() -> &'static HashMap<Backend, Vec<Regex>> {
     })
 }
 
+macro_rules! retry {
+    ($e:expr) => {{
+        use std::thread::sleep;
+        use std::time::Duration;
+        use tracing::error;
+        let mut retries = 0;
+        let max_retries: u32 = 10;
+        let backoff_factor: u64 = 2;
+        loop {
+            match $e {
+                Ok(value) => break Ok(value),
+                Err(e) if retries < max_retries => {
+                    retries += 1;
+                    error!("Error: {}. Retrying...", e);
+                    // exponential backoff, capped at 10 seconds.
+                    let backoff_time =
+                        Duration::from_millis(10_000.min(100 * backoff_factor.pow(retries)));
+                    sleep(backoff_time);
+                }
+                Err(e) => break Err(e),
+            }
+        }
+    }};
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct S3Config {
     endpoint: Option<String>,
@@ -502,11 +527,9 @@ impl StorageProvider {
         &self,
         path: P,
     ) -> Result<impl tokio::io::AsyncRead, StorageError> {
-        let path: String = path.into();
-        let bytes = self
-            .object_store
-            .get(&path.into())
-            .await
+        let path: Path = path.into().into();
+
+        let bytes = retry!(self.object_store.get(&path).await)
             .map_err(|e| Into::<StorageError>::into(e))?
             .into_stream();
 
@@ -519,9 +542,12 @@ impl StorageProvider {
         bytes: Vec<u8>,
     ) -> Result<String, StorageError> {
         let path = path.into().into();
-        self.object_store
-            .put(&self.qualify_path(&path), bytes.into())
-            .await?;
+        let bytes: Bytes = bytes.into();
+        retry!(
+            self.object_store
+                .put(&self.qualify_path(&path), bytes.clone())
+                .await
+        )?;
 
         Ok(format!("{}/{}", self.canonical_url, path))
     }
@@ -546,12 +572,11 @@ impl StorageProvider {
     }
 
     pub async fn start_multipart(&self, path: &Path) -> Result<MultipartId, StorageError> {
-        Ok(self
-            .object_store
-            .initiate_multipart_upload(path)
-            .await
-            .map_err(|e| Into::<StorageError>::into(e))?
-            .0)
+        Ok(
+            retry!(self.object_store.initiate_multipart_upload(path).await)
+                .map_err(|e| Into::<StorageError>::into(e))?
+                .0,
+        )
     }
 
     pub async fn add_multipart(
@@ -561,13 +586,14 @@ impl StorageProvider {
         part_number: usize,
         bytes: Bytes,
     ) -> Result<PartId, StorageError> {
-        Ok(self
-            .object_store
-            .get_put_part(path, multipart_id)
-            .await?
-            .put_part(bytes, part_number)
-            .await
-            .map_err(|e| Into::<StorageError>::into(e))?)
+        Ok(retry!(
+            self.object_store
+                .get_put_part(path, multipart_id)
+                .await?
+                .put_part(bytes.clone(), part_number)
+                .await
+        )
+        .map_err(|e| Into::<StorageError>::into(e))?)
     }
 
     pub async fn close_multipart(
@@ -576,13 +602,14 @@ impl StorageProvider {
         multipart_id: &MultipartId,
         parts: Vec<PartId>,
     ) -> Result<(), StorageError> {
-        Ok(self
-            .object_store
-            .get_put_part(path, multipart_id)
-            .await?
-            .complete(parts)
-            .await
-            .map_err(|e| Into::<StorageError>::into(e))?)
+        Ok(retry!(
+            self.object_store
+                .get_put_part(path, multipart_id)
+                .await?
+                .complete(parts.clone())
+                .await
+        )
+        .map_err(|e| Into::<StorageError>::into(e))?)
     }
 
     /// Produces a URL representation of this path that can be read by other systems,