From 235c3d4b6f9c840ceb92fe0fac6f5a8744e7f863 Mon Sep 17 00:00:00 2001 From: Martin Raszyk Date: Thu, 4 Dec 2025 17:35:44 +0100 Subject: [PATCH 1/3] refactor(migration-canister): cache last event for every migration --- rs/migration_canister/src/canister_state.rs | 44 +++++++++++++++------ rs/migration_canister/src/lib.rs | 37 +++++++++++++++++ 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/rs/migration_canister/src/canister_state.rs b/rs/migration_canister/src/canister_state.rs index 085003eba2fb..59a74f24c6df 100644 --- a/rs/migration_canister/src/canister_state.rs +++ b/rs/migration_canister/src/canister_state.rs @@ -10,7 +10,7 @@ use ic_stable_structures::{ memory_manager::{MemoryId, MemoryManager, VirtualMemory}, }; -use crate::{Event, MAX_ONGOING_VALIDATIONS, RequestState}; +use crate::{CanisterMigrationArgs, Event, MAX_ONGOING_VALIDATIONS, RequestState}; type Memory = VirtualMemory; @@ -38,9 +38,15 @@ thread_local! { /// with timestamps as keys and their counts as values. static LIMITER: RefCell> = RefCell::new(BTreeMap::init(MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(2))))); - static HISTORY: RefCell> = + /// Stores all events indexed by their sequence numbers + /// in the order of creation. + static HISTORY: RefCell> = RefCell::new(BTreeMap::init(MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(3))))); + /// Caches the index of the last event for a given pair of source and target canisters. + static LAST_EVENT: RefCell> = + RefCell::new(BTreeMap::init(MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(4))))); + // TODO: consider a fail counter for active requests. // This way we see if a request never makes progress which would // indicate a bug in this canister or a problem with a subnet. @@ -117,8 +123,8 @@ pub mod requests { // ============================== Events API ============================== // pub mod events { use crate::{ - Event, EventType, - canister_state::{HISTORY, LIMITER}, + CanisterMigrationArgs, Event, EventType, + canister_state::{HISTORY, LAST_EVENT, LIMITER}, }; use candid::Principal; use ic_cdk::api::time; @@ -135,16 +141,32 @@ pub mod events { }); } let event = Event { time, event }; - HISTORY.with_borrow_mut(|h| h.insert(event, ())); + let args: CanisterMigrationArgs = (&event).into(); + let idx = HISTORY.with_borrow_mut(|h| { + let idx = h.len(); + h.insert(idx, event); + idx + }); + LAST_EVENT.with_borrow_mut(|l| { + l.insert(args, idx); + }); } pub fn find_last_event(source: Principal, target: Principal) -> Option { - // TODO: should do a range scan for efficiency. - HISTORY.with_borrow(|r| { - r.keys() - .rev() - .find(|x| x.event.request().source == source && x.event.request().target == target) - }) + let idx = LAST_EVENT.with_borrow(|l| { + let args = CanisterMigrationArgs { source, target }; + l.get(&args) + }); + if let Some(idx) = idx { + HISTORY.with_borrow(|h| { + let event = h + .get(&idx) + .expect("Missing event in history (this is a bug!)"); + Some(event) + }) + } else { + None + } } } diff --git a/rs/migration_canister/src/lib.rs b/rs/migration_canister/src/lib.rs index 86c79d4398cc..0c0f52ee0753 100644 --- a/rs/migration_canister/src/lib.rs +++ b/rs/migration_canister/src/lib.rs @@ -77,6 +77,12 @@ pub enum ValidationError { }, } +#[derive(Clone, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)] +struct CanisterMigrationArgs { + pub source: Principal, + pub target: Principal, +} + #[derive(Clone, PartialOrd, Ord, PartialEq, Eq, Serialize, Deserialize)] pub struct Request { source: Principal, @@ -137,6 +143,15 @@ impl Display for Request { } } +impl From<&Request> for CanisterMigrationArgs { + fn from(request: &Request) -> Self { + Self { + source: request.source, + target: request.target, + } + } +} + /// Represents the state a `Request` is currently in and contains all data necessary /// to transition to the next state (and sometimes data for a future state). /// @@ -305,6 +320,28 @@ impl Display for Event { } } +impl From<&Event> for CanisterMigrationArgs { + fn from(x: &Event) -> Self { + x.event.request().into() + } +} + +impl Storable for CanisterMigrationArgs { + fn to_bytes(&self) -> Cow<'_, [u8]> { + Cow::Owned(to_vec(&self).expect("Canister migration argument serialization failed")) + } + + fn into_bytes(self) -> Vec { + self.to_bytes().to_vec() + } + + fn from_bytes(bytes: Cow<[u8]>) -> Self { + from_slice(&bytes).expect("Canister migration argument deserialization failed") + } + + const BOUND: Bound = Bound::Unbounded; +} + impl Storable for Request { fn to_bytes(&self) -> Cow<'_, [u8]> { Cow::Owned(to_vec(&self).expect("Request serialization failed")) From 7bb17b88326816f5a9bb9490324a31a641eb5bd0 Mon Sep 17 00:00:00 2001 From: mraszyk <31483726+mraszyk@users.noreply.github.com> Date: Fri, 5 Dec 2025 14:19:27 +0100 Subject: [PATCH 2/3] Update rs/migration_canister/src/canister_state.rs Co-authored-by: michael-weigelt <122277901+michael-weigelt@users.noreply.github.com> --- rs/migration_canister/src/canister_state.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rs/migration_canister/src/canister_state.rs b/rs/migration_canister/src/canister_state.rs index 59a74f24c6df..16367bf14bcb 100644 --- a/rs/migration_canister/src/canister_state.rs +++ b/rs/migration_canister/src/canister_state.rs @@ -141,7 +141,7 @@ pub mod events { }); } let event = Event { time, event }; - let args: CanisterMigrationArgs = (&event).into(); + let args = CanisterMigrationArgs::from(&event); let idx = HISTORY.with_borrow_mut(|h| { let idx = h.len(); h.insert(idx, event); From 1491c9c026c902132d27197dce443dd61b316bf5 Mon Sep 17 00:00:00 2001 From: Martin Raszyk Date: Mon, 8 Dec 2025 16:51:22 +0100 Subject: [PATCH 3/3] log error instead of panicking --- rs/migration_canister/src/canister_state.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/rs/migration_canister/src/canister_state.rs b/rs/migration_canister/src/canister_state.rs index 16367bf14bcb..f51cde79a7b9 100644 --- a/rs/migration_canister/src/canister_state.rs +++ b/rs/migration_canister/src/canister_state.rs @@ -159,10 +159,11 @@ pub mod events { }); if let Some(idx) = idx { HISTORY.with_borrow(|h| { - let event = h - .get(&idx) - .expect("Missing event in history (this is a bug!)"); - Some(event) + let event = h.get(&idx); + if event.is_none() { + println!("Missing event for source={} and target={} with idx={} in history! This is a bug!", source, target, idx); + } + event }) } else { None