refactor(migration-canister): cache last event for every migration

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<DefaultMemoryImpl>;
 
@@ -38,9 +38,15 @@ thread_local! {
     /// with timestamps as keys and their counts as values.
     static LIMITER: RefCell<BTreeMap<u64, u64, Memory>> = RefCell::new(BTreeMap::init(MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(2)))));
 
-    static HISTORY: RefCell<BTreeMap<Event, (), Memory>> =
+    /// Stores all events indexed by their sequence numbers
+    /// in the order of creation.
+    static HISTORY: RefCell<BTreeMap<u64, Event, Memory>> =
         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<BTreeMap<CanisterMigrationArgs, u64, Memory>> =
+        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,33 @@ pub mod events {
             });
         }
         let event = Event { time, event };
-        HISTORY.with_borrow_mut(|h| h.insert(event, ()));
+        let args = CanisterMigrationArgs::from(&event);
+        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<Event> {
-        // 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);
+                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
+        }
     }
 }
 

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<u8> {
+        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"))