feat(migration): configurable concurrency (#267)

xDarksome · web-flow · commit 9dc2d9c71ad1 · 2025-10-09T08:29:47.000Z
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-251007.0
+251008.0
diff --git a/crates/cluster/Cargo.toml b/crates/cluster/Cargo.toml
@@ -19,6 +19,7 @@ sharding = { path = "../sharding" }
 derive_more = { workspace = true, features = ["as_ref", "display", "from", "into", "try_from"] }
 derive-where = { workspace = true }
 thiserror = { workspace = true }
+strum = { workspace = true , features = ["derive"] }
 
 # General purpose
 base64.workspace = true
diff --git a/crates/cluster/src/settings.rs b/crates/cluster/src/settings.rs
@@ -4,6 +4,7 @@ use {
     crate::smart_contract,
     serde::{Deserialize, Serialize},
     std::{ops::RangeInclusive, time::Duration},
+    strum::{EnumDiscriminants, FromRepr, IntoDiscriminant, IntoStaticStr},
     time::OffsetDateTime,
 };
 
@@ -41,6 +42,21 @@ pub struct Settings {
     /// success rate of storage operations during a data migration process, due
     /// to `KeyspaceVersionMismatch` errors.
     pub clock_skew: Duration,
+
+    /// Specifies how many shards are allowed to be pulled by a
+    /// [`NodeOperator`] at the same time during data migration.
+    pub migration_concurrency: u16,
+}
+
+impl Default for Settings {
+    fn default() -> Self {
+        Self {
+            max_node_operator_data_bytes: 4096,
+            event_propagation_latency: Duration::from_secs(5),
+            clock_skew: Duration::from_millis(500),
+            migration_concurrency: 10,
+        }
+    }
 }
 
 impl Settings {
@@ -58,6 +74,22 @@ impl Settings {
     ) -> RangeInclusive<OffsetDateTime> {
         time - self.clock_skew..=time + self.clock_skew
     }
+
+    fn extra(&self) -> [Setting; 3] {
+        [
+            Setting::EventPropagationLatency(self.event_propagation_latency),
+            Setting::ClockSkew(self.clock_skew),
+            Setting::MigrationConcurrency(self.migration_concurrency),
+        ]
+    }
+
+    fn apply_extra_setting(&mut self, setting: Setting) {
+        match setting {
+            Setting::EventPropagationLatency(setting) => self.event_propagation_latency = setting,
+            Setting::ClockSkew(setting) => self.clock_skew = setting,
+            Setting::MigrationConcurrency(setting) => self.migration_concurrency = setting,
+        };
+    }
 }
 
 impl TryFrom<Settings> for smart_contract::Settings {
@@ -95,13 +127,13 @@ impl TryFrom<smart_contract::Settings> for Settings {
         let schema_version = bytes[0];
         let bytes = &bytes[1..];
 
-        let extra = match schema_version {
-            0 => postcard::from_bytes::<ExtraV0>(bytes),
+        let mut settings = match schema_version {
+            0 => postcard::from_bytes::<ExtraV0>(bytes).map(Settings::from),
+            1 => postcard::from_bytes::<ExtraV1>(bytes).map(Settings::from),
             ver => return Err(TryFromSmartContractError::UnknownSchemaVersion(ver)),
         }
         .map_err(TryFromSmartContractError::from_postcard)?;
 
-        let mut settings = Settings::from(extra);
         settings.max_node_operator_data_bytes = sc_settings.max_node_operator_data_bytes;
 
         Ok(settings)
@@ -142,7 +174,108 @@ impl From<ExtraV0> for Settings {
                 extra.event_propagation_latency_ms.into(),
             ),
             clock_skew: Duration::from_millis(extra.clock_skew_ms.into()),
+            migration_concurrency: 10,
+        }
+    }
+}
+
+// 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, Debug, Serialize, Deserialize)]
+pub(crate) struct ExtraV1 {
+    // NOTE: Or
+    encoded: Vec<Vec<u8>>,
+}
+
+#[derive(EnumDiscriminants)]
+#[strum_discriminants(name(SettingId))]
+#[strum_discriminants(derive(IntoStaticStr, FromRepr))]
+#[repr(u8)]
+enum Setting {
+    EventPropagationLatency(Duration) = 0,
+    ClockSkew(Duration) = 1,
+    MigrationConcurrency(u16) = 2,
+}
+
+impl SettingId {
+    fn from_idx(idx: usize) -> Option<Self> {
+        Self::from_repr(u8::try_from(idx).ok()?)
+    }
+}
+
+impl Setting {
+    fn encode(&self) -> Vec<u8> {
+        let encode_duration = |value: &Duration| {
+            u32::try_from(value.as_millis())
+                .unwrap_or(u32::MAX)
+                .to_be_bytes()
+                .into()
+        };
+
+        match self {
+            Self::EventPropagationLatency(latency) => encode_duration(latency),
+            Self::ClockSkew(skew) => encode_duration(skew),
+            Self::MigrationConcurrency(concurrency) => concurrency.to_be_bytes().into(),
+        }
+    }
+
+    fn decode(id: SettingId, value: &[u8]) -> Option<Self> {
+        let decode_duration = || {
+            let millis = u32::from_be_bytes(value.try_into().ok()?);
+            Some(Duration::from_millis(millis.into()))
+        };
+
+        let decode_u16 = || Some(u16::from_be_bytes(value.try_into().ok()?));
+
+        Some(match id {
+            SettingId::EventPropagationLatency => Self::EventPropagationLatency(decode_duration()?),
+            SettingId::ClockSkew => Self::ClockSkew(decode_duration()?),
+            SettingId::MigrationConcurrency => Self::MigrationConcurrency(decode_u16()?),
+        })
+    }
+}
+
+impl From<Settings> for ExtraV1 {
+    fn from(settings: Settings) -> Self {
+        let extra_settings = settings.extra();
+
+        let highest_idx = extra_settings
+            .iter()
+            .map(|setting| setting.discriminant() as usize)
+            .max()
+            .unwrap_or_default();
+
+        let mut encoded = vec![vec![]; highest_idx + 1];
+
+        for setting in settings.extra() {
+            encoded[setting.discriminant() as usize] = setting.encode();
         }
+
+        ExtraV1 { encoded }
+    }
+}
+
+impl From<ExtraV1> for Settings {
+    fn from(extra: ExtraV1) -> Self {
+        let mut settings = Settings::default();
+
+        for (idx, value) in extra.encoded.iter().enumerate() {
+            let Some(id) = SettingId::from_idx(idx) else {
+                tracing::warn!("Unknown SettingId({idx}), ignoring");
+                continue;
+            };
+
+            let Some(setting) = Setting::decode(id, value) else {
+                let name: &'static str = id.into();
+                tracing::warn!("Invalid value `{value:?}` for `{name}` setting, ignoring");
+                continue;
+            };
+
+            settings.apply_extra_setting(setting);
+        }
+
+        settings
     }
 }
 
@@ -177,3 +310,26 @@ impl TryFromSmartContractError {
         Self::Codec(format!("{err:?}"))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn extra_v1() {
+        let max_node_operator_data_bytes = 10000;
+
+        let expected_settings = Settings {
+            max_node_operator_data_bytes,
+            event_propagation_latency: Duration::from_secs(42),
+            clock_skew: Duration::from_secs(10),
+            migration_concurrency: 1000,
+        };
+
+        let mut settings: Settings = ExtraV1::from(expected_settings).into();
+        // being set outside of `From` impl
+        settings.max_node_operator_data_bytes = max_node_operator_data_bytes;
+
+        assert_eq!(expected_settings, settings);
+    }
+}
diff --git a/crates/migration/src/manager/mod.rs b/crates/migration/src/manager/mod.rs
@@ -50,10 +50,6 @@ pub trait Config:
         &self,
         node: &'a Self::Node,
     ) -> &'a Self::OutboundReplicaConnection;
