refactor: make the Chunkable trait generic

Author: rumenov

rs/interfaces/src/p2p/state_sync.rs

@@ -15,7 +15,7 @@ pub trait StateSyncClient: Send + Sync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable + Send + Sync>>;
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>;
     /// 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.

rs/p2p/state_sync_manager/src/ongoing.rs

@@ -70,7 +70,7 @@ struct OngoingStateSync {
     downloading_chunks: JoinMap<ChunkId, DownloadResult>,
     // State sync
     state_sync: Arc<dyn StateSyncClient>,
-    tracker: Arc<Mutex<Box<dyn Chunkable + Send + Sync>>>,
+    tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
     state_sync_finished: bool,
 }
 
@@ -88,7 +88,7 @@ pub(crate) fn start_ongoing_state_sync(
     log: ReplicaLogger,
     rt: &Handle,
     metrics: OngoingStateSyncMetrics,
-    tracker: Arc<Mutex<Box<dyn Chunkable + Send + Sync>>>,
+    tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
     artifact_id: StateSyncArtifactId,
     state_sync: Arc<dyn StateSyncClient>,
     transport: Arc<dyn Transport>,
@@ -296,7 +296,7 @@ impl OngoingStateSync {
     async fn download_chunk_task(
         peer_id: NodeId,
         client: Arc<dyn Transport>,
-        tracker: Arc<Mutex<Box<dyn Chunkable + Send + Sync>>>,
+        tracker: Arc<Mutex<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>>,
         artifact_id: StateSyncArtifactId,
         chunk_id: ChunkId,
         metrics: OngoingStateSyncMetrics,
@@ -409,7 +409,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::default();
+            let mut c = MockChunkable::<StateSyncMessage>::default();
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));
 
@@ -448,7 +448,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::default();
+            let mut c = MockChunkable::<StateSyncMessage>::default();
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));
             c.expect_add_chunk()
@@ -494,7 +494,7 @@ mod tests {
                     .body(compress_empty_bytes())
                     .unwrap())
             });
-            let mut c = MockChunkable::default();
+            let mut c = MockChunkable::<StateSyncMessage>::default();
             // Endless iterator
             c.expect_chunks_to_download()
                 .returning(|| Box::new(std::iter::once(ChunkId::from(1))));

rs/p2p/state_sync_manager/tests/common.rs

