fix(block-reverter): Fix reverting snapshot files

core/bin/block_reverter/src/main.rs

@@ -69,6 +69,9 @@ enum Command {
         /// Flag that specifies if RocksDB with state keeper cache should be rolled back.
         #[arg(long)]
         rollback_sk_cache: bool,
+        /// Flag that specifies if snapshot files in GCS should be rolled back.
+        #[arg(long, requires = "rollback_postgres")]
+        rollback_snapshots: bool,
         /// Flag that allows to roll back already executed blocks. It's ultra dangerous and required only for fixing external nodes.
         #[arg(long)]
         allow_executed_block_reversion: bool,
@@ -187,6 +190,7 @@ async fn main() -> anyhow::Result<()> {
             rollback_postgres,
             rollback_tree,
             rollback_sk_cache,
+            rollback_snapshots,
             allow_executed_block_reversion,
         } => {
             if !rollback_tree && rollback_postgres {
@@ -219,13 +223,15 @@ async fn main() -> anyhow::Result<()> {
 
             if rollback_postgres {
                 block_reverter.enable_rolling_back_postgres();
-                let object_store_config = SnapshotsObjectStoreConfig::from_env()
-                    .context("SnapshotsObjectStoreConfig::from_env()")?;
-                block_reverter.enable_rolling_back_snapshot_objects(
-                    ObjectStoreFactory::new(object_store_config.0)
-                        .create_store()
-                        .await,
-                );
+                if rollback_snapshots {
+                    let object_store_config = SnapshotsObjectStoreConfig::from_env()
+                        .context("SnapshotsObjectStoreConfig::from_env()")?;
+                    block_reverter.enable_rolling_back_snapshot_objects(
+                        ObjectStoreFactory::new(object_store_config.0)
+                            .create_store()
+                            .await,
+                    );
+                }
             }
             if rollback_tree {
                 block_reverter.enable_rolling_back_merkle_tree(db_config.merkle_tree.path);

core/lib/object_store/Cargo.toml

@@ -22,9 +22,11 @@ google-cloud-auth.workspace = true
 http.workspace = true
 serde_json.workspace = true
 flate2.workspace = true
+rand.workspace = true
 tokio = { workspace = true, features = ["full"] }
 tracing.workspace = true
 prost.workspace = true
 
 [dev-dependencies]
+assert_matches.workspace = true
 tempfile.workspace = true

core/lib/object_store/src/file.rs

@@ -9,7 +9,10 @@ impl From<io::Error> for ObjectStoreError {
     fn from(err: io::Error) -> Self {
         match err.kind() {
             io::ErrorKind::NotFound => ObjectStoreError::KeyNotFound(err.into()),
-            _ => ObjectStoreError::Other(err.into()),
+            kind => ObjectStoreError::Other {
+                is_transient: matches!(kind, io::ErrorKind::Interrupted | io::ErrorKind::TimedOut),
+                source: err.into(),
+            },
         }
     }
 }
@@ -20,7 +23,7 @@ pub(crate) struct FileBackedObjectStore {
 }
 
 impl FileBackedObjectStore {
-    pub async fn new(base_dir: String) -> Self {
+    pub async fn new(base_dir: String) -> Result<Self, ObjectStoreError> {
         for bucket in &[
             Bucket::ProverJobs,
             Bucket::WitnessInput,
@@ -36,13 +39,9 @@ impl FileBackedObjectStore {
             Bucket::TeeVerifierInput,
         ] {
             let bucket_path = format!("{base_dir}/{bucket}");
-            fs::create_dir_all(&bucket_path)
-                .await
-                .unwrap_or_else(|err| {
-                    panic!("failed creating bucket `{bucket_path}`: {err}");
-                });
+            fs::create_dir_all(&bucket_path).await?;
         }
-        FileBackedObjectStore { base_dir }
+        Ok(FileBackedObjectStore { base_dir })
     }
 
     fn filename(&self, bucket: Bucket, key: &str) -> String {
@@ -87,12 +86,12 @@ mod test {
     async fn test_get() {
         let dir = TempDir::new().unwrap();
         let path = dir.into_path().into_os_string().into_string().unwrap();
-        let object_store = FileBackedObjectStore::new(path).await;
+        let object_store = FileBackedObjectStore::new(path).await.unwrap();
         let expected = vec![9, 0, 8, 9, 0, 7];
-        let result = object_store
+        object_store
             .put_raw(Bucket::ProverJobs, "test-key.bin", expected.clone())
-            .await;
-        assert!(result.is_ok(), "result must be OK");
+            .await
+            .unwrap();
         let bytes = object_store
             .get_raw(Bucket::ProverJobs, "test-key.bin")
             .await
@@ -104,26 +103,26 @@ mod test {
     async fn test_put() {
         let dir = TempDir::new().unwrap();
         let path = dir.into_path().into_os_string().into_string().unwrap();
-        let object_store = FileBackedObjectStore::new(path).await;
+        let object_store = FileBackedObjectStore::new(path).await.unwrap();
         let bytes = vec![9, 0, 8, 9, 0, 7];
-        let result = object_store
+        object_store
             .put_raw(Bucket::ProverJobs, "test-key.bin", bytes)
-            .await;
-        assert!(result.is_ok(), "result must be OK");
+            .await
+            .unwrap();
     }
 
     #[tokio::test]
     async fn test_remove() {
         let dir = TempDir::new().unwrap();
         let path = dir.into_path().into_os_string().into_string().unwrap();
-        let object_store = FileBackedObjectStore::new(path).await;
-        let result = object_store
+        let object_store = FileBackedObjectStore::new(path).await.unwrap();
+        object_store
             .put_raw(Bucket::ProverJobs, "test-key.bin", vec![0, 1])
-            .await;
-        assert!(result.is_ok(), "result must be OK");
-        let result = object_store
+            .await
+            .unwrap();
+        object_store
             .remove_raw(Bucket::ProverJobs, "test-key.bin")
-            .await;
-        assert!(result.is_ok(), "result must be OK");
+            .await
+            .unwrap();
     }
 }