-
-    /// Specifies how many shards migration [`Manager`] is allowed to transfer
-    /// at the same time.
-    fn concurrency(&self) -> usize;
 }
 
 /// WCN Migration Manager.
@@ -249,7 +245,7 @@ where
             })
             .pipe(stream::iter)
             .for_each_concurrent(
-                Some(self.manager.config.concurrency()),
+                Some(cluster_view.settings().migration_concurrency.into()),
                 |(shard_id, source)| {
                     retry(move || {
                         self.transfer_shard(shard_id, source, keyspace_version)
diff --git a/crates/migration/src/manager/test.rs b/crates/migration/src/manager/test.rs
@@ -57,10 +57,6 @@ impl super::Config for Config {
     ) -> &'a Self::OutboundReplicaConnection {
         &node.storage
     }
-
-    fn concurrency(&self) -> usize {
-        100
-    }
 }
 
 #[derive(AsRef, Clone)]
@@ -98,9 +94,10 @@ async fn transfers_data_and_writes_to_smart_contract() {
             max_node_operator_data_bytes: 1024,
             event_propagation_latency: Duration::from_secs(1),
             clock_skew: Duration::from_millis(100),
+            migration_concurrency: 100,
         },
         (0..8)
-            .map(|idx| cluster::testing::node_operator(idx as u8))
+            .map(|idx: u8| cluster::testing::node_operator(idx))
             .collect(),
     )
     .await
diff --git a/crates/node/src/lib.rs b/crates/node/src/lib.rs
@@ -284,10 +284,6 @@ impl wcn_migration::manager::Config for AppConfig {
     ) -> &'a Self::OutboundReplicaConnection {
         &node.replica_low_prio_connection
     }
