Support persistent split with s3 provider

core/src/block_strider/starter/cold_boot.rs

@@ -1159,7 +1159,9 @@ impl StarterInner {
 
         // NOTE: Intentionally dont spawn yet
         let remove_state_file = async move {
-            if let Err(e) = tokio::fs::remove_file(&state_file_path).await {
+            if let Err(e) = tokio::fs::remove_file(&state_file_path).await
+                && e.kind() != std::io::ErrorKind::NotFound
+            {
                 tracing::warn!(
                     path = %state_file_path.display(),
                     "failed to remove downloaded queue state: {e:?}",
@@ -1302,7 +1304,9 @@ impl StarterInner {
                 }
             }
 
-            if local_meta.as_ref() != Some(&remote_meta) {
+            // remove previously downloaded persistent shard state files
+            // if local meta doe not match the remote one
+            if kind == PersistentStateKind::Shard && local_meta.as_ref() != Some(&remote_meta) {
                 let old_prefixes = local_meta
                     .as_ref()
                     .into_iter()
@@ -1314,9 +1318,7 @@ impl StarterInner {
                     &old_prefixes.chain(new_prefixes).collect::<Vec<_>>(),
                 )
                 .await;
-                if kind == PersistentStateKind::Shard {
-                    remote_meta.write_to_file(meta_file.path())?;
-                }
+                remote_meta.write_to_file(meta_file.path())?;
             }
 
             if is_downloaded_persistent_state_ready(state_file, &remote_meta) {
@@ -1414,7 +1416,9 @@ async fn remove_downloaded_persistent_state_files(
             return;
         }
 
-        if let Err(e) = tokio::fs::remove_file(file_path).await {
+        if let Err(e) = tokio::fs::remove_file(file_path).await
+            && e.kind() != std::io::ErrorKind::NotFound
+        {
             tracing::warn!(
                 file_path = %file_path.display(),
                 "failed to remove downloaded shard state: {e:?}",

core/src/block_strider/starter/starter_client.rs

@@ -143,27 +143,45 @@ mod s3 {
             block_id: &'a BlockId,
             kind: PersistentStateKind,
         ) -> Result<FoundState<'a>> {
-            let Some(_) = self
+            let Some(info) = self
                 .s3_client
                 .get_persistent_state_info(block_id, kind)
                 .await?
             else {
                 anyhow::bail!("not found");
             };
 
+            let info_for_part = info.clone();
+
             Ok(FoundState {
-                split_depth: 0,
-                parts: Vec::new(),
+                split_depth: info.split_depth,
+                parts: info
+                    .parts
+                    .iter()
+                    .map(|part| FoundStatePart {
+                        prefix: part.prefix,
+                    })
+                    .collect(),
                 download: Box::new(move |output| {
+                    let info = info.clone();
                     Box::pin(async move {
                         let output = self
                             .s3_client
-                            .download_persistent_state(block_id, kind, output)
+                            .download_persistent_state(info, None, output)
                             .await?;
                         Ok(output)
                     })
                 }),
-                download_part: None,
+                download_part: Some(Box::new(move |part, output| {
+                    let info = info_for_part.clone();
+                    Box::pin(async move {
+                        let output = self
+                            .s3_client
+                            .download_persistent_state(info, Some(part.prefix), output)
+                            .await?;
+                        Ok(output)
+                    })
+                })),
             })
         }
 
@@ -394,4 +412,97 @@ mod s3 {
             }
         }
     }
+
+    #[cfg(test)]
+    mod tests {
+        use bytes::Bytes;
+        use object_store::ObjectStoreExt;
+        use object_store::memory::InMemory;
+        use tycho_storage::StorageContext;
+        use tycho_types::cell::HashBytes;
+        use tycho_types::models::ShardIdent;
+        use tycho_util::fs::MappedFile;
+
+        use super::*;
+        use crate::storage::{CoreStorageConfig, PersistentStateMeta};
+
+        #[tokio::test]
+        async fn s3_starter_client_returns_split_found_state_and_downloads() -> Result<()> {
+            let store = Arc::new(InMemory::new());
+            let client = S3Client::new_for_tests(store.clone());
+            let (ctx, _tmp_dir) = StorageContext::new_temp().await?;
+            let storage = CoreStorage::open(ctx, CoreStorageConfig::new_potato()).await?;
+            let block_id = BlockId {
+                shard: ShardIdent::BASECHAIN,
+                seqno: 42,
+                root_hash: HashBytes::from([1; 32]),
+                file_hash: HashBytes::from([2; 32]),
+            };
+
+            // publish the split fixture
+            let prefixes = vec![0x2000000000000000, 0xa000000000000000];
+            let main = b"split main";
+            let parts = [b"first part".as_slice(), b"second part".as_slice()];
+            let meta = PersistentStateMeta::new(2, prefixes.clone());
+
+            store
+                .put(
+                    &client.make_state_meta_key(&block_id),
+                    Bytes::from(meta.to_bytes()?).into(),
+                )
+                .await?;
+            store
+                .put(
+                    &client.make_state_key(&block_id, PersistentStateKind::Shard, None)?,
+                    tycho_util::compression::zstd_compress_simple(main).into(),
+                )
+                .await?;
+            for (prefix, part) in prefixes.iter().zip(parts) {
+                store
+                    .put(
+                        &client.make_state_key(
+                            &block_id,
+                            PersistentStateKind::Shard,
+                            Some(*prefix),
+                        )?,
+                        tycho_util::compression::zstd_compress_simple(part).into(),
+                    )
+                    .await?;
+            }
+
+            let starter_client = S3StarterClient::new(client, storage.clone());
+
+            // discover the split state
+            let mut found = starter_client
+                .find_persistent_state(&block_id, PersistentStateKind::Shard)
+                .await?;
+            assert_eq!(found.split_depth, 2);
+            assert_eq!(
+                found
+                    .parts
+                    .iter()
+                    .map(|part| part.prefix)
+                    .collect::<Vec<_>>(),
+                prefixes,
+            );
+
+            // download the main state file
+            let main_file =
+                (found.download)(storage.context().temp_files().unnamed_file().open()?).await?;
+            let main_file = MappedFile::from_existing_file(main_file)?;
+            assert_eq!(main_file.as_slice(), main);
+
+            // download one declared part
+            let download_part = found.download_part.take().expect("split part downloader");
+            let part_file = download_part(
+                found.parts[0].clone(),
+                storage.context().temp_files().unnamed_file().open()?,
+            )
+            .await?;
+            let part_file = MappedFile::from_existing_file(part_file)?;
+            assert_eq!(part_file.as_slice(), parts[0]);
+
+            Ok(())
+        }
+    }
 }

core/src/blockchain_rpc/providers.rs

@@ -229,8 +229,8 @@ mod s3_impl {
                 .map(|info| PersistentStateInfo {
                     size: info.size,
                     chunk_size: self.chunk_size,
-                    split_depth: 0,
-                    parts: Vec::new(),
+                    split_depth: info.split_depth,
+                    parts: info.parts,
                 }))
         }
 
@@ -241,10 +241,6 @@ mod s3_impl {
             kind: PersistentStateKind,
             part_shard_prefix: Option<u64>,
         ) -> Result<Option<Bytes>> {
-            if part_shard_prefix.is_some() {
-                return Ok(None);
-            }
-
             self.check_rate_limit()?;
             self.check_bandwidth_limit()?;
 
@@ -254,7 +250,13 @@ mod s3_impl {
                 return Ok(None);
             }
 
-            let path = self.client.make_state_key(block_id, kind);
+            let path = match self
+                .client
+                .make_state_key(block_id, kind, part_shard_prefix)
+            {
+                Ok(path) => path,
+                Err(_) => return Ok(None),
+            };
             let client = self.client.client();
 
             let range = std::ops::Range {
@@ -287,6 +289,83 @@ mod s3_impl {
             Ok(())
         }
     }
+
+    #[cfg(test)]
+    mod tests {
+        use object_store::memory::InMemory;
+        use tycho_storage::StorageContext;
+        use tycho_types::cell::HashBytes;
+        use tycho_types::models::ShardIdent;
+
+        use super::*;
+        use crate::storage::{CoreStorageConfig, PersistentStateMeta};
+
+        #[tokio::test]
+        async fn s3_rpc_provider_advertises_split_info_and_reads_declared_part() -> Result<()> {
+            let store = Arc::new(InMemory::new());
+            let client = S3Client::new_for_tests(store.clone());
+            let (ctx, _tmp_dir) = StorageContext::new_temp().await?;
+            let storage = CoreStorage::open(ctx, CoreStorageConfig::new_potato()).await?;
+            let provider =
+                S3RpcDataProvider::new(client.clone(), storage, &S3ProxyConfig::default());
+            let block_id = BlockId {
+                shard: ShardIdent::BASECHAIN,
+                seqno: 42,
+                root_hash: HashBytes::from([1; 32]),
+                file_hash: HashBytes::from([2; 32]),
+            };
+
+            // publish the split fixture
+            let prefix = 0x2000000000000000;
+            let main = vec![1; client.chunk_size().get() as usize];
+            let part = vec![2; client.chunk_size().get() as usize];
+            let meta = PersistentStateMeta::new(2, vec![prefix]);
+
+            store
+                .put(
+                    &client.make_state_meta_key(&block_id),
+                    Bytes::from(meta.to_bytes()?).into(),
+                )
+                .await?;
+            store
+                .put(
+                    &client.make_state_key(&block_id, PersistentStateKind::Shard, None)?,
+                    Bytes::from(main.clone()).into(),
+                )
+                .await?;
+            store
+                .put(
+                    &client.make_state_key(&block_id, PersistentStateKind::Shard, Some(prefix))?,
+                    Bytes::from(part.clone()).into(),
+                )
+                .await?;
+
+            // request split state info
+            let info = provider
+                .get_persistent_state_info(&block_id, PersistentStateKind::Shard)
+                .await?
+                .expect("split persistent state must be available");
+            assert_eq!(info.split_depth, 2);
+            assert_eq!(info.size.get(), main.len() as u64);
+            assert_eq!(info.parts.len(), 1);
+            assert_eq!(info.parts[0].prefix, prefix);
+            assert_eq!(info.parts[0].size.get(), part.len() as u64);
+
+            // read the aligned declared part chunk
+            let part_chunk = provider
+                .get_persistent_state_chunk(&block_id, 0, PersistentStateKind::Shard, Some(prefix))
+                .await?;
+            assert_eq!(part_chunk, Some(Bytes::from(part)));
+
+            // request an unaligned chunk
+            let unaligned_chunk = provider
+                .get_persistent_state_chunk(&block_id, 1, PersistentStateKind::Shard, Some(prefix))
+                .await?;
+            assert_eq!(unaligned_chunk, None);
+
+            Ok(())
+        }
+    }
 }
 
 // === Hybrid Implementation ===