feat(cluster): add private IP address to Node (#253)

Author: xDarksome

Cargo.lock

@@ -1 +1 @@
-250917.0
+250923.0

VERSION

@@ -47,7 +47,7 @@ tokio = { workspace = true, features = ["time", "rt"] }
 tokio-stream = { workspace = true }
 
 # Domain specific
-libp2p-identity = { workspace = true, features = ["serde"] }
+libp2p-identity = { workspace = true, features = ["serde", "rand"] }
 xxhash-rust = { workspace = true }
 multihash = { version = "0.19", default-features = false }
 alloy = { workspace = true, features = [

crates/cluster/Cargo.toml

@@ -68,6 +68,18 @@ pub(crate) struct V0 {
     pub secondary_port: u16,
 }
 
+// NOTE: The on-chain serialization is non self-describing!
+// This `struct` can not be changed, a `struct` with a new schema version should
+// be created instead.
+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+pub(crate) struct V1 {
+    pub peer_id: PeerId,
+    pub ipv4_addr: Ipv4Addr,
+    pub private_ipv4_addr: Option<Ipv4Addr>,
+    pub primary_port: u16,
+    pub secondary_port: u16,
+}
+
 impl From<Node> for V0 {
     fn from(node: Node) -> Self {
         Self {
@@ -90,6 +102,29 @@ impl From<V0> for Node {
     }
 }
 
+impl From<Node> for V1 {
+    fn from(node: Node) -> Self {
+        Self {
+            peer_id: node.peer_id,
+            ipv4_addr: node.ipv4_addr,
+            private_ipv4_addr: None,
+            primary_port: node.primary_port,
+            secondary_port: node.secondary_port,
+        }
+    }
+}
+
+impl From<V1> for Node {
+    fn from(node: V1) -> Self {
+        Self {
+            peer_id: node.peer_id,
+            ipv4_addr: node.ipv4_addr,
+            primary_port: node.primary_port,
+            secondary_port: node.secondary_port,
+        }
+    }
+}
+
 struct Ipv4AddrNumeral(FlexibleNumeralString);
 
 impl From<Ipv4Addr> for Ipv4AddrNumeral {

crates/cluster/src/node.rs

@@ -3,10 +3,14 @@
 use {
     crate::{self as cluster, client, node, smart_contract, Client, Config, EncryptionKey, Node},
     derive_more::derive::AsRef,
+    libp2p_identity::PeerId,
     serde::{Deserialize, Serialize},
-    std::sync::{
-        atomic::{self, AtomicUsize},
-        Arc,
+    std::{
+        net::Ipv4Addr,
+        sync::{
+            atomic::{self, AtomicUsize},
+            Arc,
+        },
     },
     tap::Tap,
 };
@@ -189,11 +193,22 @@ impl NodeOperator {
     }
 }
 
-impl<C: Config> NodeOperator<C> {
-    pub(super) fn try_from_sc(
+impl<N> NodeOperator<N> {
+    pub(super) fn from_sc(
         operator: smart_contract::NodeOperator,
-        cfg: &C,
-    ) -> Result<NodeOperator<C::Node>, TryFromSmartContractError> {
+        cfg: &impl Config<Node = N>,
+    ) -> Self {
+        let id = operator.id;
+
+        NodeOperator::try_from_sc(operator, cfg)
+            .map_err(|err| tracing::warn!(%id, ?err, "Invalid NodeOperator"))
+            .unwrap_or_else(|_| NodeOperator::new_corrupted(id, cfg))
+    }
+
+    fn try_from_sc(
+        operator: smart_contract::NodeOperator,
+        cfg: &impl Config<Node = N>,
+    ) -> Result<Self, TryFromSmartContractError> {
         use TryFromSmartContractError as Error;
 
         let data_bytes = operator.data;
@@ -205,22 +220,57 @@ impl<C: Config> NodeOperator<C> {
         let schema_version = data_bytes[0];
         let bytes = &data_bytes[1..];
 
-        let data = match schema_version {
-            0 => postcard::from_bytes::<DataV0>(bytes),
-            ver => return Err(Error::UnknownSchemaVersion(ver)),
+        match schema_version {
+            0 => {
+                let data = postcard::from_bytes::<DataV0>(bytes).map_err(Error::from_postcard)?;
+
+                let nodes = data
+                    .nodes
+                    .into_iter()
+                    .map(|v0| Node::from(v0).tap_mut(|node| node.decrypt(cfg.as_ref())))
+                    .map(|node| cfg.new_node(operator.id, node))
+                    .collect();
+
+                let clients = data.clients.into_iter().map(Into::into).collect();
+
+                Ok(NodeOperator::new(operator.id, data.name, nodes, clients)?)
+            }
+            1 => {
+                let data = postcard::from_bytes::<DataV1>(bytes).map_err(Error::from_postcard)?;
+
+                let nodes = data
+                    .nodes
+                    .into_iter()
+                    .map(|v1| Node::from(v1).tap_mut(|node| node.decrypt(cfg.as_ref())))
+                    .map(|node| cfg.new_node(operator.id, node))
+                    .collect();
+
+                let clients = data.clients.into_iter().map(Into::into).collect();
+
+                Ok(NodeOperator::new(operator.id, data.name, nodes, clients)?)
+            }
+            ver => Err(Error::UnknownSchemaVersion(ver)),
         }
-        .map_err(Error::from_postcard)?;
+    }
 
-        let nodes = data
-            .nodes
-            .into_iter()
-            .map(|v0| Node::from(v0).tap_mut(|node| node.decrypt(cfg.as_ref())))
-            .map(|node| cfg.new_node(operator.id, node))
+    fn new_corrupted(id: Id, cfg: &impl Config<Node = N>) -> Self {
+        let nodes = (0..MIN_NODES)
+            .map(|_| Node {
+                peer_id: PeerId::random(),
+                ipv4_addr: Ipv4Addr::UNSPECIFIED,
+                primary_port: 0,
+                secondary_port: 0,
+            })
+            .map(|node| cfg.new_node(id, node))
             .collect();
 
-        let clients = data.clients.into_iter().map(Into::into).collect();
-
-        Ok(NodeOperator::new(operator.id, data.name, nodes, clients)?)
+        Self {
+            id,
+            name: Name("CORRUPTED_DATA".to_string()),
+            clients: vec![],
+            nodes,
+            counter: Default::default(),
+        }
     }
 }
 
@@ -234,6 +284,16 @@ struct DataV0 {
     clients: Vec<client::V0>,
 }
 
+// NOTE: The on-chain serialization is non self-describing!
+// This `struct` can not be changed, a `struct` with a new schema version should
+// be created instead.
+#[derive(Debug, Clone, Deserialize, Serialize)]
+struct DataV1 {
+    name: Name,
+    nodes: Vec<node::V1>,
+    clients: Vec<client::V0>,
+}
+
 #[derive(Debug, thiserror::Error)]
 pub enum TryIntoSmartContractError {
     #[error("Codec: {0}")]

crates/cluster/src/node_operator.rs

@@ -20,7 +20,6 @@ use {
         Settings,
     },
     derive_where::derive_where,
-    itertools::Itertools,
     std::sync::Arc,
     tap::Pipe as _,
     time::OffsetDateTime,
@@ -174,6 +173,7 @@ impl<C: Config> View<C> {
 
         let time = migration.started_at
             + self.settings.event_propagation_latency
+            + self.settings.clock_skew
             + migration::PULL_DATA_LEEWAY_TIME;
 
         Some(time)
@@ -232,16 +232,11 @@ impl<C: Config> View<C> {
     /// Checks whether a data pull with the specified [`keyspace::Version`] can
     /// be executed at this moment.
     pub fn validate_data_pull(&self, keyspace_version: keyspace::Version) -> bool {
-        if !self.migration.is_in_progress() {
+        let Some(scheduled_after) = self.data_pull_scheduled_after() else {
             return false;
-        }
-
-        let allowed_after = self.migration.started_at
-            + self.settings.event_propagation_latency
-            + self.settings.clock_skew
-            + migration::PULL_DATA_LEEWAY_TIME;
+        };
 
-        OffsetDateTime::now_utc() >= allowed_after && self.keyspace_version == keyspace_version
+        OffsetDateTime::now_utc() >= scheduled_after && self.keyspace_version == keyspace_version
     }
 
     pub(super) fn require_no_migration(&self) -> Result<(), migration::InProgressError> {
@@ -320,11 +315,7 @@ impl<C: Config> View<C> {
         let node_operators = view
             .node_operators
             .into_iter()
-            .map(|slot| {
-                slot.map(|operator| NodeOperator::try_from_sc(operator, cfg))
-                    .transpose()
-            })
-            .try_collect::<_, Vec<_>, _>()?
+            .map(|slot| slot.map(|operator| NodeOperator::from_sc(operator, cfg)))
             .pipe(NodeOperators::from_slots)?;
 
         let ownership = Ownership::new(view.owner);
@@ -470,10 +461,8 @@ impl smart_contract::event::NodeOperatorAdded {
             .require_not_exists(&self.operator.id)?
             .require_free_slot(self.idx)?;
 
-        view.node_operators.set(
-            self.idx,
-            Some(NodeOperator::try_from_sc(self.operator, cfg)?),
-        );
+        view.node_operators
+            .set(self.idx, Some(NodeOperator::from_sc(self.operator, cfg)));
 
         Ok(())
     }
@@ -483,7 +472,7 @@ impl smart_contract::event::NodeOperatorUpdated {
     pub(super) fn apply<C: Config>(self, cfg: &C, view: &mut View<C>) -> Result<()> {
         let idx = view.node_operators.require_idx(&self.operator.id)?;
         view.node_operators
-            .set(idx, Some(NodeOperator::try_from_sc(self.operator, cfg)?));
+            .set(idx, Some(NodeOperator::from_sc(self.operator, cfg)));
 
         Ok(())
     }

crates/cluster/src/view.rs

@@ -19,6 +19,9 @@ backoff = { workspace = true }
 tap = { workspace = true }
 time = { workspace = true }
 
+# Collections
+smallvec = { workspace = true }
+
 # Async
 futures = { workspace = true }
 futures-concurrency = { workspace = true }

crates/migration/Cargo.toml

@@ -3,6 +3,7 @@ use {
     backoff::ExponentialBackoffBuilder,
     futures::{stream, FutureExt as _, StreamExt, TryFutureExt},
     futures_concurrency::future::Race,
+    smallvec::SmallVec,
     std::{
         fmt,
         future::{self, Future},
@@ -14,7 +15,7 @@ use {
     time::OffsetDateTime,
     wcn_cluster::{
         self as cluster,
-        keyspace,
+        keyspace::{self, ReplicaSet},
         migration,
         node_operator,
         smart_contract::{self, Write},
@@ -219,41 +220,43 @@ where
             .secondary_keyspace_shards()
             .ok_or_else(|| anyhow!("Missing secondary keyspace shards"))?;
 
-        let replica_idx = |shard: &keyspace::Shard<&NodeOperator<_>>| {
-            shard
-                .replica_set()
-                .iter()
-                .position(|op| op.id == self.manager.node_operator_id)
-        };
+        // a separate closure because otherwise rustfmt breaks and refuses to format
+        // this function
+        let log_err = |err, shard_id, source_id| tracing::warn!(?err, %shard_id, %source_id, "Failed to transfer shard");
 
         primary_shards
             .zip(secondary_shards)
             .filter_map(|((shard_id, primary), (_, secondary))| {
-                // If operator is not in the new replica set skip this shard.
-                let secondary_idx = replica_idx(&secondary)?;
+                let old = primary.replica_set();
+                let new = secondary.replica_set();
+
+                let added_operators = replica_set_difference(new, old);
+
+                let idx = added_operators
+                    .iter()
+                    .position(|op| op.id == self.manager.node_operator_id)?;
 
-                // Also skip if operator is in the old replica set. Meaning that it's in both,
-                // so we don't need to transfer any data.
-                let None = replica_idx(&primary) else {
-                    return None;
-                };
+                // Skip the shard if this node operator is not being newly added to it.
+                let removed_operators = replica_set_difference(old, new);
 
-                // Pull data only from the operator previously occupying the same replica index
-                // to guarantee consistency.
+                // Pull data only from the operator being replaced to ensure consistency.
                 // TODO: We should have a way to override this behaviour for special
                 // circumstances, for example if a node operator is completely dead and we are
                 // forcefully removing it.
-                let source = primary.replica_set()[secondary_idx];
+                let source = removed_operators[idx];
 
                 Some((shard_id, source))
             })
             .pipe(stream::iter)
-            .for_each_concurrent(Some(self.manager.config.concurrency()), |(shard_id, source)| {
-                retry(move || {
-                    self.transfer_shard(shard_id, source, keyspace_version)
-                        .map_err(move |err| tracing::warn!(?err, %shard_id, source = %source.id, "Failed to transfer shard"))
-                })
-            })
+            .for_each_concurrent(
+                Some(self.manager.config.concurrency()),
+                |(shard_id, source)| {
+                    retry(move || {
+                        self.transfer_shard(shard_id, source, keyspace_version)
+                            .map_err(move |err| log_err(err, shard_id, source.id))
+                    })
+                },
+            )
             .await;
 
         Ok(State::CompletingMigration(migration_id))
@@ -406,3 +409,19 @@ where
     // emitted.
     backoff::future::retry(backoff, f).await.unwrap()
 }
+
+// Returns `NodeOperator`s which are in `a`, but not in `b`.
+fn replica_set_difference<'a, T>(
+    a: &ReplicaSet<&'a NodeOperator<T>>,
+    b: &ReplicaSet<&NodeOperator<T>>,
+) -> SmallVec<[&'a NodeOperator<T>; 5]> {
+    let mut diff = SmallVec::new();
+
+    for operator in a {
+        if !b.iter().any(|op| op.id == operator.id) {
+            diff.push(*operator);
+        }
+    }
+
+    diff
+}