feat(state-dumper): rewrite the state dumper logic

chain/chain/src/store/mod.rs

@@ -1020,16 +1020,35 @@ impl ChainStore {
         key
     }
 
-    /// Retrieves STATE_SYNC_DUMP for the given shard.
-    pub fn get_state_sync_dump_progress(
-        &self,
-        shard_id: ShardId,
-    ) -> Result<StateSyncDumpProgress, Error> {
-        option_to_not_found(
-            self.store
-                .get_ser(DBCol::BlockMisc, &ChainStore::state_sync_dump_progress_key(shard_id)),
-            format!("STATE_SYNC_DUMP:{}", shard_id),
-        )
+    /// For each value stored, this returs an (EpochId, bool), where the bool tells whether it's finished
+    /// because those are the only fields we really care about.
+    pub fn iter_state_sync_dump_progress<'a>(
+        &'a self,
+    ) -> impl Iterator<Item = io::Result<(ShardId, (EpochId, bool))>> + 'a {
+        self.store
+            .iter_prefix_ser::<StateSyncDumpProgress>(DBCol::BlockMisc, STATE_SYNC_DUMP_KEY)
+            .map(|item| {
+                item.and_then(|(key, progress)| {
+                    // + 1 for the ':'
+                    let prefix_len = STATE_SYNC_DUMP_KEY.len() + 1;
+                    let int_part = &key[prefix_len..];
+                    let int_part = int_part.try_into().map_err(|_| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidData,
+                            format!("Bad StateSyncDump columnn key length: {}", key.len()),
+                        )
+                    })?;
+                    let shard_id = ShardId::from_le_bytes(int_part);
+                    Ok((
+                        shard_id,
+                        match progress {
+                            StateSyncDumpProgress::AllDumped { epoch_id, .. } => (epoch_id, true),
+                            StateSyncDumpProgress::InProgress { epoch_id, .. } => (epoch_id, false),
+                            StateSyncDumpProgress::Skipped { epoch_id, .. } => (epoch_id, true),
+                        },
+                    ))
+                })
+            })
     }
 
     /// Updates STATE_SYNC_DUMP for the given shard.

chain/client/src/sync/state/shard.rs

@@ -4,7 +4,7 @@ use crate::metrics;
 use crate::sync::state::chain_requests::ChainFinalizationRequest;
 use crate::sync::state::util::query_epoch_id_and_height_for_block;
 use futures::{StreamExt, TryStreamExt};
-use near_async::futures::{FutureSpawner, FutureSpawnerExt};
+use near_async::futures::{respawn_for_parallelism, FutureSpawner};
 use near_async::messaging::AsyncSender;
 use near_chain::types::RuntimeAdapter;
 use near_chain::BlockHeader;
@@ -247,21 +247,3 @@ async fn apply_state_part(
     )?;
     Ok(())
 }
-
-/// Given a future, respawn it as an equivalent future but which does not block the
-/// driver of the future. For example, if the given future directly performs
-/// computation, normally the whoever drives the future (such as a buffered_unordered)
-/// would be blocked by the computation, thereby not allowing computation of other
-/// futures driven by the same driver to proceed. This function respawns the future
-/// onto the FutureSpawner, so the driver of the returned future would not be blocked.
-fn respawn_for_parallelism<T: Send + 'static>(
-    future_spawner: &dyn FutureSpawner,
-    name: &'static str,
-    f: impl std::future::Future<Output = T> + Send + 'static,
-) -> impl std::future::Future<Output = T> + Send + 'static {
-    let (sender, receiver) = tokio::sync::oneshot::channel();
-    future_spawner.spawn(name, async move {
-        sender.send(f.await).ok();
-    });
-    async move { receiver.await.unwrap() }
-}

core/async/src/futures.rs

@@ -42,6 +42,24 @@ impl FutureSpawnerExt for dyn FutureSpawner + '_ {
     }
 }
 
+/// Given a future, respawn it as an equivalent future but which does not block the
+/// driver of the future. For example, if the given future directly performs
+/// computation, normally the whoever drives the future (such as a buffered_unordered)
+/// would be blocked by the computation, thereby not allowing computation of other
+/// futures driven by the same driver to proceed. This function respawns the future
+/// onto the FutureSpawner, so the driver of the returned future would not be blocked.
+pub fn respawn_for_parallelism<T: Send + 'static>(
+    future_spawner: &dyn FutureSpawner,
+    name: &'static str,
+    f: impl std::future::Future<Output = T> + Send + 'static,
+) -> impl std::future::Future<Output = T> + Send + 'static {
+    let (sender, receiver) = tokio::sync::oneshot::channel();
+    future_spawner.spawn(name, async move {
+        sender.send(f.await).ok();
+    });
+    async move { receiver.await.unwrap() }
+}
+
 /// A FutureSpawner that hands over the future to Actix.
 pub struct ActixFutureSpawner;
 

integration-tests/src/test_loop/builder.rs

@@ -719,6 +719,7 @@ impl TestLoopBuilder {
                 future_spawner.spawn_boxed("state_sync_dumper", future);
                 Box::new(|| {})
             }),
+            future_spawner: Arc::new(self.test_loop.future_spawner()),
             handle: None,
         };
         let state_sync_dumper_handle = self.test_loop.data.register_data(state_sync_dumper);

integration-tests/src/tests/client/state_dump.rs

@@ -1,5 +1,6 @@
 use assert_matches::assert_matches;
 
+use near_async::futures::ActixFutureSpawner;
 use near_async::time::{Clock, Duration};
 use near_chain::near_chain_primitives::error::QueryError;
 use near_chain::{ChainGenesis, ChainStoreAccess, Provenance};
@@ -66,6 +67,7 @@ fn test_state_dump() {
         runtime,
         validator,
         dump_future_runner: StateSyncDumper::arbiter_dump_future_runner(),
+        future_spawner: Arc::new(ActixFutureSpawner),
         handle: None,
     };
     state_sync_dumper.start().unwrap();
@@ -171,6 +173,7 @@ fn run_state_sync_with_dumped_parts(
         runtime,
         validator,
         dump_future_runner: StateSyncDumper::arbiter_dump_future_runner(),
+        future_spawner: Arc::new(ActixFutureSpawner),
         handle: None,
     };
     state_sync_dumper.start().unwrap();

nearcore/Cargo.toml

@@ -44,6 +44,7 @@ strum.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
+tokio-stream.workspace = true
 tracing.workspace = true
 xz2.workspace = true
 

nearcore/src/lib.rs

@@ -388,6 +388,7 @@ pub fn start_with_config_and_synchronization(
     let state_sync_runtime =
         Arc::new(tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap());
 
+    let state_sync_spawner = Arc::new(TokioRuntimeFutureSpawner(state_sync_runtime.clone()));
     let StartClientResult { client_actor, client_arbiter_handle, resharding_handle } = start_client(
         Clock::real(),
         config.client_config.clone(),
@@ -396,7 +397,7 @@ pub fn start_with_config_and_synchronization(
         shard_tracker.clone(),
         runtime.clone(),
         node_id,
-        Arc::new(TokioRuntimeFutureSpawner(state_sync_runtime.clone())),
+        state_sync_spawner.clone(),
         network_adapter.as_multi_sender(),
         shards_manager_adapter.as_sender(),
         config.validator_signer.clone(),
@@ -433,6 +434,7 @@ pub fn start_with_config_and_synchronization(
         runtime,
         validator: config.validator_signer.clone(),
         dump_future_runner: StateSyncDumper::arbiter_dump_future_runner(),
+        future_spawner: state_sync_spawner.clone(),
         handle: None,
     };
     state_sync_dumper.start()?;