chore(migration-canister): return only a single migration status

packages/pocket-ic/tests/icp_features.rs

@@ -128,14 +128,14 @@ fn test_canister_migration() {
     res.0.unwrap();
 
     let status = || {
-        let res = update_candid::<_, (Vec<MigrationStatus>,)>(
+        let res = update_candid::<_, (Option<MigrationStatus>,)>(
             &pic,
             canister_migration_orchestrator,
             "migration_status",
             (migrate_canister_args.clone(),),
         )
         .unwrap();
-        res.0.last().unwrap().clone()
+        res.0.unwrap()
     };
     loop {
         pic.advance_time(Duration::from_secs(10));

rs/migration_canister/migration_canister.did

@@ -33,7 +33,5 @@ service : (MigrationCanisterInitArgs) -> {
   disable_api : () -> (MigrationCanisterResult);
   enable_api : () -> (MigrationCanisterResult);
   migrate_canister : (MigrateCanisterArgs) -> (ValidationResult);
-  // The same `(canister_id, replace_canister_id)` pair might be processed again after a failure
-  // and thus we return a vector.
-  migration_status : (MigrateCanisterArgs) -> (vec MigrationStatus) query;
+  migration_status : (MigrateCanisterArgs) -> (opt MigrationStatus) query;
 }

rs/migration_canister/src/canister_state.rs

@@ -94,13 +94,20 @@ pub mod requests {
         REQUESTS.with_borrow(|req| req.keys().filter(predicate).collect())
     }
 
-    pub fn find_request(source: Principal, target: Principal) -> Vec<RequestState> {
-        // TODO: should do a range scan for efficiency.
-        REQUESTS.with_borrow(|r| {
+    pub fn find_request(source: Principal, target: Principal) -> Option<RequestState> {
+        // We perform a linear scan here which is fine since
+        // there can only be at most `RATE_LIMIT` (50) requests
+        // at any time.
+        let requests: Vec<_> = REQUESTS.with_borrow(|r| {
             r.keys()
                 .filter(|x| x.request().source == source && x.request().target == target)
                 .collect()
-        })
+        });
+        assert!(
+            requests.len() <= 1,
+            "There should only be a single request for a given pair of canister IDs."
+        );
+        requests.first().cloned()
     }
 }
 
@@ -136,14 +143,12 @@ pub mod events {
         })
     }
 
-    pub fn find_event(source: Principal, target: Principal) -> Vec<Event> {
+    pub fn find_last_event(source: Principal, target: Principal) -> Option<Event> {
         // TODO: should do a range scan for efficiency.
         HISTORY.with_borrow(|r| {
             r.keys()
-                .filter(|x| {
-                    x.event.request().source == source && x.event.request().target == target
-                })
-                .collect()
+                .rev()
+                .find(|x| x.event.request().source == source && x.event.request().target == target)
         })
     }
 }

rs/migration_canister/src/migration_canister.rs

@@ -13,7 +13,7 @@ use crate::{
     RequestState, ValidationError,
     canister_state::{
         ValidationGuard, caller_allowed,
-        events::find_event,
+        events::find_last_event,
         migrations_disabled,
         requests::{find_request, insert_request},
         set_allowlist,
@@ -100,25 +100,22 @@ pub enum MigrationStatus {
 }
 
 #[query]
-/// The same (canister_id, replace_canister_id) pair might be present in the `HISTORY`, and valid to process again, so
-/// we return a vector.
-fn migration_status(args: MigrateCanisterArgs) -> Vec<MigrationStatus> {
-    let mut active: Vec<MigrationStatus> = find_request(args.canister_id, args.replace_canister_id)
-        .into_iter()
-        .map(|r| MigrationStatus::InProgress {
-            status: r.name().to_string(),
-        })
-        .collect();
-    let events: Vec<MigrationStatus> = find_event(args.canister_id, args.replace_canister_id)
-        .into_iter()
-        .map(|event| match event.event {
+fn migration_status(args: MigrateCanisterArgs) -> Option<MigrationStatus> {
+    if let Some(request_status) = find_request(args.canister_id, args.replace_canister_id) {
+        let migration_status = MigrationStatus::InProgress {
+            status: request_status.name().to_string(),
+        };
+        Some(migration_status)
+    } else if let Some(event) = find_last_event(args.canister_id, args.replace_canister_id) {
+        let migration_status = match event.event {
             crate::EventType::Succeeded { .. } => MigrationStatus::Succeeded { time: event.time },
             crate::EventType::Failed { reason, .. } => MigrationStatus::Failed {
                 reason,
                 time: event.time,
             },
-        })
-        .collect();
-    active.extend(events);
-    active
+        };
+        Some(migration_status)
+    } else {
+        None
+    }
 }

rs/migration_canister/src/tests.rs

@@ -28,7 +28,7 @@ fn test() {
         caller,
     );
     insert_request(RequestState::Accepted { request });
-    assert!(find_request(source, target).len() == 1);
+    assert!(find_request(source, target).is_some());
 }
 
 #[test]

rs/migration_canister/tests/tests.rs

@@ -251,7 +251,7 @@ async fn get_status(
     pic: &PocketIc,
     sender: Principal,
     args: &MigrateCanisterArgs,
-) -> Vec<MigrationStatus> {
+) -> Option<MigrationStatus> {
     let res = pic
         .update_call(
             MIGRATION_CANISTER_ID.into(),
@@ -261,7 +261,7 @@ async fn get_status(
         )
         .await
         .unwrap();
-    Decode!(&res, Vec<MigrationStatus>).unwrap()
+    Decode!(&res, Option<MigrationStatus>).unwrap()
 }
 
 /// Advances time by a second and executes enough ticks that the state machine
@@ -593,7 +593,7 @@ async fn replay_call_after_migration() {
 
     loop {
         let status = get_status(&pic, sender, &args).await;
-        if let MigrationStatus::Succeeded { .. } = status[0] {
+        if let MigrationStatus::Succeeded { .. } = status.unwrap() {
             break;
         }
         // We proceed in small steps here so that
@@ -1156,7 +1156,7 @@ async fn status_correct() {
 
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "Accepted".to_string()
         }
@@ -1165,7 +1165,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "ControllersChanged".to_string()
         }
@@ -1174,7 +1174,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "StoppedAndReady".to_string()
         }
@@ -1183,7 +1183,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "RenamedTarget".to_string()
         }
@@ -1192,7 +1192,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "UpdatedRoutingTable".to_string()
         }
@@ -1201,7 +1201,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "RoutingTableChangeAccepted".to_string()
         }
@@ -1210,7 +1210,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "SourceDeleted".to_string()
         }
@@ -1219,7 +1219,7 @@ async fn status_correct() {
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
     assert_eq!(
-        status[0],
+        status.unwrap(),
         MigrationStatus::InProgress {
             status: "RestoredControllers".to_string()
         }
@@ -1249,7 +1249,7 @@ async fn after_validation_source_not_stopped() {
     advance(&pic).await;
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Failed { ref reason, .. } = status[0] else {
+    let MigrationStatus::Failed { ref reason, .. } = status.unwrap() else {
         panic!()
     };
     assert_eq!(reason, &"Source is not stopped.".to_string());
@@ -1278,7 +1278,7 @@ async fn after_validation_target_not_stopped() {
     advance(&pic).await;
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Failed { ref reason, .. } = status[0] else {
+    let MigrationStatus::Failed { ref reason, .. } = status.unwrap() else {
         panic!()
     };
     assert_eq!(reason, &"Target is not stopped.".to_string());
@@ -1319,7 +1319,7 @@ async fn after_validation_target_has_snapshot() {
     advance(&pic).await;
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Failed { ref reason, .. } = status[0] else {
+    let MigrationStatus::Failed { ref reason, .. } = status.unwrap() else {
         panic!()
     };
     assert_eq!(reason, &"Target has snapshots.".to_string());
@@ -1361,7 +1361,7 @@ async fn after_validation_insufficient_cycles() {
     advance(&pic).await;
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Failed { ref reason, .. } = status[0] else {
+    let MigrationStatus::Failed { ref reason, .. } = status.unwrap() else {
         panic!()
     };
     assert!(reason.contains("Source does not have sufficient cycles"));
@@ -1391,7 +1391,7 @@ async fn failure_controllers_restored() {
     advance(&pic).await;
     advance(&pic).await;
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Failed { .. } = status[0] else {
+    let MigrationStatus::Failed { .. } = status.unwrap() else {
         panic!()
     };
     let mut source_controllers_after = pic.get_controllers(source).await;
@@ -1430,8 +1430,8 @@ async fn success_controllers_restored() {
         advance(&pic).await;
     }
     let status = get_status(&pic, sender, &args).await;
-    let MigrationStatus::Succeeded { .. } = status[0] else {
-        panic!("status: {:?}", status[0]);
+    let MigrationStatus::Succeeded { .. } = status.as_ref().unwrap() else {
+        panic!("status: {:?}", status.unwrap());
     };
     let mut source_controllers_after = pic.get_controllers(source).await;
     source_controllers_after.sort();
@@ -1494,7 +1494,7 @@ async fn parallel_migrations() {
             },
         )
         .await;
-        let MigrationStatus::InProgress { ref status } = status[0] else {
+        let MigrationStatus::InProgress { ref status } = status.unwrap() else {
             panic!()
         };
         assert_eq!(status, "SourceDeleted");

rs/tests/execution/canister_migration_test.rs

@@ -431,13 +431,13 @@ async fn test_async(env: TestEnv) {
 
     assert_eq!(decoded_result, Ok(()));
 
-    // The migration canister has a step where it waits for 5 minutes, so we give it a minute more than that.
-    println!("Wait over 5 minutes for processing.");
+    // The migration canister has a step where it waits for 6 minutes, so we give it a minute more than that.
+    println!("Wait 7 minutes for processing.");
 
     retry_with_msg_async!(
-        "Wait 5m for migration canister to process",
+        "Wait 7m for migration canister to process",
         &logger,
-        Duration::from_secs(360),
+        Duration::from_secs(420),
         Duration::from_secs(10),
         || async {
             let status = nns_agent
@@ -446,13 +446,14 @@ async fn test_async(env: TestEnv) {
                 .call_and_wait()
                 .await
                 .expect("Failed to call migration_status.");
-            let decoded_status = Decode!(&status, Vec<MigrationStatus>)
-                .expect("Failed to decode response from migration_status.");
+            let decoded_status = Decode!(&status, Option<MigrationStatus>)
+                .expect("Failed to decode response from migration_status.")
+                .expect("There should be a migration status available.");
 
-            if matches!(decoded_status[0], MigrationStatus::Succeeded { .. }) {
+            if matches!(decoded_status, MigrationStatus::Succeeded { .. }) {
                 Ok(())
             } else {
-                bail!("Not ready. Status: {:?}", decoded_status[0])
+                bail!("Not ready. Status: {:?}", decoded_status)
             }
         }
     )