Merge pull request fedimint#4719 from m1sterc001guy/migration_bytes

test: Allow Client DB Migration tests to write arbitrary bytes

db/migrations/dummy-client-v0/IDENTITY

@@ -1 +1 @@
-0c1eec31-86a3-492c-adbd-8d32f6eb6699
+3a603dbb-c6b2-49f0-9084-ea906b18aa80

db/migrations/lightning-client-v0/IDENTITY

@@ -1 +1 @@
-d044f224-f20c-40d6-97fa-967a7994df0b
+c4445175-0471-435b-afbb-c2c5aa644f01

fedimint-client/src/db.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 
 use fedimint_core::api::ApiVersionSet;
 use fedimint_core::config::{ClientConfig, FederationId};
-use fedimint_core::core::{ModuleInstanceId, OperationId};
+use fedimint_core::core::{IntoDynInstance, ModuleInstanceId, OperationId};
 use fedimint_core::db::{
     migrate_database_version, Database, DatabaseTransaction, DatabaseValue, DatabaseVersion,
     DatabaseVersionKey, IDatabaseTransactionOpsCoreTyped,
@@ -24,7 +24,7 @@ use crate::sm::executor::{
     ActiveStateKey, ActiveStateKeyBytes, ActiveStateKeyPrefixBytes, InactiveStateKey,
     InactiveStateKeyBytes, InactiveStateKeyPrefixBytes,
 };
-use crate::sm::{ActiveStateMeta, DynState, InactiveStateMeta};
+use crate::sm::{ActiveStateMeta, DynState, InactiveStateMeta, State};
 
 #[repr(u8)]
 #[derive(Clone, EnumIter, Debug)]
@@ -557,3 +557,58 @@ pub async fn remove_old_and_persist_new_inactive_states(
         dbtx.insert_new_entry(&state, &inactive_state).await;
     }
 }
+
+/// Migrates a particular state by looping over all active and inactive states.
+/// If the `migrate` closure returns `None`, this state was not migrated and
+/// should be added to the new state machine vectors.
+pub async fn migrate_state<S: State + IntoDynInstance<DynType = DynState>>(
+    module_instance_id: ModuleInstanceId,
+    active_states: Vec<(Vec<u8>, OperationId)>,
+    inactive_states: Vec<(Vec<u8>, OperationId)>,
+    decoders: ModuleDecoderRegistry,
+    migrate: fn(
+        &[u8],
+        ModuleInstanceId,
+        &ModuleDecoderRegistry,
+    ) -> anyhow::Result<Option<DynState>>,
+) -> anyhow::Result<Option<(Vec<DynState>, Vec<DynState>)>> {
+    let mut new_active_states = Vec::with_capacity(active_states.len());
+    for (active_state, _) in active_states {
+        let bytes = active_state.as_slice();
+        let state = match migrate(bytes, module_instance_id, &decoders)? {
+            Some(state) => state,
+            None => {
+                // Try to decode the bytes as a `DynState`
+                let dynstate = DynState::from_bytes(bytes, &decoders)?;
+                let state_machine = dynstate
+                    .as_any()
+                    .downcast_ref::<S>()
+                    .expect("Unexpected DynState supplied to migration function");
+                state_machine.clone().into_dyn(module_instance_id)
+            }
+        };
+
+        new_active_states.push(state);
+    }
+
+    let mut new_inactive_states = Vec::with_capacity(inactive_states.len());
+    for (inactive_state, _) in inactive_states {
+        let bytes = inactive_state.as_slice();
+        let state = match migrate(bytes, module_instance_id, &decoders)? {
+            Some(state) => state,
+            None => {
+                // Try to decode the bytes as a `DynState`
+                let dynstate = DynState::from_bytes(bytes, &decoders)?;
+                let state_machine = dynstate
+                    .as_any()
+                    .downcast_ref::<S>()
+                    .expect("Unexpected DynState supplied to migration function");
+                state_machine.clone().into_dyn(module_instance_id)
+            }
+        };
+
+        new_inactive_states.push(state);
+    }
+
+    Ok(Some((new_active_states, new_inactive_states)))
+}

fedimint-client/src/sm/mod.rs

