feat(backend): User data migration

src/backend/src/lib.rs

@@ -17,6 +17,7 @@ use ic_cdk::api::management_canister::ecdsa::{
     SignWithEcdsaArgument,
 };
 use ic_cdk::api::time;
+use ic_cdk::eprintln;
 use ic_cdk_macros::{export_candid, init, post_upgrade, query, update};
 use ic_cdk_timers::{clear_timer, set_timer_interval};
 use ic_stable_structures::{
@@ -625,10 +626,22 @@ fn migration_stop_timer() -> Result<(), String> {
     })
 }
 
-/// Steps the migration
+/// Steps the migration.
+///
+/// On error, the migration is marked as failed and the timer is cleared.
 #[update(guard = "caller_is_allowed")]
 async fn step_migration() {
-    migrate::step_migration().await;
+    let result = migrate::step_migration().await;
+    eprintln!("Stepped migration: {:?}", result);
+    if let Err(err) = result {
+        mutate_state(|s| {
+            if let Some(migration) = &mut s.migration {
+                migration.progress = MigrationProgress::Failed(err);
+                clear_timer(migration.timer_id);
+            }
+            eprintln!("Migration failed: {err:?}");
+        });
+    };
 }
 
 /// Computes the parity bit allowing to recover the public key from the signature.

src/backend/src/migrate.rs

@@ -2,13 +2,23 @@ use crate::{
     mutate_state, read_state,
     types::{Candid, StoredPrincipal},
 };
-use candid::{decode_one, CandidType, Principal};
+use candid::{decode_one, encode_one, CandidType, Principal};
+use ic_cdk::eprintln;
 use ic_cdk_timers::clear_timer;
 use serde::Deserialize;
-use shared::types::{
-    custom_token::CustomToken, token::UserToken, user_profile::StoredUserProfile,
-    MigrationProgress, Timestamp,
+use shared::{
+    backend_api::Service,
+    types::{
+        custom_token::CustomToken, token::UserToken, user_profile::StoredUserProfile,
+        MigrationError, MigrationProgress, Timestamp,
+    },
 };
+use std::ops::Bound;
+use steps::{
+    assert_target_empty, assert_target_has_all_data, lock_migration_target, make_this_readonly,
+    unlock_local, unlock_target,
+};
+pub mod steps;
 
 /// A chunk of data to be migrated.
 ///
@@ -67,8 +77,110 @@ pub fn bulk_up(data: &[u8]) {
     }
 }
 