@@ -197,7 +197,7 @@ impl StateSyncClient for FakeStateSync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable + Send + Sync>> {
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>> {
         if !self.uses_global() && id.height > self.current_height() && !self.disconnected() {
             return Some(Box::new(FakeChunkable::new(
                 self.local_state.clone(),
@@ -263,7 +263,7 @@ impl FakeChunkable {
     }
 }
 
-impl Chunkable for FakeChunkable {
+impl Chunkable<StateSyncMessage> for FakeChunkable {
     /// Returns iterator for chunks to download.
     /// Tries to first download metamanifest, then manifest and then chunks. Does not advance if previous didn't complete.
     fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>> {
@@ -311,7 +311,7 @@ impl Chunkable for FakeChunkable {
 
 #[derive(Clone, Default)]
 pub struct SharableMockChunkable {
-    mock: Arc<Mutex<MockChunkable>>,
+    mock: Arc<Mutex<MockChunkable<StateSyncMessage>>>,
     chunks_to_download_calls: Arc<AtomicUsize>,
     add_chunks_calls: Arc<AtomicUsize>,
 }
@@ -323,7 +323,7 @@ impl SharableMockChunkable {
             ..Default::default()
         }
     }
-    pub fn get_mut(&self) -> MutexGuard<'_, MockChunkable> {
+    pub fn get_mut(&self) -> MutexGuard<'_, MockChunkable<StateSyncMessage>> {
         self.mock.lock().unwrap()
     }
     pub fn add_chunks_calls(&self) -> usize {
@@ -335,7 +335,7 @@ impl SharableMockChunkable {
     }
 }
 
-impl Chunkable for SharableMockChunkable {
+impl Chunkable<StateSyncMessage> for SharableMockChunkable {
     fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>> {
         self.chunks_to_download_calls.fetch_add(1, Ordering::SeqCst);
         self.mock.lock().unwrap().chunks_to_download()
@@ -390,7 +390,7 @@ impl StateSyncClient for SharableMockStateSync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable + Send + Sync>> {
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>> {
         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

@@ -26,7 +26,7 @@ mock! {
         fn start_state_sync(
             &self,
             id: &StateSyncArtifactId,
-        ) -> Option<Box<dyn Chunkable + Send + Sync>>;
+        ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>>;
 
         fn should_cancel(&self, id: &StateSyncArtifactId) -> bool;
 
@@ -58,11 +58,11 @@ mock! {
 }
 
 mock! {
-    pub Chunkable {}
+    pub Chunkable<T> {}
 
-    impl Chunkable for Chunkable{
+    impl<T> Chunkable<T> for Chunkable<T> {
         fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>>;
-        fn add_chunk(&mut self, chunk_id: ChunkId, chunk: Chunk) -> Result<StateSyncMessage, ArtifactErrorCode>;
+        fn add_chunk(&mut self, chunk_id: ChunkId, chunk: Chunk) -> Result<T, ArtifactErrorCode>;
     }
 }
 

rs/state_manager/src/state_sync.rs

@@ -34,7 +34,7 @@ impl StateSync {
     pub fn create_chunkable_state(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Box<dyn Chunkable + Send + Sync> {
+    ) -> Box<dyn Chunkable<StateSyncMessage> + Send + Sync> {
         info!(self.log, "Starting state sync @{}", id.height);
 
         Box::new(crate::state_sync::chunkable::IncompleteState::new(
@@ -239,7 +239,7 @@ impl StateSyncClient for StateSync {
     fn start_state_sync(
         &self,
         id: &StateSyncArtifactId,
-    ) -> Option<Box<dyn Chunkable + Send + Sync>> {
+    ) -> Option<Box<dyn Chunkable<StateSyncMessage> + Send + Sync>> {
         if self.get_priority_function()(id, &()) != Priority::Fetch {
             return None;
         }

rs/state_manager/src/state_sync/chunkable.rs

@@ -1125,7 +1125,7 @@ impl IncompleteState {
     }
 }
 
-impl Chunkable for IncompleteState {
+impl Chunkable<StateSyncMessage> for IncompleteState {
     fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>> {
         match self.state {
             DownloadState::Blank => Box::new(std::iter::once(META_MANIFEST_CHUNK)),

rs/state_manager/tests/common/mod.rs

@@ -390,7 +390,10 @@ pub enum StateSyncErrorCode {
     OtherChunkVerificationFailed,
 }
 
-pub fn pipe_state_sync(src: StateSyncMessage, mut dst: Box<dyn Chunkable>) -> StateSyncMessage {
+pub fn pipe_state_sync(
+    src: StateSyncMessage,
+    mut dst: Box<dyn Chunkable<StateSyncMessage>>,
+) -> StateSyncMessage {
     pipe_partial_state_sync(&src, &mut *dst, &Default::default(), false)
         .expect("State sync not completed.")
 }
@@ -414,7 +417,7 @@ fn alter_chunk_data(chunk: &mut Chunk) {
 /// Alter the chunk data if `use_bad_chunk` is set to true.
 pub fn pipe_meta_manifest(
     src: &StateSyncMessage,
-    dst: &mut dyn Chunkable,
+    dst: &mut dyn Chunkable<StateSyncMessage>,
     use_bad_chunk: bool,
 ) -> Result<StateSyncMessage, StateSyncErrorCode> {
     let ids: Vec<_> = dst.chunks_to_download().collect();
@@ -445,7 +448,7 @@ pub fn pipe_meta_manifest(
 /// Alter the data of the chunk in the middle position if `use_bad_chunk` is set to true.
 pub fn pipe_manifest(
     src: &StateSyncMessage,
-    dst: &mut dyn Chunkable,
+    dst: &mut dyn Chunkable<StateSyncMessage>,
     use_bad_chunk: bool,
 ) -> Result<StateSyncMessage, StateSyncErrorCode> {
     let ids: Vec<_> = dst.chunks_to_download().collect();
@@ -484,7 +487,7 @@ pub fn pipe_manifest(
 /// Alter the data of the chunk in the middle position if `use_bad_chunk` is set to true.
 pub fn pipe_partial_state_sync(
     src: &StateSyncMessage,
-    dst: &mut dyn Chunkable,
+    dst: &mut dyn Chunkable<StateSyncMessage>,
     omit: &HashSet<ChunkId>,
     use_bad_chunk: bool,
 ) -> Result<StateSyncMessage, StateSyncErrorCode> {

rs/types/types/src/chunkable.rs

@@ -11,7 +11,6 @@
 //!
 //! For more context please check `<https://youtu.be/WaNJINjGleg>`
 //!
-use crate::state_sync::StateSyncMessage;
 use phantom_newtype::Id;
 
 pub type Chunk = Vec<u8>;
@@ -27,11 +26,7 @@ pub enum ArtifactErrorCode {
 pub struct ChunkIdTag;
 pub type ChunkId = Id<ChunkIdTag, u32>;
 
-pub trait Chunkable {
+pub trait Chunkable<T> {
     fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>>;
-    fn add_chunk(
-        &mut self,
-        chunk_id: ChunkId,
-        chunk: Chunk,
-    ) -> Result<StateSyncMessage, ArtifactErrorCode>;
+    fn add_chunk(&mut self, chunk_id: ChunkId, chunk: Chunk) -> Result<T, ArtifactErrorCode>;
 }