refactor: make the P2P/statesync impl generic over the StateSyncMessage type

Author: rumenov

rs/interfaces/src/p2p/state_sync.rs

@@ -1,10 +1,11 @@
 use ic_types::{
     artifact::StateSyncArtifactId,
     chunkable::{Chunk, ChunkId, Chunkable},
-    state_sync::StateSyncMessage,
 };
 
 pub trait StateSyncClient: Send + Sync {
+    type Message;
+
     /// Returns a list of all states available.
     fn available_states(&self) -> Vec<StateSyncArtifactId>;
     /// Initiates new state sync for the specified Id. Returns None if the state should not be synced.
@@ -15,13 +16,13 @@ pub trait StateSyncClient: Send + Sync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>;
+    ) -> Option<Box<dyn Chunkable<Self::Message> + Send>>;
     /// Returns true if a state sync with the specified Id can be cancelled because a newer state is available.
     /// The result of this function is only meaningful the Id refers to a active state sync started with `start_state_sync`.
     /// TODO: (NET-1469) In the future this API should be made safer by only allowing the id of a previously initiated state sync.
     fn should_cancel(&self, id: &StateSyncArtifactId) -> bool;
     /// Get a specific chunk from the specified state.
     fn chunk(&self, id: &StateSyncArtifactId, chunk_id: ChunkId) -> Option<Chunk>;
     /// Finish a state sync by delivering the `StateSyncMessage` returned in `Chunkable::add_chunks`.
-    fn deliver_state_sync(&self, msg: StateSyncMessage);
+    fn deliver_state_sync(&self, msg: Self::Message);
 }

rs/p2p/state_sync_manager/src/lib.rs

@@ -51,8 +51,8 @@ const ADVERT_BROADCAST_INTERVAL: Duration = Duration::from_secs(5);
 const ADVERT_BROADCAST_TIMEOUT: Duration =
     ADVERT_BROADCAST_INTERVAL.saturating_sub(Duration::from_secs(2));
 
