Repository: mintlayer/mintlayer-core
Pull request #777 (2023-03-21 23:21:12)

Add the `get_locator_from_height` function


Patch:
diff --git a/chainstate/src/detail/query.rs b/chainstate/src/detail/query.rs
index f5ae243f6..d25429e47 100644
--- a/chainstate/src/detail/query.rs
+++ b/chainstate/src/detail/query.rs
@@ -110,7 +110,13 @@ impl<'a, S: BlockchainStorageRead, V: TransactionVerificationStrategy> Chainstat
             .get_best_block_index()?
             .ok_or(PropertyQueryError::BestBlockIndexNotFound)?;
         let height = best_block_index.block_height();
+        self.get_locator_from_height(height)
+    }
 
+    pub fn get_locator_from_height(
+        &self,
+        height: BlockHeight,
+    ) -> Result<Locator, PropertyQueryError> {
         let headers = locator_tip_distances()
             .map_while(|dist| height - dist)
             .map(|ht| self.chainstate_ref.get_block_id_by_height(&ht));
diff --git a/chainstate/src/interface/chainstate_interface.rs b/chainstate/src/interface/chainstate_interface.rs
index ee35b9e0d..6f0351b5d 100644
--- a/chainstate/src/interface/chainstate_interface.rs
+++ b/chainstate/src/interface/chainstate_interface.rs
@@ -63,6 +63,9 @@ pub trait ChainstateInterface: Send {
     /// to quickly find a common ancestor between different chains.
     fn get_locator(&self) -> Result<Locator, ChainstateError>;
 
+    /// Returns a locator starting from the specified height.
+    fn get_locator_from_height(&self, height: BlockHeight) -> Result<Locator, ChainstateError>;
+
     /// Returns a list of block headers starting from the last locator's block that is in the main
     /// chain.
     ///
diff --git a/chainstate/src/interface/chainstate_interface_impl.rs b/chainstate/src/interface/chainstate_interface_impl.rs
index 8a76fb779..a32bc4d47 100644
--- a/chainstate/src/interface/chainstate_interface_impl.rs
+++ b/chainstate/src/interface/chainstate_interface_impl.rs
@@ -143,6 +143,14 @@ impl<S: BlockchainStorage, V: TransactionVerificationStrategy> ChainstateInterfa
             .map_err(ChainstateError::FailedToReadProperty)
     }
 
+    fn get_locator_from_height(&self, height: BlockHeight) -> Result<Locator, ChainstateError> {
+        self.chainstate
+            .query()
+            .map_err(ChainstateError::from)?
+            .get_locator_from_height(height)
+            .map_err(ChainstateError::FailedToReadProperty)
+    }
+
     fn get_headers(
         &self,
         locator: Locator,
diff --git a/chainstate/src/interface/chainstate_interface_impl_delegation.rs b/chainstate/src/interface/chainstate_interface_impl_delegation.rs
index 1e1789a1d..1fe1e2c6b 100644
--- a/chainstate/src/interface/chainstate_interface_impl_delegation.rs
+++ b/chainstate/src/interface/chainstate_interface_impl_delegation.rs
@@ -106,6 +106,10 @@ impl<
         self.deref().get_locator()
     }
 
+    fn get_locator_from_height(&self, height: BlockHeight) -> Result<Locator, ChainstateError> {
+        self.deref().get_locator_from_height(height)
+    }
+
     fn get_headers(
         &self,
         locator: Locator,
diff --git a/mocks/src/chainstate.rs b/mocks/src/chainstate.rs
index 19de73fb2..fac32d95d 100644
--- a/mocks/src/chainstate.rs
+++ b/mocks/src/chainstate.rs
@@ -64,6 +64,7 @@ mockall::mock! {
         ) -> Result<Option<Id<GenBlock>>, ChainstateError>;
         fn get_block(&self, block_id: Id<Block>) -> Result<Option<Block>, ChainstateError>;
         fn get_locator(&self) -> Result<Locator, ChainstateError>;
+        fn get_locator_from_height(&self, height: BlockHeight) -> Result<Locator, ChainstateError>;
         fn get_headers(
             &self,
             locator: Locator,
diff --git a/p2p/src/sync/peer.rs b/p2p/src/sync/peer.rs
index 342474870..13d516248 100644
--- a/p2p/src/sync/peer.rs
+++ b/p2p/src/sync/peer.rs
@@ -277,7 +277,8 @@ where
         }
         log::trace!("Received headers: {headers:#?}");
 
-        // We are in sync with this peer.
+        // TODO: Should the empty headers response be treated as misbehavior if we are going to
+        // send a locator starting with the block preceding the tip?
         if headers.is_empty() {
             return Ok(());
         }


Test patch:
diff --git a/chainstate/test-suite/src/tests/syncing_tests.rs b/chainstate/test-suite/src/tests/syncing_tests.rs
index f79bc721a..63ad708be 100644
--- a/chainstate/test-suite/src/tests/syncing_tests.rs
+++ b/chainstate/test-suite/src/tests/syncing_tests.rs
@@ -75,6 +75,39 @@ fn get_locator(#[case] seed: Seed) {
     });
 }
 
