apollo_propeller: replace local Stake type alias with StakingWeight from starknet_api (#13037)

Author: sirandreww-starkware

Cargo.lock

@@ -20,6 +20,7 @@ prost.workspace = true
 rand.workspace = true
 reed-solomon-simd.workspace = true
 sha2.workspace = true
+starknet_api.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 tracing.workspace = true

crates/apollo_propeller/Cargo.toml

@@ -23,13 +23,13 @@ use libp2p::swarm::{
     THandlerOutEvent,
     ToSwarm,
 };
+use starknet_api::staking::StakingWeight;
 use tokio::sync::{mpsc, oneshot};
 
 use crate::config::Config;
 use crate::engine::{Engine, EngineCommand, EngineOutput};
 use crate::handler::Handler;
 use crate::metrics::PropellerMetrics;
-use crate::tree::Stake;
 use crate::types::{CommitteeId, CommitteeSetupError, Event, ShardPublishError};
 
 /// The Propeller network behaviour.
@@ -76,7 +76,7 @@ impl Behaviour {
     pub fn register_committee_peers(
         &self,
         committee_id: CommitteeId,
-        peers: Vec<(PeerId, Stake)>,
+        peers: Vec<(PeerId, StakingWeight)>,
     ) -> oneshot::Receiver<Result<(), CommitteeSetupError>> {
         self.register_committee_peers_and_optional_keys(
             committee_id,
@@ -91,7 +91,7 @@ impl Behaviour {
     pub fn register_committee_peers_and_optional_keys(
         &self,
         committee_id: CommitteeId,
-        peers: Vec<(PeerId, Stake, Option<PublicKey>)>,
+        peers: Vec<(PeerId, StakingWeight, Option<PublicKey>)>,
     ) -> oneshot::Receiver<Result<(), CommitteeSetupError>> {
         let (response_tx, response_rx) = oneshot::channel();
         let command =

crates/apollo_propeller/src/behaviour.rs

@@ -8,6 +8,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use libp2p::identity::{Keypair, PeerId, PublicKey};
+use starknet_api::staking::StakingWeight;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, error, trace, warn};
 
@@ -18,7 +19,7 @@ use crate::metrics::PropellerMetrics;
 use crate::sharding::create_units_to_publish;
 use crate::signature;
 use crate::time_cache::TimeCache;
-use crate::tree::{PropellerScheduleManager, Stake};
+use crate::tree::PropellerScheduleManager;
 use crate::types::{CommitteeId, CommitteeSetupError, Event, MessageRoot, ShardPublishError};
 use crate::unit::PropellerUnit;
 
@@ -36,7 +37,7 @@ struct MessageKey {
 pub enum EngineCommand {
     RegisterCommitteePeers {
         committee_id: CommitteeId,
-        peers: Vec<(PeerId, Stake, Option<PublicKey>)>,
+        peers: Vec<(PeerId, StakingWeight, Option<PublicKey>)>,
         response: oneshot::Sender<Result<(), CommitteeSetupError>>,
     },
     // TODO(AndrewL): remove this variant once unregister is no longer needed.
@@ -135,7 +136,7 @@ impl Engine {
     pub fn register_committee(
         &mut self,
         committee_id: CommitteeId,
-        peers: Vec<(PeerId, Stake, Option<PublicKey>)>,
+        peers: Vec<(PeerId, StakingWeight, Option<PublicKey>)>,
     ) -> Result<(), CommitteeSetupError> {
         if self.committees.contains_key(&committee_id) {
             warn!(?committee_id, "Committee already registered, ignoring re-registration");

crates/apollo_propeller/src/engine.rs

@@ -40,7 +40,7 @@ pub use config::Config;
 pub use handler::{Handler, HandlerIn, HandlerOut};
 pub use merkle::{MerkleHash, MerkleProof, MerkleTree};
 pub use metrics::PropellerMetrics;
-pub use tree::{PropellerScheduleManager, Stake};
+pub use tree::PropellerScheduleManager;
 pub use types::{
     CommitteeId,
     CommitteeSetupError,

crates/apollo_propeller/src/lib.rs

@@ -5,12 +5,11 @@
 //! based on the publisher and shard ID, making the network resilient to targeted attacks.
 
 use libp2p::identity::PeerId;
+use starknet_api::staking::StakingWeight;
 
 use crate::types::{CommitteeSetupError, ScheduleError, ShardIndex};
 use crate::ShardValidationError;
 
-pub type Stake = u64;
-
 // TODO(AndrewL): add the concept of shard_owner when naming
 
 /// Propeller tree manager that computes tree topology on-demand for each publisher.
@@ -25,7 +24,7 @@ pub type Stake = u64;
 #[derive(Debug, Clone)]
 pub struct PropellerScheduleManager {
     /// All nodes in the committee with their stake, sorted by peer_id
-    committee_nodes: Vec<(PeerId, Stake)>,
+    committee_nodes: Vec<(PeerId, StakingWeight)>,
     /// This node's peer ID.
     local_peer_id: PeerId,
     /// This node's index in the nodes vector.
@@ -40,7 +39,7 @@ impl PropellerScheduleManager {
     // TODO(AndrewL): What should I name the error type?
     pub fn new(
         local_peer_id: PeerId,
-        mut nodes: Vec<(PeerId, Stake)>,
+        mut nodes: Vec<(PeerId, StakingWeight)>,
     ) -> Result<Self, CommitteeSetupError> {
         // Check that local peer is in the list before sorting
         if !nodes.iter().any(|(peer_id, _)| *peer_id == local_peer_id) {
@@ -81,7 +80,7 @@ impl PropellerScheduleManager {
         self.committee_nodes.len()
     }
 
-    pub fn get_nodes(&self) -> &[(PeerId, Stake)] {
+    pub fn get_nodes(&self) -> &[(PeerId, StakingWeight)] {
         &self.committee_nodes
     }
 

crates/apollo_propeller/src/tree.rs

@@ -1,5 +1,6 @@
 use libp2p::PeerId;
 use rstest::rstest;
+use starknet_api::staking::StakingWeight;
 
 use crate::tree::PropellerScheduleManager;
 use crate::types::{CommitteeSetupError, ScheduleError, ShardIndex, ShardValidationError};
@@ -18,7 +19,7 @@ fn make_schedule_manager(index: u8, num_nodes: u8) -> PropellerScheduleManager {
     let local_peer = peers[usize::from(index)];
     let scheduler = PropellerScheduleManager::new(
         local_peer,
-        peers.into_iter().map(|peer| (peer, 1)).collect(),
+        peers.into_iter().map(|peer| (peer, StakingWeight(1))).collect(),
     )
     .unwrap();
     assert_eq!(scheduler.get_nodes()[usize::from(index)].0, local_peer);
@@ -34,7 +35,7 @@ fn test_create_empty_schedule_manager() {
 #[test]
 fn test_create_schedule_manager() {
     let peer = PeerId::random();
-    let manager = PropellerScheduleManager::new(peer, vec![(peer, 100)]).unwrap();
+    let manager = PropellerScheduleManager::new(peer, vec![(peer, StakingWeight(100))]).unwrap();
     assert_eq!(manager.get_local_peer_id(), peer);
     assert_eq!(manager.get_node_count(), 1);
 }
@@ -73,7 +74,7 @@ fn test_should_build_and_receive(
 fn test_new_schedule_manager_without_local_peer() {
     let peer1 = PeerId::random();
     let peer2 = PeerId::random();
-    let result = PropellerScheduleManager::new(peer1, vec![(peer2, 100)]);
+    let result = PropellerScheduleManager::new(peer1, vec![(peer2, StakingWeight(100))]);
     assert_eq!(result.unwrap_err(), CommitteeSetupError::LocalPeerNotInCommittee);
 }
 
@@ -206,7 +207,12 @@ fn test_get_my_shard_index_given_publisher() {
     let (peer0, peer1, peer2, peer3) = (peers[0], peers[1], peers[2], peers[3]);
     let manager = PropellerScheduleManager::new(
         peer2,
-        vec![(peer0, 100), (peer1, 75), (peer2, 50), (peer3, 25)],
+        vec![
+            (peer0, StakingWeight(100)),
+            (peer1, StakingWeight(75)),
+            (peer2, StakingWeight(50)),
+            (peer3, StakingWeight(25)),
+        ],
     )
     .unwrap();
     // When peer0 is publisher, peer2 is at sorted position 2, so shard index is 1

crates/apollo_propeller/src/tree_test.rs

@@ -4,6 +4,7 @@ use std::sync::Arc;
 use libp2p::identity::Keypair;
 use libp2p::PeerId;
 use rstest::rstest;
+use starknet_api::staking::StakingWeight;
 
 use crate::types::ShardSignatureVerificationError;
 use crate::{
@@ -46,8 +47,8 @@ fn build_env(num_shards_per_peer: usize) -> TestEnv {
     let local_peer = PeerId::random();
     let other_peers = (0..(NUM_PEERS - 2)).map(|_| PeerId::random()).collect::<Vec<_>>();
 
-    let mut peers = vec![(local_peer, 10), (publisher, 10)];
-    peers.extend(other_peers.iter().copied().map(|peer| (peer, 10)));
+    let mut peers = vec![(local_peer, StakingWeight(10)), (publisher, StakingWeight(10))];
+    peers.extend(other_peers.iter().copied().map(|peer| (peer, StakingWeight(10))));
 
     let schedule_manager = Arc::new(PropellerScheduleManager::new(local_peer, peers).unwrap());
 

crates/apollo_propeller/src/unit_validator_test.rs

@@ -8,6 +8,7 @@ use libp2p::swarm::SwarmEvent;
 use libp2p::{PeerId, Swarm};
 use libp2p_swarm_test::SwarmExt;
 use rstest::rstest;
+use starknet_api::staking::StakingWeight;
 use tokio::sync::mpsc;
 use tracing_test::traced_test;
 
@@ -45,8 +46,10 @@ impl TestSetup {
     }
 
     async fn register_committee(&mut self, committee_id: CommitteeId, node_indices: &[NodeIndex]) {
-        let peers: Vec<(PeerId, u64)> =
-            node_indices.iter().map(|&i| (*self.swarms[i].local_peer_id(), 1)).collect();
+        let peers: Vec<(PeerId, StakingWeight)> = node_indices
+            .iter()
+            .map(|&i| (*self.swarms[i].local_peer_id(), StakingWeight(1)))
+            .collect();
         for &i in node_indices {
             self.swarms[i]
                 .behaviour_mut()