feat: migration 2.0 (#232)

Author: xDarksome

Cargo.lock

@@ -55,11 +55,14 @@ tokio-serde-postcard = { git = "https://github.com/xDarksome/tokio-serde-postcar
 postcard = { version = "1.0", default-features = false }
 itertools = "0.12"
 futures = "0.3"
+futures-concurrency = "7.6"
 backoff = { version = "0.4", features = ["tokio"] }
 tracing = "0.1"
 tokio-stream = "0.1"
 strum = "0.27"
 xxhash-rust = { version = "0.8", features = ["xxh3", "const_xxh3"] }
+thiserror = "1"
+anyhow = "1"
 
 [workspace.lints.clippy]
 all = { level = "deny", priority = -1 }

Cargo.toml

@@ -1 +1 @@
-250729.0
+250804.0

VERSION

@@ -6,8 +6,7 @@ use {
     derivative::Derivative,
     derive_more::TryFrom,
     serde::{Deserialize, Serialize},
-    sharding::ShardId,
-    std::collections::HashSet,
+    std::{collections::HashSet, ops::RangeInclusive},
     xxhash_rust::xxh3::Xxh3Builder,
 };
 
@@ -38,10 +37,18 @@ pub struct Keyspace<S = ()> {
 #[derive(Clone)]
 pub struct Shards(sharding::Keyspace<node_operator::Idx, { REPLICATION_FACTOR as usize }>);
 
+/// ID of a [`Shard`].
+pub type ShardId = u16;
+
+/// Returns the keyrange the provided shard is resposible for.
+pub fn keyrange(id: ShardId) -> RangeInclusive<u64> {
+    sharding::ShardId(id).key_range()
+}
+
 /// A single [`Shard`] within a [`Keyspace`].
 #[derive(Clone, Copy, Debug)]
-pub struct Shard {
-    replica_set: [node_operator::Idx; REPLICATION_FACTOR as usize],
+pub struct Shard<T = node_operator::Idx> {
+    pub(crate) replica_set: [T; REPLICATION_FACTOR as usize],
 }
 
 /// Strategy of distributing [`Shard`]s to [`node_operator`]s.
@@ -95,9 +102,20 @@ impl Keyspace<Shards> {
     /// Returns the [`Shard`] that contains the specified key.
     pub fn shard(&self, key: u64) -> Shard {
         Shard {
-            replica_set: *self.shards.0.shard_replicas(ShardId::from_key(key)),
+            replica_set: *self
+                .shards
+                .0
+                .shard_replicas(sharding::ShardId::from_key(key)),
         }
     }
+
+    /// Returns all [`Shard`]s of this [`Keyspace`].
+    pub fn shards(&self) -> impl Iterator<Item = (ShardId, Shard)> + '_ {
+        self.shards
+            .0
+            .shards()
+            .map(|(id, &replica_set)| (id.0, Shard { replica_set }))
+    }
 }
 
 impl<S> Keyspace<S> {
@@ -144,10 +162,10 @@ impl<S> Keyspace<S> {
     }
 }
 
-impl Shard {
+impl<T: Copy> Shard<T> {
     /// Returns [`ReplicaSet`] assigned to this [`Shard`].
-    pub fn replica_set(&self) -> ReplicaSet {
-        self.replica_set
+    pub fn replica_set(&self) -> &ReplicaSet<T> {
+        &self.replica_set
     }
 }
 

crates/cluster/src/keyspace.rs

@@ -74,7 +74,7 @@ pub trait Config: Send + Sync + 'static {
     /// data/logic into it.
     ///
     /// If no additional logic is required - just specify [`Node`].
-    type Node: Clone + Send + Sync + 'static;
+    type Node: AsRef<PeerId> + Clone + Send + Sync + 'static;
 
     /// Creates a new [`Config::Node`].
     fn new_node(&self, operator_id: node_operator::Id, node: Node) -> Self::Node;
@@ -96,7 +96,9 @@ struct Inner<C: Config> {
     config: C,
     smart_contract: C::SmartContract,
     view: ArcSwap<View<C>>,
-    watch: watch::Receiver<()>,
+
+    watch_tx: watch::Sender<()>,
+    watch_rx: watch::Receiver<()>,
 }
 
 /// Version of a WCN [`Cluster`].
@@ -187,7 +189,8 @@ where
             config: cfg,
             smart_contract: contract,
             view: ArcSwap::new(Arc::new(view)),
-            watch: rx,
+            watch_tx: tx.clone(),
+            watch_rx: rx,
         });
 
         let guard = Task {
@@ -225,13 +228,55 @@ impl<C: Config> Cluster<C> {
     pub fn updates(&self) -> impl Stream<Item = ()> + Send + 'static {
         // TODO: periodically check with the SC to prevent drift
 
-        tokio_stream::wrappers::WatchStream::new(self.inner.watch.clone())
+        tokio_stream::wrappers::WatchStream::new(self.inner.watch_rx.clone())
+    }
+
+    /// Creates a new [`Watch`] of this [`Cluster`].
+    pub fn watch(&self) -> Watch {
+        let mut rx = self.inner.watch_rx.clone();
+        rx.mark_changed();
+
+        Watch {
+            inner: rx,
+            _tx: self.inner.watch_tx.clone(),
+        }
     }
 
     /// Returns reference to the underlying [`SmartContract`].
     pub fn smart_contract(&self) -> &C::SmartContract {
         &self.inner.smart_contract
     }
+
+    /// Indicates whether this [`Cluster`] contains a [`Node`] with the provided
+    /// ID.
+    pub fn contains_node(&self, peer_id: &PeerId) -> bool
+    where
+        C::Node: AsRef<PeerId>,
+    {
+        self.using_view(|view| view.node_operators().contains_node(peer_id))
+    }
+
+    /// Checks whether the provided [`keyspace`] version is compatible with
+    /// current state of the [`Cluster`].
+    pub fn validate_keyspace_version(&self, version: u64) -> bool {
+        self.using_view(|view| view.validate_keyspace_version(version))
+    }
+}
+
+/// Handle to track [`Cluster`] updates.
+pub struct Watch {
+    inner: watch::Receiver<()>,
+    _tx: watch::Sender<()>,
+}
+
+impl Watch {
+    /// Resolves when the next [`Cluster`] update occurs.
+    ///
+    /// First time after creation resolves immediately.
+    pub async fn cluster_updated(&mut self) {
+        // NOTE(unwrap): we are holding the `Sender`, so it won't ever error.
+        self.inner.changed().await.unwrap()
+    }
 }
 
 impl<C: Config> Cluster<C>

crates/cluster/src/lib.rs

@@ -1,6 +1,7 @@
 //! Node within a WCN cluster.
 
 use {
+    derive_more::derive::AsRef,
     libp2p::identity::PeerId,
     serde::{Deserialize, Serialize},
     std::net::SocketAddrV4,
@@ -11,12 +12,13 @@ use {
 /// The IP address is currently being encrypted using a format-preserving
 /// encryption algorithm.
 // TODO: encrypt
-#[derive(Debug, Clone, Copy)]
+#[derive(AsRef, Debug, Clone, Copy)]
 pub struct Node {
     /// [`PeerId`] of the [`Node`].
     ///
     /// Used for authentication. Multiple nodes managed by the same
     /// node operator are allowed to have the same [`PeerId`].
+    #[as_ref]
     pub peer_id: PeerId,
 
     /// [`SocketAddrV4`] of the [`Node`].

crates/cluster/src/node.rs

@@ -115,12 +115,31 @@ impl<N> NodeOperator<N> {
         &self.nodes[n % self.nodes.len()]
     }
 
-    /// List of [`Node`]s of the [`NodeOperator`].
+    /// Returns [`Node`]s of this [`NodeOperator`].
     ///
     /// [`NodeOperator`] is guaranteed to always have at least 2 nodes.
     pub fn nodes(&self) -> &[N] {
         &self.nodes
     }
+
+    /// Returns an [`Iterator`] of [`Node`]s of this [`NodeOperator`].
+    ///
+    /// Iterates over all [`Node`]s starting from an arbitrary position.
+    /// Intended for load balancing purposes.
+    ///
+    /// [`NodeOperator`] is guaranteed to always have at least 2 nodes.
+    pub fn nodes_lb_iter(&self) -> impl Iterator<Item = &N> {
+        // TODO: this is suboptimal for > 2 nodes, because in case of a node failure the
+        // next node in the list will receive all of the load of it's neighbor.
+        //
+        // Naive solution to this is to reallocate and shuffle. A better solution may be
+        // to split the `Vec` on chunks of size <=N and then iterate over the chunks and
+        // shuffle the chunks in-place.
+
+        let n = self.counter.fetch_add(1, atomic::Ordering::Relaxed);
+        let (left, right) = self.nodes.split_at(n % self.nodes.len());
+        right.iter().chain(left.iter())
+    }
 }
 
 /// [`NodeOperator`] with serialized [`Data`].

crates/cluster/src/node_operator.rs

@@ -3,6 +3,7 @@
 use {
     crate::{self as cluster, node_operator, Node, NodeOperator},
     indexmap::IndexMap,
+    libp2p::PeerId,
     std::sync::{
         atomic::{self, AtomicUsize},
         Arc,
@@ -59,6 +60,20 @@ impl<N> NodeOperators<N> {
         })
     }
 
+    /// Indicates whether any of the [`NodeOperators`] contains a node with the
+    /// provided ID.
+    pub fn contains_node(&self, peer_id: &PeerId) -> bool
+    where
+        N: AsRef<PeerId>,
+    {
+        // TODO: Consider optimizing by building a lookup table.
+        self.slots.iter().any(|opt| {
+            opt.as_ref()
+                .map(|op| op.nodes().iter().any(|node| node.as_ref() == peer_id))
+                .unwrap_or_default()
+        })
+    }
+
     /// Returns a [`NodeOperator`] responsible for the next request.
     ///
     /// [`NodeOperator`]s are being iterated in round-robin fashion for

crates/cluster/src/node_operators.rs

@@ -96,7 +96,8 @@ pub trait Write {
     /// [`migration::Completed`] event MUST be emitted.
     /// Otherwise the data pull MUST be marked as completed for the
     /// [`node::Operator`] and [`migration::DataPullCompleted`] MUST be emitted.
-    fn complete_migration(&self, id: migration::Id) -> impl Future<Output = WriteResult<()>>;
+    fn complete_migration(&self, id: migration::Id)
+        -> impl Future<Output = WriteResult<()>> + Send;
 
     /// Aborts the ongoing data [`migration`] process restoring the WCN cluster
     /// to the original state it had before the migration had started.

crates/cluster/src/smart_contract/mod.rs

@@ -48,7 +48,7 @@ impl<C: Config<KeyspaceShards = keyspace::Shards>> View<C> {
         // therefore guaranteeing that every `NodeOperatorIdx` is valid.
         self.keyspace
             .shard(key)
-            .replica_set()
+            .replica_set
             .map(|idx| self.node_operators.get_by_idx(idx).unwrap())
     }
 
@@ -61,10 +61,54 @@ impl<C: Config<KeyspaceShards = keyspace::Shards>> View<C> {
         self.migration()?
             .keyspace()
             .shard(key)
-            .replica_set()
+            .replica_set
             .map(|idx| self.node_operators.get_by_idx(idx).unwrap())
             .pipe(Some)
     }
+
+    /// Returns [`keyspace::Shard`]s of the primary [`Keyspace`].
+    pub fn primary_keyspace_shards(
+        &self,
+    ) -> impl Iterator<Item = (keyspace::ShardId, keyspace::Shard<&NodeOperator<C::Node>>)> {
+        // NOTE(unwrap): we use `Keyspace::validate` every time it enters the system,
+        // therefore guaranteeing that every `NodeOperatorIdx` is valid.
+        self.keyspace.shards().map(|(id, shard)| {
+            (id, keyspace::Shard {
+                replica_set: shard
+                    .replica_set
+                    .map(|idx| self.node_operators.get_by_idx(idx).unwrap()),
+            })
+        })
+    }
+
+    /// Returns [`keyspace::Shard`]s of the secondary [`Keyspace`].
+    pub fn secondary_keyspace_shards(
+        &self,
+    ) -> Option<impl Iterator<Item = (keyspace::ShardId, keyspace::Shard<&NodeOperator<C::Node>>)>>
+    {
+        // NOTE(unwrap): we use `Keyspace::validate` every time it enters the system,
+        // therefore guaranteeing that every `NodeOperatorIdx` is valid.
+        self.migration()?
+            .keyspace()
+            .shards()
+            .map(|(id, shard)| {
+                (id, keyspace::Shard {
+                    replica_set: shard
+                        .replica_set
+                        .map(|idx| self.node_operators.get_by_idx(idx).unwrap()),
+                })
+            })
+            .pipe(Some)
+    }
+
+    /// Indicates whether the specified [`node_operator`] is still in process of
+    /// pulling data as part of a data [`migration`] process.
+    pub fn is_pulling(&self, operator_id: &node_operator::Id) -> bool {
+        self.node_operators()
+            .get_idx(operator_id)
+            .and_then(|operator_idx| self.migration().map(|mig| mig.is_pulling(operator_idx)))
+            .unwrap_or_default()
+    }
 }
 
 impl<C: Config> View<C> {
@@ -98,6 +142,16 @@ impl<C: Config> View<C> {
         &self.node_operators
     }
 
+    /// Checks whether the provided [`keyspace`] version is compatible with
+    /// current state of the cluster.
+    pub fn validate_keyspace_version(&self, version: u64) -> bool {
+        if let Some(migration) = self.migration() {
+            migration.keyspace().version() == version
+        } else {
+            self.keyspace().version() == version
+        }
+    }
+
     pub(super) fn require_no_migration(&self) -> Result<(), migration::InProgressError> {
         if let Some(migration) = self.migration() {
             return Err(migration::InProgressError(migration.id()));