@@ -10,8 +10,8 @@ mod notifier;
 
 pub use dbtx::ClientSMDatabaseTransaction;
 pub use executor::{
-    ActiveStateKeyPrefix, ActiveStateMeta, Executor, ExecutorBuilder, InactiveStateKeyPrefix,
-    InactiveStateMeta,
+    ActiveStateKeyBytes, ActiveStateKeyPrefix, ActiveStateMeta, Executor, ExecutorBuilder,
+    InactiveStateKeyBytes, InactiveStateKeyPrefix, InactiveStateMeta,
 };
 pub use notifier::{ModuleNotifier, Notifier, NotifierSender};
 pub use state::{Context, DynContext, DynState, IState, OperationState, State, StateTransition};

fedimint-testing/src/db.rs

@@ -5,14 +5,14 @@ use std::path::{Path, PathBuf};
 use std::{env, fs, io};
 
 use anyhow::{bail, format_err, Context};
-use fedimint_client::db::{
-    apply_migrations_client, remove_old_and_persist_new_active_states,
-    remove_old_and_persist_new_inactive_states,
-};
+use fedimint_client::db::apply_migrations_client;
 use fedimint_client::module::init::DynClientModuleInit;
 use fedimint_client::module::ClientModule;
-use fedimint_client::sm::{ActiveStateKeyPrefix, InactiveStateKeyPrefix};
-use fedimint_core::core::IntoDynInstance;
+use fedimint_client::sm::{
+    ActiveStateKeyBytes, ActiveStateKeyPrefix, ActiveStateMeta, InactiveStateKeyBytes,
+    InactiveStateKeyPrefix, InactiveStateMeta,
+};
+use fedimint_core::core::OperationId;
 use fedimint_core::db::{
     apply_migrations, apply_migrations_server, Database, DatabaseVersion,
     IDatabaseTransactionOpsCoreTyped, ServerMigrationFn,
@@ -165,16 +165,15 @@ where
 /// namespace. `state_machine_prepare` creates client state machine data that
 /// can be used for testing state machine migrations. This is created in the
 /// global namespace.
-pub async fn snapshot_db_migrations_client<'a, F, S, I, T>(
+pub async fn snapshot_db_migrations_client<'a, F, S, I>(
     snapshot_name: &str,
     data_prepare: F,
     state_machine_prepare: S,
 ) -> anyhow::Result<()>
 where
     F: Fn(Database) -> BoxFuture<'a, ()> + Send + Sync,
-    S: Fn() -> (Vec<T::States>, Vec<T::States>) + Send + Sync,
+    S: Fn() -> (Vec<Vec<u8>>, Vec<Vec<u8>>) + Send + Sync,
     I: CommonModuleInit,
-    T: ClientModule,
 {
     let project_root = get_project_root().unwrap();
     let snapshot_dir = project_root.join("db/migrations").join(snapshot_name);
@@ -188,30 +187,39 @@ where
             data_prepare(isolated_db).await;
 
             let (active_states, inactive_states) = state_machine_prepare();
-            let new_active_states = active_states
-                .into_iter()
-                .map(|state| state.into_dyn(TEST_MODULE_INSTANCE_ID))
-                .collect::<Vec<_>>();
-            let new_inactive_states = inactive_states
-                .into_iter()
-                .map(|state| state.into_dyn(TEST_MODULE_INSTANCE_ID))
-                .collect::<Vec<_>>();
-
             let mut global_dbtx = db.begin_transaction().await;
-            remove_old_and_persist_new_active_states(
-                &mut global_dbtx.to_ref_nc(),
-                new_active_states,
-                Vec::new(),
-                TEST_MODULE_INSTANCE_ID,
-            )
-            .await;
-            remove_old_and_persist_new_inactive_states(
-                &mut global_dbtx.to_ref_nc(),
-                new_inactive_states,
-                Vec::new(),
-                TEST_MODULE_INSTANCE_ID,
-            )
-            .await;
+
+            for state in active_states {
+                global_dbtx
+                    .insert_new_entry(
+                        &ActiveStateKeyBytes {
+                            operation_id: OperationId::new_random(),
+                            module_instance_id: TEST_MODULE_INSTANCE_ID,
+                            state,
+                        },
+                        &ActiveStateMeta {
+                            created_at: fedimint_core::time::now(),
+                        },
+                    )
+                    .await;
+            }
+
+            for state in inactive_states {
+                global_dbtx
+                    .insert_new_entry(
+                        &InactiveStateKeyBytes {
+                            operation_id: OperationId::new_random(),
+                            module_instance_id: TEST_MODULE_INSTANCE_ID,
+                            state,
+                        },
+                        &InactiveStateMeta {
+                            created_at: fedimint_core::time::now(),
+                            exited_at: fedimint_core::time::now(),
+                        },
+                    )
+                    .await;
+            }
+
             global_dbtx.commit_tx().await;
         }
         .boxed()