-
-    fn concurrency(&self) -> usize {
-        100
-    }
 }
 
 impl wcn_metrics::aggregator::Config for AppConfig {
diff --git a/crates/replication/src/coordinator/test.rs b/crates/replication/src/coordinator/test.rs
@@ -90,10 +90,11 @@ impl Context {
                 max_node_operator_data_bytes: 1024,
                 event_propagation_latency: Duration::from_secs(3),
                 clock_skew: Duration::from_secs(1),
+                migration_concurrency: 100,
             },
             (0..8)
-                .map(|idx| {
-                    wcn_cluster::testing::node_operator(idx as u8).tap_mut(|operator| {
+                .map(|idx: u8| {
+                    wcn_cluster::testing::node_operator(idx).tap_mut(|operator| {
                         if idx == 0 {
                             operator.clients = vec![Client {
                                 peer_id: client_peer_id,
diff --git a/crates/replication/src/replica/test.rs b/crates/replication/src/replica/test.rs
@@ -71,9 +71,10 @@ impl Context {
                 max_node_operator_data_bytes: 1024,
                 event_propagation_latency: Duration::from_secs(4),
                 clock_skew: Duration::from_secs(2),
+                migration_concurrency: 100,
             },
             (0..8)
-                .map(|idx| wcn_cluster::testing::node_operator(idx as u8))
+                .map(|idx: u8| wcn_cluster::testing::node_operator(idx))
                 .collect(),
         )
         .await
diff --git a/crates/testing/src/lib.rs b/crates/testing/src/lib.rs
@@ -92,6 +92,7 @@ impl TestCluster {
             max_node_operator_data_bytes: 4096,
             event_propagation_latency: Duration::from_secs(1),
             clock_skew: Duration::from_millis(100),
+            migration_concurrency: 100,
         };
 
         tracing::info!(port = %anvil.port(), "Anvil launched");
diff --git a/docs/node-operator-onboarding.md b/docs/node-operator-onboarding.md
@@ -40,7 +40,7 @@ You can have a separate keypair per Node/Database, or you can use a single keypa
 To generate a keypair pull the repository and run the following:
 
 ```bash
- cargo run -p wcn_cluster -F cli key generate
+ cargo run -p wcn_operator -- key generate
 ```
 
 It will show you the following (**DO NOT use these example values for your deployments**)

Original file line number	Diff line number	Diff line change
`@@ -284,10 +284,6 @@ impl wcn_migration::manager::Config for AppConfig {`
`284`	`284`	`) -> &'a Self::OutboundReplicaConnection {`
`285`	`285`	`&node.replica_low_prio_connection`
`286`	`286`	`}`
`287`		`-`
`288`		`- fn concurrency(&self) -> usize {`
`289`		`- 100`
`290`		`- }`
`291`	`287`	`}`
`292`	`288`
`293`	`289`	`impl wcn_metrics::aggregator::Config for AppConfig {`