-pub fn build_axum_router(
-    state_sync: Arc<dyn StateSyncClient>,
+pub fn build_axum_router<T: 'static>(
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
     log: ReplicaLogger,
     metrics_registry: &MetricsRegistry,
 ) -> (
@@ -81,12 +81,12 @@ pub fn build_axum_router(
     (app, rx)
 }
 
-pub fn start_state_sync_manager(
+pub fn start_state_sync_manager<T: Send + 'static>(
     log: ReplicaLogger,
     metrics: &MetricsRegistry,
     rt: &Handle,
     transport: Arc<dyn Transport>,
-    state_sync: Arc<dyn StateSyncClient>,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
     advert_receiver: tokio::sync::mpsc::Receiver<(StateSyncArtifactId, NodeId)>,
 ) -> JoinHandle<()> {
     let state_sync_manager_metrics = StateSyncManagerMetrics::new(metrics);
@@ -102,17 +102,17 @@ pub fn start_state_sync_manager(
     rt.spawn(manager.run())
 }
 
-struct StateSyncManager {
+struct StateSyncManager<T> {
     log: ReplicaLogger,
     rt: Handle,
     metrics: StateSyncManagerMetrics,
     transport: Arc<dyn Transport>,
-    state_sync: Arc<dyn StateSyncClient>,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
     advert_receiver: tokio::sync::mpsc::Receiver<(StateSyncArtifactId, NodeId)>,
     ongoing_state_sync: Option<OngoingStateSyncHandle>,
 }
 
-impl StateSyncManager {
+impl<T: 'static + Send> StateSyncManager<T> {
     async fn run(mut self) {
         let mut interval = tokio::time::interval(ADVERT_BROADCAST_INTERVAL);
         let mut advertise_task = JoinSet::new();
@@ -186,7 +186,7 @@ impl StateSyncManager {
 
     async fn send_state_adverts(
         rt: Handle,
-        state_sync: Arc<dyn StateSyncClient>,
+        state_sync: Arc<dyn StateSyncClient<Message = T>>,
         transport: Arc<dyn Transport>,
         metrics: StateSyncManagerMetrics,
     ) {

rs/p2p/state_sync_manager/src/metrics.rs

@@ -1,7 +1,6 @@
 use ic_metrics::{
     buckets::decimal_buckets, tokio_metrics_collector::TokioTaskMetricsCollector, MetricsRegistry,
 };
-use ic_types::state_sync::StateSyncMessage;
 use prometheus::{Histogram, IntCounter, IntCounterVec, IntGauge};
 use tokio_metrics::TaskMonitor;
 
@@ -119,10 +118,7 @@ impl OngoingStateSyncMetrics {
     }
 
     /// Utility to record metrics for download result.
-    pub fn record_chunk_download_result(
-        &self,
-        res: &Result<Option<StateSyncMessage>, DownloadChunkError>,
-    ) {
+    pub fn record_chunk_download_result<T>(&self, res: &Result<Option<T>, DownloadChunkError>) {
         match res {
             // Received chunk
             Ok(Some(_)) => {

rs/p2p/state_sync_manager/src/ongoing.rs

@@ -27,7 +27,6 @@ use ic_types::{
     artifact::StateSyncArtifactId,
     chunkable::ChunkId,
     chunkable::{ArtifactErrorCode, Chunkable},
-    state_sync::StateSyncMessage,
     NodeId,
 };
 use rand::{
@@ -53,7 +52,7 @@ const CHUNK_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(10);
 /// to last for 1000 blocks (two checkpoint intervals) -> 1000b/0.1b/s = 10000s
 const STATE_SYNC_TIMEOUT: Duration = Duration::from_secs(10000);
 
-struct OngoingStateSync {
+struct OngoingStateSync<T: Send> {
     log: ReplicaLogger,
     rt: Handle,
     artifact_id: StateSyncArtifactId,
@@ -67,10 +66,10 @@ struct OngoingStateSync {
     allowed_downloads: usize,
     chunks_to_download: Box<dyn Iterator<Item = ChunkId> + Send>,
     // Event tasks
-    downloading_chunks: JoinMap<ChunkId, DownloadResult>,
+    downloading_chunks: JoinMap<ChunkId, DownloadResult<T>>,
     // State sync
-    state_sync: Arc<dyn StateSyncClient>,
-    tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
+    tracker: Arc<Mutex<Box<dyn Chunkable<T> + Send>>>,
     state_sync_finished: bool,
 }
 
@@ -79,18 +78,18 @@ pub(crate) struct OngoingStateSyncHandle {
     pub jh: JoinHandle<()>,
 }
 
-pub(crate) struct DownloadResult {
+pub(crate) struct DownloadResult<T> {
     peer_id: NodeId,
-    result: Result<Option<StateSyncMessage>, DownloadChunkError>,
+    result: Result<Option<T>, DownloadChunkError>,
 }
 
-pub(crate) fn start_ongoing_state_sync(
+pub(crate) fn start_ongoing_state_sync<T: Send + 'static>(
     log: ReplicaLogger,
     rt: &Handle,
     metrics: OngoingStateSyncMetrics,
-    tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
+    tracker: Arc<Mutex<Box<dyn Chunkable<T> + Send>>>,
     artifact_id: StateSyncArtifactId,
-    state_sync: Arc<dyn StateSyncClient>,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
     transport: Arc<dyn Transport>,
 ) -> OngoingStateSyncHandle {
     let (new_peers_tx, new_peers_rx) = tokio::sync::mpsc::channel(ONGOING_STATE_SYNC_CHANNEL_SIZE);
@@ -117,7 +116,7 @@ pub(crate) fn start_ongoing_state_sync(
     }
 }
 
-impl OngoingStateSync {
+impl<T: 'static + Send> OngoingStateSync<T> {
     pub async fn run(mut self) {
         let state_sync_timeout = tokio::time::sleep(STATE_SYNC_TIMEOUT);
         tokio::pin!(state_sync_timeout);
@@ -193,7 +192,7 @@ impl OngoingStateSync {
 
     async fn handle_downloaded_chunk_result(
         &mut self,
-        DownloadResult { peer_id, result }: DownloadResult,
+        DownloadResult { peer_id, result }: DownloadResult<T>,
     ) {
         self.metrics.record_chunk_download_result(&result);
         match result {
@@ -296,11 +295,11 @@ impl OngoingStateSync {
     async fn download_chunk_task(
         peer_id: NodeId,
         client: Arc<dyn Transport>,
-        tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
+        tracker: Arc<Mutex<Box<dyn Chunkable<T> + Send>>>,
         artifact_id: StateSyncArtifactId,
         chunk_id: ChunkId,
         metrics: OngoingStateSyncMetrics,
-    ) -> DownloadResult {
+    ) -> DownloadResult<T> {
         let _timer = metrics.chunk_download_duration.start_timer();
 
         let response_result = tokio::time::timeout(
@@ -386,6 +385,9 @@ mod tests {
 
     use super::*;
 
+    #[derive(Clone)]
+    struct TestMessage;
+
     fn compress_empty_bytes() -> Bytes {
         let mut raw = BytesMut::new();
         Bytes::new()
@@ -409,7 +411,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::<StateSyncMessage>::default();
+            let mut c = MockChunkable::<TestMessage>::default();
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));
 
@@ -448,7 +450,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::<StateSyncMessage>::default();
+            let mut c = MockChunkable::<TestMessage>::default();
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));
             c.expect_add_chunk()
@@ -494,7 +496,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::<StateSyncMessage>::default();
+            let mut c = MockChunkable::<TestMessage>::default();
             // Endless iterator
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));

rs/p2p/state_sync_manager/src/routes/chunk.rs

@@ -22,16 +22,16 @@ pub const STATE_SYNC_CHUNK_PATH: &str = "/state-sync/chunk";
 /// State sync uses 1Mb chunks. To be safe we use 8Mib here same as transport.
 const MAX_CHUNK_SIZE: usize = 8 * 1024 * 1024;
 
-pub(crate) struct StateSyncChunkHandler {
+pub(crate) struct StateSyncChunkHandler<T> {
     _log: ReplicaLogger,
-    state_sync: Arc<dyn StateSyncClient>,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
     metrics: StateSyncManagerHandlerMetrics,
 }
 
-impl StateSyncChunkHandler {
+impl<T> StateSyncChunkHandler<T> {
     pub fn new(
         log: ReplicaLogger,
-        state_sync: Arc<dyn StateSyncClient>,
+        state_sync: Arc<dyn StateSyncClient<Message = T>>,
         metrics: StateSyncManagerHandlerMetrics,
     ) -> Self {
         Self {
@@ -42,8 +42,8 @@ impl StateSyncChunkHandler {
     }
 }
 
-pub(crate) async fn state_sync_chunk_handler(
-    State(state): State<Arc<StateSyncChunkHandler>>,
+pub(crate) async fn state_sync_chunk_handler<T: 'static>(
+    State(state): State<Arc<StateSyncChunkHandler<T>>>,
     payload: Bytes,
 ) -> Result<Bytes, StatusCode> {
     // Parse payload

rs/p2p/state_sync_manager/tests/common.rs

@@ -183,6 +183,8 @@ impl FakeStateSync {
 }
 
 impl StateSyncClient for FakeStateSync {
+    type Message = StateSyncMessage;
+
     fn available_states(&self) -> Vec<StateSyncArtifactId> {
         if self.disconnected.load(Ordering::SeqCst) {
             return vec![];
@@ -197,7 +199,7 @@ impl StateSyncClient for FakeStateSync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>> {
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send>> {
         if !self.uses_global() && id.height > self.current_height() && !self.disconnected() {
             return Some(Box::new(FakeChunkable::new(
                 self.local_state.clone(),
@@ -352,7 +354,7 @@ impl Chunkable<StateSyncMessage> for SharableMockChunkable {
 
 #[derive(Clone, Default)]
 pub struct SharableMockStateSync {
-    mock: Arc<Mutex<MockStateSync>>,
+    mock: Arc<Mutex<MockStateSync<StateSyncMessage>>>,
     available_states_calls: Arc<AtomicUsize>,
     start_state_sync_calls: Arc<AtomicUsize>,
     should_cancel_calls: Arc<AtomicUsize>,
@@ -367,7 +369,7 @@ impl SharableMockStateSync {
             ..Default::default()
         }
     }
-    pub fn get_mut(&self) -> MutexGuard<'_, MockStateSync> {
+    pub fn get_mut(&self) -> MutexGuard<'_, MockStateSync<StateSyncMessage>> {
         self.mock.lock().unwrap()
     }
     pub fn start_state_sync_calls(&self) -> usize {
@@ -383,14 +385,16 @@ impl SharableMockStateSync {
 }
 
 impl StateSyncClient for SharableMockStateSync {
+    type Message = StateSyncMessage;
+
     fn available_states(&self) -> Vec<StateSyncArtifactId> {
         self.available_states_calls.fetch_add(1, Ordering::SeqCst);
         self.mock.lock().unwrap().available_states()
     }
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>> {
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send>> {
         self.start_state_sync_calls.fetch_add(1, Ordering::SeqCst);
         self.mock.lock().unwrap().start_state_sync(id)
     }

rs/p2p/test_utils/src/mocks.rs

@@ -12,27 +12,28 @@ use ic_types::{
     artifact::StateSyncArtifactId,
     chunkable::{ArtifactErrorCode, Chunkable},
     chunkable::{Chunk, ChunkId},
-    state_sync::StateSyncMessage,
     NodeId,
 };
 use mockall::mock;
 
 mock! {
-    pub StateSync {}
+    pub StateSync<T: Send> {}
+
+    impl<T: Send + Sync> StateSyncClient for StateSync<T> {
+        type Message = T;
 
-    impl StateSyncClient for StateSync {
         fn available_states(&self) -> Vec<StateSyncArtifactId>;
 
         fn start_state_sync(
             &self,
             id: &StateSyncArtifactId,
-        ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>;
+        ) -> Option<Box<dyn Chunkable<T> + Send>>;
 
         fn should_cancel(&self, id: &StateSyncArtifactId) -> bool;
 
         fn chunk(&self, id: &StateSyncArtifactId, chunk_id: ChunkId) -> Option<Chunk>;
 
-        fn deliver_state_sync(&self, msg: StateSyncMessage);
+        fn deliver_state_sync(&self, msg: T);
     }
 }
 

rs/p2p/test_utils/src/turmoil.rs

@@ -27,6 +27,7 @@ use ic_logger::ReplicaLogger;
 use ic_metrics::MetricsRegistry;
 use ic_peer_manager::SubnetTopology;
 use ic_quic_transport::{QuicTransport, Transport};
+use ic_types::state_sync::StateSyncMessage;
 use ic_types::{artifact::UnvalidatedArtifactMutation, NodeId, RegistryVersion};
 use ic_types_test_utils::ids::SUBNET_1;
 use quinn::{
@@ -271,7 +272,7 @@ pub fn add_transport_to_sim<F>(
     conn_checker: Option<Router>,
     crypto: Option<Arc<dyn TlsConfig + Send + Sync>>,
     sev: Option<Arc<dyn ValidateAttestedStream<Box<dyn TlsStream>> + Send + Sync>>,
-    state_sync_client: Option<Arc<dyn StateSyncClient>>,
+    state_sync_client: Option<Arc<dyn StateSyncClient<Message = StateSyncMessage>>>,
     consensus_manager: Option<TestConsensus<U64Artifact>>,
     post_setup_future: F,
 ) where

rs/replica/setup_ic_network/src/lib.rs

@@ -65,6 +65,7 @@ use ic_types::{
     malicious_flags::MaliciousFlags,
     p2p::GossipAdvert,
     replica_config::ReplicaConfig,
+    state_sync::StateSyncMessage,
     NodeId, SubnetId,
 };
 use std::{
@@ -147,7 +148,7 @@ pub fn setup_consensus_and_p2p(
     state_reader: Arc<dyn StateReader<State = ReplicatedState>>,
     consensus_pool: Arc<RwLock<ConsensusPoolImpl>>,
     catch_up_package: CatchUpPackage,
-    state_sync_client: Arc<dyn StateSyncClient>,
+    state_sync_client: Arc<dyn StateSyncClient<Message = StateSyncMessage>>,
     xnet_payload_builder: Arc<dyn XNetPayloadBuilder>,
     self_validating_payload_builder: Arc<dyn SelfValidatingPayloadBuilder>,
     query_stats_payload_builder: Box<dyn BatchPayloadBuilder>,