modules/fedimint-dummy-tests/tests/tests.rs

@@ -148,6 +148,7 @@ mod fedimint_migration_tests {
     use fedimint_core::db::{
         Database, DatabaseVersion, DatabaseVersionKeyV0, IDatabaseTransactionOpsCoreTyped,
     };
+    use fedimint_core::encoding::Encodable;
     use fedimint_core::module::DynServerModuleInit;
     use fedimint_core::{Amount, BitcoinHash, OutPoint, TransactionId};
     use fedimint_dummy_client::db::{
@@ -163,7 +164,7 @@ mod fedimint_migration_tests {
     use fedimint_logging::TracingSetup;
     use fedimint_testing::db::{
         snapshot_db_migrations, snapshot_db_migrations_client, validate_migrations_client,
-        validate_migrations_server, BYTE_32,
+        validate_migrations_server, BYTE_32, TEST_MODULE_INSTANCE_ID,
     };
     use futures::StreamExt;
     use rand::rngs::OsRng;
@@ -207,19 +208,66 @@ mod fedimint_migration_tests {
         dbtx.commit_tx().await;
     }
 
-    fn create_client_states() -> (Vec<DummyStateMachine>, Vec<DummyStateMachine>) {
+    fn create_client_states() -> (Vec<Vec<u8>>, Vec<Vec<u8>>) {
         // Create an active state and inactive state that will not be migrated.
         let input_operation_id = OperationId::new_random();
-        let txid = TransactionId::from_slice(&BYTE_32).unwrap();
-        let input_state =
-            DummyStateMachine::Input(Amount::from_sats(1000), txid, input_operation_id);
+        let input_state = {
+            let mut input_state = Vec::<u8>::new();
+            Amount::from_sats(1000)
+                .consensus_encode(&mut input_state)
+                .expect("Amount is encodable");
+            TransactionId::from_slice(&BYTE_32)
+                .expect("Couldn't create TransactionId")
+                .consensus_encode(&mut input_state)
+                .expect("TransactionId is encodable");
+            input_operation_id
+                .consensus_encode(&mut input_state)
+                .expect("OperationId is encodable");
+            input_state
+        };
+
+        let input_variant = {
+            let mut input_variant = Vec::<u8>::new();
+            TEST_MODULE_INSTANCE_ID
+                .consensus_encode(&mut input_variant)
+                .expect("u16 is encodable");
+            0u64.consensus_encode(&mut input_variant)
+                .expect("u64 is encodable"); // Input variant
+            input_state
+                .consensus_encode(&mut input_variant)
+                .expect("input state is encodable");
+            input_variant
+        };
 
         // Create and active state and inactive state that will be migrated.
-        let operation_id = OperationId::new_random();
-        let state = DummyStateMachine::Unreachable(operation_id, Amount::from_sats(1000));
+        let unreachable_operation_id = OperationId::new_random();
+        let unreachable_state = {
+            let mut unreachable = Vec::<u8>::new();
+            unreachable_operation_id
+                .consensus_encode(&mut unreachable)
+                .expect("OperationId is encodable");
+            Amount::from_sats(1000)
+                .consensus_encode(&mut unreachable)
+                .expect("Amount is encodable");
+            unreachable
+        };
+
+        let unreachable_variant = {
+            let mut unreachable = Vec::<u8>::new();
+            TEST_MODULE_INSTANCE_ID
+                .consensus_encode(&mut unreachable)
+                .expect("u16 is encodable");
+            5u64.consensus_encode(&mut unreachable)
+                .expect("u64 is encodable"); // Unreachable variant
+            unreachable_state
+                .consensus_encode(&mut unreachable)
+                .expect("unreachable state is encodable");
+            unreachable
+        };
+
         (
-            vec![state.clone(), input_state.clone()],
-            vec![state, input_state],
+            vec![input_variant.clone(), unreachable_variant.clone()],
+            vec![input_variant, unreachable_variant],
         )
     }
 
@@ -278,7 +326,7 @@ mod fedimint_migration_tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn snapshot_client_db_migrations() -> anyhow::Result<()> {
-        snapshot_db_migrations_client::<_, _, DummyCommonInit, DummyClientModule>(
+        snapshot_db_migrations_client::<_, _, DummyCommonInit>(
             "dummy-client-v0",
             |db| Box::pin(async move { create_client_db_with_v0_data(db).await }),
             create_client_states,