From b9b569bdb8af80dd57251e9722e20a738be4a63b Mon Sep 17 00:00:00 2001 From: Vitaly Terekhov Date: Wed, 12 Aug 2026 08:56:42 +0000 Subject: [PATCH 1/2] fix(core): do not spam with downloaded persistent state remove warnings --- core/src/block_strider/starter/cold_boot.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/core/src/block_strider/starter/cold_boot.rs b/core/src/block_strider/starter/cold_boot.rs index 35c5a2c9b..9a19ccc0c 100644 --- a/core/src/block_strider/starter/cold_boot.rs +++ b/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::>(), ) .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:?}", From 090edba81be13932695cc3891b2a7b30bd1b66f9 Mon Sep 17 00:00:00 2001 From: Vitaly Terekhov Date: Mon, 10 Aug 2026 14:37:50 +0000 Subject: [PATCH 2/2] feature(core): support persistent shard state split in s3 provider --- .../block_strider/starter/starter_client.rs | 121 ++++- core/src/blockchain_rpc/providers.rs | 93 +++- core/src/s3/client.rs | 474 +++++++++++++++++- core/src/storage/mod.rs | 4 +- .../persistent_state/shard_state/writer.rs | 48 +- core/src/storage/persistent_state/tests.rs | 20 + util/src/serde_helpers.rs | 28 ++ 7 files changed, 736 insertions(+), 52 deletions(-) diff --git a/core/src/block_strider/starter/starter_client.rs b/core/src/block_strider/starter/starter_client.rs index 5a5409ca9..0b49b22a2 100644 --- a/core/src/block_strider/starter/starter_client.rs +++ b/core/src/block_strider/starter/starter_client.rs @@ -143,7 +143,7 @@ mod s3 { block_id: &'a BlockId, kind: PersistentStateKind, ) -> Result> { - let Some(_) = self + let Some(info) = self .s3_client .get_persistent_state_info(block_id, kind) .await? @@ -151,19 +151,37 @@ mod s3 { 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::>(), + 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(()) + } + } } diff --git a/core/src/blockchain_rpc/providers.rs b/core/src/blockchain_rpc/providers.rs index fcb95b0cb..b4dd47b54 100644 --- a/core/src/blockchain_rpc/providers.rs +++ b/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, ) -> Result> { - 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 === diff --git a/core/src/s3/client.rs b/core/src/s3/client.rs index 24e8d3579..9e842d81b 100644 --- a/core/src/s3/client.rs +++ b/core/src/s3/client.rs @@ -3,15 +3,20 @@ use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize}; use std::sync::Arc; use std::time::{Duration, Instant}; +use anyhow::Context; use bytes::Bytes; use bytesize::ByteSize; +use futures_util::StreamExt; use object_store::path::Path; use object_store::{DynObjectStore, Error, ObjectStoreExt}; use serde::{Deserialize, Serialize}; use tycho_block_util::archive::ArchiveVerifier; use tycho_types::models::BlockId; -use crate::storage::PersistentStateKind; +use crate::storage::{ + PersistentStateKind, PersistentStateMeta, PersistentStatePartInfo, ShardStateWriter, + validate_persistent_state_split_metadata, +}; use crate::util::downloader::{DownloaderError, DownloaderResponseHandle, download_and_decompress}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -129,8 +134,17 @@ impl S3Client { self.inner.make_archive_key(archive_id) } - pub fn make_state_key(&self, block_id: &BlockId, kind: PersistentStateKind) -> Path { - self.inner.make_state_key(block_id, kind) + pub fn make_state_key( + &self, + block_id: &BlockId, + kind: PersistentStateKind, + part_prefix: Option, + ) -> anyhow::Result { + self.inner.make_state_key(block_id, kind, part_prefix) + } + + pub fn make_state_meta_key(&self, block_id: &BlockId) -> Path { + self.inner.make_state_meta_key(block_id) } pub async fn get_archive_info( @@ -220,30 +234,111 @@ impl S3Client { block_id: &BlockId, kind: PersistentStateKind, ) -> Result, Error> { - let path = self.inner.make_state_key(block_id, kind); - let meta = match self.inner.client.head(&path).await { - Ok(meta) if meta.size > 0 => meta, - Ok(_) | Err(object_store::Error::NotFound { .. }) => return Ok(None), + // read and validate the optional shard manifest before looking up main object + let meta_path = self.inner.make_state_meta_key(block_id); + let meta = match kind { + PersistentStateKind::Queue => None, + PersistentStateKind::Shard => match self.inner.client.get(&meta_path).await { + Ok(result) => { + let meta = PersistentStateMeta::from_bytes(&result.bytes().await?) + .map_err(persistent_state_error)? + .expect("S3 manifest bytes are always present"); + validate_persistent_state_split_metadata( + block_id.shard, + meta.split_depth, + meta.parts.iter().copied(), + ) + .map_err(persistent_state_error)?; + Some(meta) + } + Err(object_store::Error::NotFound { .. }) => None, + Err(e) => return Err(e), + }, + }; + + // read main object after manifest: is valid or not exists + let main_path = self + .inner + .make_state_key(block_id, kind, None) + .map_err(persistent_state_error)?; + let main_size = match self.inner.client.head(&main_path).await { + Ok(meta) => NonZeroU64::new(meta.size), + Err(Error::NotFound { .. }) => None, Err(e) => return Err(e), }; + let Some(main_size) = main_size else { + return Ok(None); + }; - Ok(Some(BriefPersistentStateInfo { - block_id: *block_id, - kind, - size: NonZeroU64::new(meta.size).unwrap(), - })) + // return fast if no manifest or kind is queue + let Some(meta) = meta else { + return Ok(Some(BriefPersistentStateInfo { + block_id: *block_id, + kind, + size: main_size, + split_depth: 0, + parts: Vec::new(), + })); + }; + + // if a valid manifest exists then read parts info + if let Some(parts) = self + .get_persistent_state_parts_info(block_id, kind, &meta.parts) + .await? + { + return Ok(Some(BriefPersistentStateInfo { + block_id: *block_id, + kind, + size: main_size, + split_depth: meta.split_depth, + parts, + })); + }; + + Ok(None) + } + + async fn get_persistent_state_parts_info( + &self, + block_id: &BlockId, + kind: PersistentStateKind, + prefixes: &[u64], + ) -> Result>, Error> { + let client = self.inner.client.clone(); + futures_util::stream::iter(prefixes.iter().copied()) + .map(|prefix| { + let path = self + .inner + .make_state_key(block_id, kind, Some(prefix)) + .map_err(persistent_state_error); + let client = client.clone(); + async move { + let path = path?; + match client.head(&path).await { + Ok(meta) => Ok(NonZeroU64::new(meta.size) + .map(|size| PersistentStatePartInfo { prefix, size })), + Err(Error::NotFound { .. }) => Ok(None), + Err(e) => Err(e), + } + } + }) + .buffered(PARALLEL_REQUESTS.get()) + .collect::>() + .await + .into_iter() + .collect() } #[tracing::instrument(skip_all, fields( - block_id = %block_id, - kind = ?kind, + block_id = %info.block_id, + kind = ?info.kind, ))] pub async fn download_persistent_state( &self, - block_id: &BlockId, - kind: PersistentStateKind, + info: BriefPersistentStateInfo, + part_prefix: Option, output: W, - ) -> anyhow::Result + ) -> Result where W: Write + Send + 'static, { @@ -259,11 +354,23 @@ impl S3Client { let client = &self.inner.client.clone(); - let path = self.inner.make_state_key(block_id, kind); - let meta = client.head(&path).await?; - let Some(target_size) = NonZeroU64::new(meta.size) else { - return Err(empty_file_error(path)); + let target_size = match part_prefix { + Some(prefix) => info + .parts + .iter() + .find(|part| part.prefix == prefix) + .map(|part| part.size) + .ok_or_else(|| { + persistent_state_error(anyhow::anyhow!( + "persistent state part not found: {prefix:016x}" + )) + })?, + None => info.size, }; + let path = self + .inner + .make_state_key(&info.block_id, info.kind, part_prefix) + .map_err(persistent_state_error)?; download_and_decompress( target_size, @@ -320,6 +427,9 @@ pub struct BriefPersistentStateInfo { pub block_id: BlockId, pub kind: PersistentStateKind, pub size: NonZeroU64, + /// Will be non-zero for persistent shard state with parts + pub split_depth: u8, + pub parts: Vec, } struct Inner { @@ -335,11 +445,30 @@ impl Inner { Path::from(format!("{}{archive_id}", self.archive_key_prefix)) } - fn make_state_key(&self, block_id: &BlockId, kind: PersistentStateKind) -> Path { + fn make_state_key( + &self, + block_id: &BlockId, + kind: PersistentStateKind, + part_prefix: Option, + ) -> anyhow::Result { + let file_name = match part_prefix { + Some(prefix) => kind + .make_part_file_name(block_id, prefix) + .context("persistent state parts are not supported for queue state")?, + None => kind.make_file_name(block_id), + }; + Ok(Path::from(format!( + "{}{}", + self.state_key_prefix, + file_name.display() + ))) + } + + fn make_state_meta_key(&self, block_id: &BlockId) -> Path { Path::from(format!( "{}{}", self.state_key_prefix, - kind.make_file_name(block_id).display() + ShardStateWriter::meta_file_name(block_id).display() )) } } @@ -393,6 +522,13 @@ fn empty_file_error(path: impl Into) -> Error { } } +fn persistent_state_error(error: anyhow::Error) -> Error { + Error::Generic { + store: "s3 client", + source: error.into(), + } +} + struct DownloaderHandle; impl DownloaderResponseHandle for DownloaderHandle { @@ -402,3 +538,295 @@ impl DownloaderResponseHandle for DownloaderHandle { // TODO: Move into config const PARALLEL_REQUESTS: NonZeroUsize = NonZeroUsize::new(10).unwrap(); + +#[cfg(test)] +mod tests { + use object_store::memory::InMemory; + use tycho_types::cell::HashBytes; + use tycho_types::models::ShardIdent; + + use super::*; + + impl S3Client { + pub(crate) fn new_for_tests(client: Arc) -> Self { + Self { + inner: Arc::new(Inner { + client, + archive_key_prefix: String::new(), + state_key_prefix: String::new(), + chunk_size: NonZeroU32::new(1024).unwrap(), + download_retries: 1, + }), + } + } + } + + async fn make_client() -> (S3Client, Arc) { + let store = Arc::new(InMemory::new()); + (S3Client::new_for_tests(store.clone()), store) + } + + fn block_id(shard: ShardIdent) -> BlockId { + BlockId { + shard, + seqno: 42, + root_hash: HashBytes::from([1; 32]), + file_hash: HashBytes::from([2; 32]), + } + } + + async fn put(store: &InMemory, path: Path, bytes: &[u8]) { + store + .put(&path, Bytes::copy_from_slice(bytes).into()) + .await + .unwrap(); + } + + #[tokio::test] + async fn persistent_state_discovery_handles_legacy_and_split_bundles() { + let (client, store) = make_client().await; + + // discover a legacy shard state before introducing split metadata + let legacy_id = block_id(ShardIdent::BASECHAIN); + put( + &store, + client + .make_state_key(&legacy_id, PersistentStateKind::Shard, None) + .unwrap(), + b"legacy", + ) + .await; + + let legacy = client + .get_persistent_state_info(&legacy_id, PersistentStateKind::Shard) + .await + .unwrap() + .unwrap(); + assert_eq!(legacy.size.get(), 6); + assert_eq!(legacy.split_depth, 0); + assert!(legacy.parts.is_empty()); + + // keep archive discovery independent from persistent-state representation + put(&store, client.make_archive_key(42), b"archive").await; + assert_eq!( + client + .get_archive_info(42) + .await + .unwrap() + .unwrap() + .size + .get(), + 7 + ); + + // add a split bundle whose object insertion order differs from manifest order + let split_id = BlockId { + seqno: 43, + ..legacy_id + }; + let prefixes = vec![0xa000000000000000, 0x2000000000000000]; + let meta = PersistentStateMeta::new(2, prefixes.clone()); + let meta_bytes = meta.to_bytes().unwrap(); + put(&store, client.make_state_meta_key(&split_id), &meta_bytes).await; + put( + &store, + client + .make_state_key(&split_id, PersistentStateKind::Shard, None) + .unwrap(), + b"main", + ) + .await; + put( + &store, + client + .make_state_key( + &split_id, + PersistentStateKind::Shard, + Some(0x2000000000000000), + ) + .unwrap(), + b"first", + ) + .await; + put( + &store, + client + .make_state_key( + &split_id, + PersistentStateKind::Shard, + Some(0xa000000000000000), + ) + .unwrap(), + b"second", + ) + .await; + + // discover the complete split bundle in canonical manifest order + let split = client + .get_persistent_state_info(&split_id, PersistentStateKind::Shard) + .await + .unwrap() + .unwrap(); + assert_eq!(split.size.get(), 4); + assert_eq!(split.split_depth, 2); + assert_eq!( + split + .parts + .iter() + .map(|part| (part.prefix, part.size.get())) + .collect::>(), + vec![(0x2000000000000000, 5), (0xa000000000000000, 6)] + ); + } + + #[tokio::test] + async fn persistent_state_manifest_rejects_incomplete_and_invalid_bundles() { + let (client, store) = make_client().await; + let test_block_id = block_id(ShardIdent::BASECHAIN); + + // a manifest without main and declared parts is not a usable bundle + let meta = PersistentStateMeta::new(2, vec![0x2000000000000000]); + let meta_bytes = meta.to_bytes().unwrap(); + put( + &store, + client.make_state_meta_key(&test_block_id), + &meta_bytes, + ) + .await; + assert!( + client + .get_persistent_state_info(&test_block_id, PersistentStateKind::Shard) + .await + .unwrap() + .is_none() + ); + + // adding main alone leaves the split bundle incomplete + put( + &store, + client + .make_state_key(&test_block_id, PersistentStateKind::Shard, None) + .unwrap(), + b"main", + ) + .await; + assert!( + client + .get_persistent_state_info(&test_block_id, PersistentStateKind::Shard) + .await + .unwrap() + .is_none() + ); + + // reject duplicate prefixes even when the main object is absent + let invalid_id = BlockId { + seqno: 44, + ..test_block_id + }; + put( + &store, + client.make_state_meta_key(&invalid_id), + br#"{"version":1,"split_depth":2,"parts":["2000000000000000","2000000000000000"]}"#, + ) + .await; + assert!( + client + .get_persistent_state_info(&invalid_id, PersistentStateKind::Shard) + .await + .is_err() + ); + + // reject malformed and semantically invalid manifest variants + for (seqno, bytes) in [ + (45, &b"not json"[..]), + ( + 46, + br#"{"version":2,"split_depth":2,"parts":["2000000000000000"]}"#.as_slice(), + ), + ( + 47, + br#"{"version":1,"split_depth":7,"parts":["2000000000000000"]}"#.as_slice(), + ), + ( + 48, + br#"{"version":1,"split_depth":2,"parts":["0000000000000000"]}"#.as_slice(), + ), + ] { + let invalid_id = BlockId { + seqno, + ..test_block_id + }; + put(&store, client.make_state_meta_key(&invalid_id), bytes).await; + assert!( + client + .get_persistent_state_info(&invalid_id, PersistentStateKind::Shard) + .await + .is_err() + ); + } + } + + #[tokio::test] + async fn persistent_state_download_selects_main_or_declared_part() { + let (client, store) = make_client().await; + + // store the compressed main object and its declared split part + let block_id = block_id(ShardIdent::BASECHAIN); + let main = tycho_util::compression::zstd_compress_simple(b"main"); + let part = tycho_util::compression::zstd_compress_simple(b"part"); + let info = BriefPersistentStateInfo { + block_id, + kind: PersistentStateKind::Shard, + size: NonZeroU64::new(main.len() as u64).unwrap(), + split_depth: 2, + parts: vec![PersistentStatePartInfo { + prefix: 0x2000000000000000, + size: NonZeroU64::new(part.len() as u64).unwrap(), + }], + }; + put( + &store, + client + .make_state_key(&block_id, PersistentStateKind::Shard, None) + .unwrap(), + &main, + ) + .await; + put( + &store, + client + .make_state_key( + &block_id, + PersistentStateKind::Shard, + Some(0x2000000000000000), + ) + .unwrap(), + &part, + ) + .await; + + // select the main object when no part prefix is requested + assert_eq!( + client + .download_persistent_state(info.clone(), None, Vec::new()) + .await + .unwrap(), + b"main" + ); + + // select a declared part and reject prefixes outside the manifest + assert_eq!( + client + .download_persistent_state(info.clone(), Some(0x2000000000000000), Vec::new()) + .await + .unwrap(), + b"part" + ); + assert!( + client + .download_persistent_state(info, Some(0xa000000000000000), Vec::new()) + .await + .is_err() + ); + } +} diff --git a/core/src/storage/mod.rs b/core/src/storage/mod.rs index d3de8dad2..19cb75e3f 100644 --- a/core/src/storage/mod.rs +++ b/core/src/storage/mod.rs @@ -24,8 +24,8 @@ pub use self::gc::ManualGcTrigger; pub use self::node_state::{NodeStateStorage, NodeSyncState}; pub use self::persistent_state::{ BriefBocHeader, PersistentState, PersistentStateInfo, PersistentStateKind, PersistentStateMeta, - PersistentStateStorage, QueueDiffReader, QueueStateReader, QueueStateWriter, ShardStateReader, - ShardStateWriter, validate_persistent_state_split_metadata, + PersistentStatePartInfo, PersistentStateStorage, QueueDiffReader, QueueStateReader, + QueueStateWriter, ShardStateReader, ShardStateWriter, validate_persistent_state_split_metadata, }; pub use self::shard_state::{ BlockInfoForApply, InitiatedStoreState, LoadStateHint, ShardStateStorage, diff --git a/core/src/storage/persistent_state/shard_state/writer.rs b/core/src/storage/persistent_state/shard_state/writer.rs index 41a860c34..57fbdf01a 100644 --- a/core/src/storage/persistent_state/shard_state/writer.rs +++ b/core/src/storage/persistent_state/shard_state/writer.rs @@ -592,23 +592,13 @@ impl PersistentStateMeta { } pub fn write_to_file(&self, file_path: impl AsRef) -> Result<()> { - let raw = RawPersistentStateMeta { - version: Self::VERSION, - split_depth: self.split_depth, - parts: self - .parts - .iter() - .map(|prefix| format!("{prefix:016x}")) - .collect(), - }; - let file_path = file_path.as_ref(); let temp_file_path = file_path.with_extension("temp"); scopeguard::defer! { std::fs::remove_file(&temp_file_path).ok(); } - tycho_util::serde_helpers::save_json_to_file(&raw, &temp_file_path)?; + std::fs::write(&temp_file_path, self.to_bytes()?)?; std::fs::rename(&temp_file_path, file_path)?; Ok(()) @@ -626,8 +616,31 @@ impl PersistentStateMeta { return Ok(None); } - let raw: RawPersistentStateMeta = - tycho_util::serde_helpers::load_json_from_file(file_path)?; + Self::from_bytes(&std::fs::read(file_path)?) + } + + pub fn to_bytes(&self) -> Result> { + tycho_util::serde_helpers::save_json_to_vec(self.to_raw()) + } + + pub fn from_bytes(data: &[u8]) -> Result> { + let raw: RawPersistentStateMeta = tycho_util::serde_helpers::load_json_from_slice(data)?; + Self::from_raw(raw).map(Some) + } + + fn to_raw(&self) -> RawPersistentStateMeta { + RawPersistentStateMeta { + version: Self::VERSION, + split_depth: self.split_depth, + parts: self + .parts + .iter() + .map(|prefix| format!("{prefix:016x}")) + .collect(), + } + } + + fn from_raw(raw: RawPersistentStateMeta) -> Result { anyhow::ensure!( raw.version == Self::VERSION, "unsupported persistent state meta version: {}", @@ -638,10 +651,15 @@ impl PersistentStateMeta { if prefix.len() != 16 || !prefix.chars().all(|c| c.is_ascii_hexdigit()) { anyhow::bail!("invalid persistent state part prefix: {prefix}"); } - parts.push(u64::from_str_radix(&prefix, 16)?); + let prefix = u64::from_str_radix(&prefix, 16)?; + anyhow::ensure!( + !parts.contains(&prefix), + "duplicated persistent state part prefix: {prefix:016x}" + ); + parts.push(prefix); } - Ok(Some(Self::new(raw.split_depth, parts))) + Ok(Self::new(raw.split_depth, parts)) } } diff --git a/core/src/storage/persistent_state/tests.rs b/core/src/storage/persistent_state/tests.rs index e3d3a7e19..d93f2ab09 100644 --- a/core/src/storage/persistent_state/tests.rs +++ b/core/src/storage/persistent_state/tests.rs @@ -48,6 +48,26 @@ fn persistent_state_meta_roundtrip() -> Result<()> { Ok(()) } +#[test] +fn persistent_state_meta_bytes_roundtrip() -> Result<()> { + let meta = PersistentStateMeta::new(2, vec![0xa000000000000000, 0x2000000000000000]); + assert_eq!( + PersistentStateMeta::from_bytes(&meta.to_bytes()?)?, + Some(meta) + ); + assert!( + PersistentStateMeta::from_bytes(br#"{"version":1,"split_depth":2,"parts":["invalid"]}"#) + .is_err() + ); + assert!( + PersistentStateMeta::from_bytes( + br#"{"version":1,"split_depth":2,"parts":["2000000000000000","2000000000000000"]}"# + ) + .is_err() + ); + Ok(()) +} + #[test] fn shard_state_part_file_name_parser_recognizes_only_parts() -> Result<()> { let block_id = BlockId { diff --git a/util/src/serde_helpers.rs b/util/src/serde_helpers.rs index 5a30bd643..7f3aa3ee4 100644 --- a/util/src/serde_helpers.rs +++ b/util/src/serde_helpers.rs @@ -29,6 +29,21 @@ where Ok(()) } +pub fn load_json_from_slice(data: &[u8]) -> Result +where + for<'de> T: Deserialize<'de>, +{ + let de = &mut serde_json::Deserializer::from_slice(data); + serde_path_to_error::deserialize(de).map_err(Into::into) +} + +pub fn save_json_to_vec(value: T) -> Result> +where + T: Serialize, +{ + Ok(serde_json::to_vec_pretty(&value)?) +} + pub mod socket_addr { use std::net::SocketAddr; @@ -551,4 +566,17 @@ mod tests { assert_eq!(value, parsed); } } + + #[test] + fn json_slice_roundtrip() { + #[derive(Debug, Eq, PartialEq, Serialize, Deserialize)] + struct Test { + value: u64, + } + + let data = save_json_to_vec(&Test { value: 123 }).unwrap(); + assert_eq!(load_json_from_slice::(&data).unwrap(), Test { + value: 123 + }); + } }