refactor: remove the ArtifactChunk type

Author: rumenov

rs/p2p/state_sync_manager/src/ongoing.rs

@@ -329,7 +329,7 @@ impl OngoingStateSync {
 
         let chunk_add_result = tokio::task::spawn_blocking(move || {
             let chunk = parse_chunk_handler_response(response, chunk_id, metrics)?;
-            Ok(tracker.lock().unwrap().add_chunk(chunk))
+            Ok(tracker.lock().unwrap().add_chunk(chunk_id, chunk))
         })
         .await
         .map_err(|err| DownloadChunkError::RequestError {

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

@@ -13,7 +13,7 @@ use ic_logger::ReplicaLogger;
 use ic_protobuf::p2p::v1 as pb;
 use ic_types::{
     artifact::StateSyncArtifactId,
-    chunkable::{ArtifactChunk, ChunkId},
+    chunkable::{Chunk, ChunkId},
 };
 use prost::Message;
 
@@ -100,7 +100,7 @@ pub(crate) fn parse_chunk_handler_response(
     response: Response<Bytes>,
     chunk_id: ChunkId,
     metrics: OngoingStateSyncMetrics,
-) -> Result<ArtifactChunk, DownloadChunkError> {
+) -> Result<Chunk, DownloadChunkError> {
     let (parts, body) = response.into_parts();
 
     match parts.status {
@@ -127,11 +127,7 @@ pub(crate) fn parse_chunk_handler_response(
                     }
                 })?;
 
-            let chunk = ArtifactChunk {
-                chunk_id,
-                chunk: pb.data,
-            };
-            Ok(chunk)
+            Ok(pb.data)
         }
         StatusCode::NO_CONTENT => Err(DownloadChunkError::NoContent),
         StatusCode::TOO_MANY_REQUESTS => Err(DownloadChunkError::Overloaded),

rs/p2p/state_sync_manager/tests/common.rs

@@ -16,7 +16,7 @@ use ic_metrics::MetricsRegistry;
 use ic_p2p_test_utils::mocks::{MockChunkable, MockStateSync};
 use ic_types::{
     artifact::{StateSyncArtifactId, StateSyncMessage},
-    chunkable::{ArtifactChunk, ArtifactErrorCode, Chunk, ChunkId, Chunkable},
+    chunkable::{ArtifactErrorCode, Chunk, ChunkId, Chunkable},
     crypto::CryptoHash,
     state_sync::{Manifest, MetaManifest, StateSyncVersion},
     CryptoHashOfState, Height, NodeId, PrincipalId,
@@ -281,23 +281,23 @@ impl Chunkable for FakeChunkable {
 
     fn add_chunk(
         &mut self,
-        artifact_chunk: ArtifactChunk,
+        chunk_id: ChunkId,
+        chunk: Chunk,
     ) -> Result<StateSyncMessage, ArtifactErrorCode> {
         for set in self.chunk_sets.iter_mut() {
             if set.is_empty() {
                 continue;
             }
-            if set.remove(&artifact_chunk.chunk_id) {
+            if set.remove(&chunk_id) {
                 break;
             } else {
-                panic!("Downloaded chunk {} twice", artifact_chunk.chunk_id)
+                panic!("Downloaded chunk {} twice", chunk_id)
             }
         }
 
         // Add chunk to state if not part of manifest
-        if !is_manifest_chunk(artifact_chunk.chunk_id) {
-            self.local_state
-                .add_chunk(artifact_chunk.chunk_id, artifact_chunk.chunk.len())
+        if !is_manifest_chunk(chunk_id) {
+            self.local_state.add_chunk(chunk_id, chunk.len())
         }
 
         let elems = self.chunk_sets.iter().map(|set| set.len()).sum::<usize>();
@@ -342,10 +342,11 @@ impl Chunkable for SharableMockChunkable {
     }
     fn add_chunk(
         &mut self,
-        artifact_chunk: ArtifactChunk,
+        chunk_id: ChunkId,
+        chunk: Chunk,
     ) -> Result<StateSyncMessage, ArtifactErrorCode> {
         self.add_chunks_calls.fetch_add(1, Ordering::SeqCst);
-        self.mock.lock().unwrap().add_chunk(artifact_chunk)
+        self.mock.lock().unwrap().add_chunk(chunk_id, chunk)
     }
 }
 

rs/p2p/state_sync_manager/tests/test.rs

@@ -735,7 +735,7 @@ fn test_state_sync_abortion() {
         c2.clear();
         c2.get_mut()
             .expect_add_chunk()
-            .returning(|_| Err(ArtifactErrorCode::ChunkVerificationFailed));
+            .returning(|_, _| Err(ArtifactErrorCode::ChunkVerificationFailed));
         c2.get_mut()
             .expect_chunks_to_download()
             .returning(|| Box::new(vec![ChunkId::from(1)].into_iter()) as Box<_>);

rs/p2p/test_utils/src/mocks.rs

@@ -7,7 +7,7 @@ use ic_interfaces::p2p::{
     state_sync::StateSyncClient,
 };
 use ic_quic_transport::{ConnId, SendError, Transport};
-use ic_types::{artifact::PriorityFn, chunkable::ArtifactChunk};
+use ic_types::artifact::PriorityFn;
 use ic_types::{
     artifact::{StateSyncArtifactId, StateSyncMessage},
     chunkable::{ArtifactErrorCode, Chunkable},
@@ -61,7 +61,7 @@ mock! {
 
     impl Chunkable for Chunkable{
         fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>>;
-        fn add_chunk(&mut self, artifact_chunk: ArtifactChunk) -> Result<StateSyncMessage, ArtifactErrorCode>;
+        fn add_chunk(&mut self, chunk_id: ChunkId, chunk: Chunk) -> Result<StateSyncMessage, ArtifactErrorCode>;
     }
 }
 

rs/state_manager/src/state_sync/chunkable.rs

@@ -13,9 +13,8 @@ use ic_sys::mmap::ScopedMmap;
 use ic_types::{
     artifact::StateSyncMessage,
     chunkable::{
-        ArtifactChunk,
         ArtifactErrorCode::{self, ChunkVerificationFailed, ChunksMoreNeeded},
-        ChunkId, Chunkable,
+        Chunk, ChunkId, Chunkable,
     },
     malicious_flags::MaliciousFlags,
     state_sync::{
@@ -1159,21 +1158,22 @@ impl Chunkable for IncompleteState {
 
     fn add_chunk(
         &mut self,
-        artifact_chunk: ArtifactChunk,
+        chunk_id: ChunkId,
+        chunk: Chunk,
     ) -> Result<StateSyncMessage, ArtifactErrorCode> {
-        let ix = artifact_chunk.chunk_id.get();
-        let payload = &artifact_chunk.chunk;
+        let ix = chunk_id.get();
+        let payload = &chunk;
         match &mut self.state {
             DownloadState::Complete(ref artifact) => {
                 debug!(
                     self.log,
-                    "Received chunk {} on completed state {}", artifact_chunk.chunk_id, self.height
+                    "Received chunk {} on completed state {}", chunk_id, self.height
                 );
 
                 Ok(*artifact.clone())
             }
             DownloadState::Blank => {
-                if artifact_chunk.chunk_id == META_MANIFEST_CHUNK {
+                if chunk_id == META_MANIFEST_CHUNK {
                     let meta_manifest = decode_meta_manifest(payload).map_err(|err| {
                         warn!(
                             self.log,

rs/state_manager/tests/common/mod.rs

@@ -22,9 +22,8 @@ use ic_test_utilities_tmpdir::tmpdir;
 use ic_types::{
     artifact::StateSyncMessage,
     chunkable::{
-        ArtifactChunk,
         ArtifactErrorCode::{ChunkVerificationFailed, ChunksMoreNeeded},
-        ChunkId, Chunkable,
+        Chunk, ChunkId, Chunkable,
     },
     consensus::certification::{Certification, CertificationContent},
     crypto::Signed,
@@ -394,8 +393,8 @@ pub fn pipe_state_sync(src: StateSyncMessage, mut dst: Box<dyn Chunkable>) -> St
         .expect("State sync not completed.")
 }
 
-fn alter_chunk_data(chunk: &mut ArtifactChunk) {
-    let mut chunk_data = chunk.chunk.clone();
+fn alter_chunk_data(chunk: &mut Chunk) {
+    let mut chunk_data = chunk.clone();
     match chunk_data.last_mut() {
         Some(last) => {
             // Alter the last element of chunk_data.
@@ -406,7 +405,7 @@ fn alter_chunk_data(chunk: &mut ArtifactChunk) {
             chunk_data = vec![9; 100];
         }
     }
-    chunk.chunk = chunk_data;
+    *chunk = chunk_data;
 }
 
 /// Pipe the meta-manifest (chunk 0) from src to dest.
@@ -423,19 +422,16 @@ pub fn pipe_meta_manifest(
 
     let id = ids[0];
 
-    let mut chunk = ArtifactChunk {
-        chunk_id: id,
-        chunk: src
-            .clone()
-            .get_chunk(id)
-            .unwrap_or_else(|| panic!("Requested unknown chunk {}", id)),
-    };
+    let mut chunk = src
+        .clone()
+        .get_chunk(id)
+        .unwrap_or_else(|| panic!("Requested unknown chunk {}", id));
 
     if use_bad_chunk {
         alter_chunk_data(&mut chunk);
     }
 
-    match dst.add_chunk(chunk) {
+    match dst.add_chunk(id, chunk) {
         Ok(msg) => Ok(msg),
         Err(ChunksMoreNeeded) => Err(StateSyncErrorCode::ChunksMoreNeeded),
         Err(ChunkVerificationFailed) => Err(StateSyncErrorCode::MetaManifestVerificationFailed),
@@ -460,19 +456,16 @@ pub fn pipe_manifest(
     assert!(ids.iter().all(|id| manifest_chunks.contains(id)));
 
     for (index, id) in ids.iter().enumerate() {
-        let mut chunk = ArtifactChunk {
-            chunk_id: *id,
-            chunk: src
-                .clone()
-                .get_chunk(*id)
-                .unwrap_or_else(|| panic!("Requested unknown chunk {}", id)),
-        };
+        let mut chunk = src
+            .clone()
+            .get_chunk(*id)
+            .unwrap_or_else(|| panic!("Requested unknown chunk {}", id));
 
         if use_bad_chunk && index == ids.len() / 2 {
             alter_chunk_data(&mut chunk);
         }
 
-        match dst.add_chunk(chunk) {
+        match dst.add_chunk(*id, chunk) {
             Ok(msg) => {
                 return Ok(msg);
             }
@@ -506,19 +499,16 @@ pub fn pipe_partial_state_sync(
                 omitted_chunks = true;
                 continue;
             }
-            let mut chunk = ArtifactChunk {
-                chunk_id: *id,
-                chunk: src
-                    .clone()
-                    .get_chunk(*id)
-                    .unwrap_or_else(|| panic!("Requested unknown chunk {}", id)),
-            };
+            let mut chunk = src
+                .clone()
+                .get_chunk(*id)
+                .unwrap_or_else(|| panic!("Requested unknown chunk {}", id));
 
             if use_bad_chunk && index == ids.len() / 2 {
                 alter_chunk_data(&mut chunk);
             }
 
-            match dst.add_chunk(chunk) {
+            match dst.add_chunk(*id, chunk) {
                 Ok(msg) => {
                     return Ok(msg);
                 }

rs/types/types/src/chunkable.rs

@@ -12,7 +12,6 @@
 //! For more context please check `<https://youtu.be/WaNJINjGleg>`
 //!
 use crate::artifact::{Artifact, StateSyncMessage};
-use ic_protobuf::p2p::v1 as p2p_pb;
 use phantom_newtype::Id;
 
 pub type Chunk = Vec<u8>;
@@ -25,16 +24,8 @@ pub enum ArtifactErrorCode {
 }
 
 /// The chunk type.
-pub type ChunkId = Id<ArtifactChunk, u32>;
-
-/// An artifact chunk.
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
-pub struct ArtifactChunk {
-    // Chunk number/id for this chunk
-    pub chunk_id: ChunkId,
-    // Payload for the chunk
-    pub chunk: Chunk,
-}
+pub struct ChunkIdTag;
+pub type ChunkId = Id<ChunkIdTag, u32>;
 
 /// Interface providing access to artifact.
 pub trait ChunkableArtifact {
@@ -45,12 +36,7 @@ pub trait Chunkable {
     fn chunks_to_download(&self) -> Box<dyn Iterator<Item = ChunkId>>;
     fn add_chunk(
         &mut self,
-        artifact_chunk: ArtifactChunk,
+        chunk_id: ChunkId,
+        chunk: Chunk,
     ) -> Result<StateSyncMessage, ArtifactErrorCode>;
 }
-
-impl From<ArtifactChunk> for p2p_pb::StateSyncChunkResponse {
-    fn from(chunk: ArtifactChunk) -> Self {
-        Self { data: chunk.chunk }
-    }
-}