-#[allow(clippy::unused_async)] // TODO: remove this once we make real function calls
-pub async fn step_migration() {
+/// The next chunk of user tokens to be migrated.
+fn next_user_token_chunk(last_user_token: Option<Principal>) -> Vec<(Principal, Vec<UserToken>)> {
+    let chunk_size = 5;
+    let range = last_user_token.map_or((Bound::Unbounded, Bound::Unbounded), |token| {
+        (Bound::Excluded(StoredPrincipal(token)), Bound::Unbounded)
+    });
+    read_state(|state| {
+        state
+            .user_token
+            .range(range)
+            .take(chunk_size)
+            .map(|(stored_principal, token)| (stored_principal.0, token.0))
+            .collect::<Vec<_>>()
+    })
+}
+
+/// The next chunk of custom tokens to be migrated.
+fn next_custom_token_chunk(
+    last_custom_token: Option<Principal>,
+) -> Vec<(Principal, Vec<CustomToken>)> {
+    let chunk_size = 5;
+    let range = last_custom_token.map_or((Bound::Unbounded, Bound::Unbounded), |token| {
+        (Bound::Excluded(StoredPrincipal(token)), Bound::Unbounded)
+    });
+    read_state(|state| {
+        state
+            .custom_token
+            .range(range)
+            .take(chunk_size)
+            .map(|(stored_principal, token)| (stored_principal.0, token.0))
+            .collect::<Vec<_>>()
+    })
+}
+
+/// The next chunk of user profiles to be migrated.
+fn next_user_profile_chunk(
+    last_user_profile: Option<(Timestamp, Principal)>,
+) -> Vec<((Timestamp, Principal), StoredUserProfile)> {
+    let chunk_size = 5;
+    let range = last_user_profile.map_or(
+        (Bound::Unbounded, Bound::Unbounded),
+        |(timestamp, principal)| {
+            (
+                Bound::Excluded((timestamp, StoredPrincipal(principal))),
+                Bound::Unbounded,
+            )
+        },
+    );
+    read_state(|state| {
+        state
+            .user_profile
+            .range(range)
+            .take(chunk_size)
+            .map(|((timestamp, stored_principal), profile)| {
+                ((timestamp, stored_principal.0), profile.0)
+            })
+            .collect::<Vec<_>>()
+    })
+}
+
+/// The next chunk of user timestamps to be migrated.
+fn next_user_timestamp_chunk(user_maybe: Option<Principal>) -> Vec<(Principal, Timestamp)> {
+    let chunk_size = 5;
+    let range = user_maybe.map_or((Bound::Unbounded, Bound::Unbounded), |user| {
+        (Bound::Excluded(StoredPrincipal(user)), Bound::Unbounded)
+    });
+    read_state(|state| {
+        state
+            .user_profile_updated
+            .range(range)
+            .take(chunk_size)
+            .map(|(stored_principal, timestamp)| (stored_principal.0, timestamp))
+            .collect::<Vec<_>>()
+    })
+}
+
+/// Migrates a chunk of data.
+///
+/// # Returns
+/// The updated progress.
+///
+/// # Errors
+/// - Throws a `MigrationError::DataMigrationFailed` if the data transfer fails.
+macro_rules! migrate {
+    ($migration:ident, $chunk:ident, $progress_variant:ident, $chunk_variant:ident) => {{
+        let last = $chunk.last().map(|(k, _)| k).cloned();
+        let next_state = last
+            .map(|last| MigrationProgress::$progress_variant(Some(last)))
+            .unwrap_or_else(|| $migration.progress.next());
+        let migration_data = MigrationChunk::$chunk_variant($chunk);
+        let migration_bytes = encode_one(migration_data).expect("failed to encode migration data");
+        Service($migration.to)
+            .bulk_up(migration_bytes)
+            .await
+            .map_err(|e| {
+                eprintln!("Failed to transfer data {e:?}");
+                MigrationError::DataMigrationFailed
+            })?;
+        next_state
+    }};
+}
+pub(crate) use migrate;
+
+pub async fn step_migration() -> Result<MigrationProgress, MigrationError> {
     fn set_progress(progress: MigrationProgress) {
         mutate_state(|state| {
             state.migration.iter_mut().for_each(|migration| {
@@ -78,60 +190,60 @@ pub async fn step_migration() {
     }
     let migration = read_state(|s| s.migration.clone());
     let progress = match migration {
-        Some(migration) => {
-            match migration.progress {
-                MigrationProgress::Pending => {
-                    // TODO: Lock the local canister APIs.
-                    migration.progress.next()
-                }
-                MigrationProgress::LockingTarget => {
-                    // TODO: Lock the target canister APIs.
-                    migration.progress.next()
-                }
-                MigrationProgress::CheckingTarget => {
-                    // TODO: Check that the target canister is empty.
-                    migration.progress.next()
-                }
-                MigrationProgress::MigratedUserTokensUpTo(_last) => {
-                    // TODO: Migrate user tokens, then move on to the next stage.
-                    migration.progress.next()
-                }
-                MigrationProgress::MigratedCustomTokensUpTo(_last_custom_token) => {
-                    // TODO: Migrate custom tokens, then move on to the next stage.
-                    migration.progress.next()
-                }
-                MigrationProgress::MigratedUserTimestampsUpTo(_user_maybe) => {
-                    // TODO: Migrate user timestamps, then move on to the next stage.
-                    migration.progress.next()
-                }
-                MigrationProgress::MigratedUserProfilesUpTo(_last_user_profile) => {
-                    // TODO: Migrate user profiles, then move on to the next stage.
-                    migration.progress.next()
-                }
-                MigrationProgress::CheckingDataMigration => {
-                    // TODO: Check that the target canister has all the data.
-                    migration.progress.next()
-                }
-                MigrationProgress::UnlockingTarget => {
-                    // TODO: Unlock user profiles in the target canister
-                    migration.progress.next()
-                }
-                MigrationProgress::Unlocking => {
-                    // TODO: Unlock signing in this canister
-                    migration.progress.next()
-                }
-                // TODO: Add steps to unlock APIs.
-                MigrationProgress::Completed => {
-                    // Migration is complete.
-                    clear_timer(migration.timer_id);
-                    migration.progress.next()
-                }
-                MigrationProgress::Failed(_) => ic_cdk::trap("Migration error."),
+        Some(migration) => match migration.progress {
+            MigrationProgress::Pending => {
+                make_this_readonly();
+                migration.progress.next()
             }
-        }
-        None => {
-            ic_cdk::trap("migration is not in progress");
-        }
+            MigrationProgress::LockingTarget => {
+                lock_migration_target(&migration).await?;
+                migration.progress.next()
+            }
+            MigrationProgress::CheckingTarget => {
+                assert_target_empty(&migration).await?;
+                migration.progress.next()
+            }
+            MigrationProgress::MigratedUserTokensUpTo(last) => {
+                let chunk = next_user_token_chunk(last);
+                migrate!(migration, chunk, MigratedUserTokensUpTo, UserToken)
+            }
+            MigrationProgress::MigratedCustomTokensUpTo(last_custom_token) => {
+                let chunk = next_custom_token_chunk(last_custom_token);
+                migrate!(migration, chunk, MigratedCustomTokensUpTo, CustomToken)
+            }
+            MigrationProgress::MigratedUserTimestampsUpTo(user_maybe) => {
+                let chunk = next_user_timestamp_chunk(user_maybe);
+                migrate!(
+                    migration,
+                    chunk,
+                    MigratedUserTimestampsUpTo,
+                    UserProfileUpdated
+                )
+            }
+            MigrationProgress::MigratedUserProfilesUpTo(last_user_profile) => {
+                let chunk = next_user_profile_chunk(last_user_profile);
+                migrate!(migration, chunk, MigratedUserProfilesUpTo, UserProfile)
+            }
+            MigrationProgress::CheckingDataMigration => {
+                assert_target_has_all_data(&migration).await?;
+                migration.progress.next()
+            }
+            MigrationProgress::UnlockingTarget => {
+                unlock_target(&migration).await?;
+                migration.progress.next()
+            }
+            MigrationProgress::Unlocking => {
+                unlock_local();
+                migration.progress.next()
+            }
+            MigrationProgress::Completed => {
+                clear_timer(migration.timer_id);
+                migration.progress.next()
+            }
+            MigrationProgress::Failed(e) => return Err(e),
+        },
+        None => return Err(MigrationError::NoMigrationInProgress),
     };
     set_progress(progress);
+    Ok(progress)
 }

src/backend/src/migrate/steps.rs

@@ -0,0 +1,94 @@
+use crate::{modify_state_config, mutate_state};
+use ic_cdk::eprintln;
+use shared::{
+    backend_api::Service,
+    types::{ApiEnabled, Guards, Migration, MigrationError, Stats},
+};
+
+/// Makes the local canister APIs readonly.
+pub fn make_this_readonly() {
+    mutate_state(|state| {
+        modify_state_config(state, |config| {
+            config.api = Some(Guards {
+                threshold_key: ApiEnabled::ReadOnly,
+                user_data: ApiEnabled::ReadOnly,
+            });
+        });
+    });
+}
+
+/// Locks the migration target canister APIs.
+pub async fn lock_migration_target(migration: &Migration) -> Result<(), MigrationError> {
+    Service(migration.to)
+        .set_guards(Guards {
+            threshold_key: ApiEnabled::Disabled,
+            user_data: ApiEnabled::Disabled,
+        })
+        .await
+        .map_err(|e| {
+            eprintln!("Failed to lock target canister: {:?}", e);
+            MigrationError::TargetLockFailed
+        })
+}
+
+/// Verifies that there is no data in the migration target canister.
+pub async fn assert_target_empty(migration: &Migration) -> Result<(), MigrationError> {
+    let stats = Service(migration.to)
+        .stats()
+        .await
+        .map_err(|e| {
+            eprintln!("Failed to get stats from the target canister: {:?}", e);
+            MigrationError::CouldNotGetTargetPriorStats
+        })?
+        .0;
+    if stats != Stats::default() {
+        return Err(MigrationError::TargetCanisterNotEmpty(stats));
+    }
+    Ok(())
+}
+
+/// Verifies that the target canister has all the data.
+pub async fn assert_target_has_all_data(migration: &Migration) -> Result<(), MigrationError> {
+    let source_stats = crate::stats();
+    let target_stats = Service(migration.to)
+        .stats()
+        .await
+        .map_err(|e| {
+            eprintln!("Failed to get stats from the target canister: {e:?}");
+            MigrationError::CouldNotGetTargetPostStats
+        })?
+        .0;
+    if source_stats != target_stats {
+        return Err(MigrationError::TargetStatsMismatch(
+            source_stats,
+            target_stats,
+        ));
+    }
+    Ok(())
+}
+
+/// Unlocks the target canister APIs.
+pub async fn unlock_target(migration: &Migration) -> Result<(), MigrationError> {
+    Service(migration.to)
+        .set_guards(Guards {
+            threshold_key: ApiEnabled::Disabled,
+            user_data: ApiEnabled::Enabled,
+        })
+        .await
+        .map_err(|e| {
+            eprintln!("Failed to unlock target canister: {e:?}");
+            MigrationError::TargetUnlockFailed
+        })
+}
+
+/// Unlocks the local signing APIs.  Disables local user data APIs.
+pub fn unlock_local() {
+    mutate_state(|state| {
+        modify_state_config(state, |config| {
+            config.api = Some(Guards {
+                threshold_key: ApiEnabled::Enabled,
+                user_data: ApiEnabled::Disabled,
+            });
+        });
+    });
+}