+#[rstest]
+#[trace]
+#[case(Seed::from_entropy())]
+fn get_locator_from_height(#[case] seed: Seed) {
+    utils::concurrency::model(move || {
+        let mut rng = make_seedable_rng(seed);
+        let mut btf = TestFramework::builder(&mut rng).build();
+
+        let blocks = rng.gen_range(1000..2000);
+        btf.create_chain(&btf.genesis().get_id().into(), blocks, &mut rng).unwrap();
+
+        for _ in 0..8 {
+            let height = rng.gen_range(1000..blocks) as u64;
+
+            // Check the locator length.
+            let locator = btf.chainstate.get_locator_from_height(height.into()).unwrap();
+            assert_eq!(locator.len(), (height as f64).log2().ceil() as usize + 1);
+
+            // Check the locator headers.
+            assert_eq!(
+                &locator[0],
+                &btf.chainstate.get_block_id_from_height(&height.into()).unwrap().unwrap()
+            );
+            for (i, header) in locator.iter().skip(1).enumerate() {
+                let idx = BlockHeight::from(height) - BlockDistance::new(2i64.pow(i as u32));
+                let expected =
+                    btf.chainstate.get_block_id_from_height(&idx.unwrap()).unwrap().unwrap();
+                assert_eq!(&expected, header);
+            }
+        }
+    });
+}
+
 // Check that new blocks (produced after a locator is created) are returned.
 #[rstest]
 #[trace]
diff --git a/p2p/src/sync/tests/header_list_request.rs b/p2p/src/sync/tests/header_list_request.rs
index 315d40ac1..d69da848d 100644
--- a/p2p/src/sync/tests/header_list_request.rs
+++ b/p2p/src/sync/tests/header_list_request.rs
@@ -94,7 +94,11 @@ async fn valid_request(#[case] seed: Seed) {
         .with_chain_config(chain_config.as_ref().clone())
         .build();
     // Process a block to finish the initial block download.
-    tf.make_block_builder().build_and_process().unwrap().unwrap();
+    let block_index = tf.make_block_builder().build_and_process().unwrap().unwrap();
+    let locator = tf
+        .chainstate
+        .get_locator_from_height(block_index.block_height().prev_height().unwrap())
+        .unwrap();
     let (chainstate, mempool) =
         start_subsystems_with_chainstate(tf.into_chainstate(), Arc::clone(&chain_config));
 
@@ -109,11 +113,16 @@ async fn valid_request(#[case] seed: Seed) {
 
     handle.send_message(
         peer,
-        SyncMessage::HeaderListRequest(HeaderListRequest::new(Locator::new(Vec::new()))),
+        SyncMessage::HeaderListRequest(HeaderListRequest::new(locator)),
     );
 
     let (sent_to, message) = handle.message().await;
     assert_eq!(peer, sent_to);
-    assert!(matches!(message, SyncMessage::HeaderListResponse(_)));
+    let headers = match message {
+        SyncMessage::HeaderListResponse(r) => r.into_headers(),
+        m => panic!("Unexpected message: {m:?}"),
+    };
+    assert_eq!(headers.len(), 1);
+    assert_eq!(&headers[0], block_index.block_header());
     handle.assert_no_error().await;
 }