feat: fix graph copy stab/destab for multi-migration and dedup table (#5993)

The graph copy stabilization/destabilization algorithm basically does a
BFS walk to first serialize the heap to stable memory, then after an
upgrade a deserialization bringing back up the heap from stable memory.
This BFS was initially designed to just start the BFS walk out of the
stable actor record.

However, later additions to the enhanced OP incremental GC added 2
GC-only roots: the `blob deduplication table` and the `migration
functions list`. Without proper handling of the
serialization/deserialization, these objects would be lost and the new
actor version would work erroneously.

This PR modifies the BFS algorithm to add the two objects as roots in
the serialization/deserialization graph copy algorithm such that they
are first saved properly to stable memory and then restored correctly
into the heap. Other future GC roots can be treated similary.

To verify the functionality works correctly, two tests were modified to
trigger this behavior and check if the new upgraded actor works
correctly.

**Attention point:**
This PR also fixes a corner-case. For enhanced multi-migration, which
migration function (in the chain) is matched against the old actor type
is only determined at runtime (see `desugar.ml`), depending on which
migration function in the chain it's been brought at previously. The old
code in `compile_enhanced.ml` was performing compatibility checks
against the new actor `.pre`, which can be wrong for multi-migration in
case there might be multiple steps missing. Originally, the code made a
check against old actor and new actor `.pre` types early in the `start
destabilization` code path. This cannot work with the new migration.

The solution is to restore into `metadata` the old actor type from
stable memory and let the checks happen at `ICStableRead` time (via
`register_stable_type`, see `compile_enhanced.ml`). Previously, this
metadata was wiped and `ICStableRead` was not performing any checks,
just recovering the actor.

This way, the checks happen correctly at runtime for both enhanced
multi-migration and regular migration (or upgrade).

---------

Co-authored-by: Claudio Russo <claudio@dfinity.org>

Author: alexandru-uta

rts/motoko-rts-tests/src/stabilization.rs

@@ -14,7 +14,10 @@ use crate::{
 use motoko_rts::{
     memory::{alloc_array, Memory},
     stabilization::{
-        deserialization::Deserialization, graph_copy::GraphCopy, serialization::Serialization,
+        deserialization::Deserialization,
+        graph_copy::GraphCopy,
+        layout::StableValue,
+        serialization::{Serialization, SerializationRoots},
     },
     types::{Value, Words, TAG_ARRAY_M},
 };
@@ -168,14 +171,25 @@ fn test_serialization_deserialization(random: &mut Rand32, max_objects: usize, s
 
 fn serialize(old_stable_root: Value, stable_start: u64) -> u64 {
     let mut memory = TestMemory::new(Words(0));
-    let mut serialization = Serialization::start(&mut memory, old_stable_root, stable_start);
+    let roots = SerializationRoots {
+        actor: old_stable_root,
+        dedup_table: Value::from_raw(0),
+        migrations_list: Value::from_raw(0),
+    };
+    let mut serialization = Serialization::start(&mut memory, roots, stable_start);
     serialization.copy_increment(&mut memory);
     assert!(serialization.is_completed());
     serialization.serialized_data_length()
 }
 
 fn deserialize<M: Memory>(mem: &mut M, stable_start: u64, stable_size: u64) -> Value {
-    let mut deserialization = Deserialization::start(mem, stable_start, stable_size);
+    let mut deserialization = Deserialization::start(
+        mem,
+        stable_start,
+        stable_size,
+        StableValue::from_raw(0),
+        StableValue::from_raw(0),
+    );
     deserialization.copy_increment(mem);
     assert!(deserialization.is_completed());
     deserialization.get_stable_root()

rts/motoko-rts/src/persistence.rs

@@ -15,8 +15,9 @@ use crate::{
     persistence::compatibility::memory_compatible,
     region::{
         LEGACY_VERSION_NO_STABLE_MEMORY, LEGACY_VERSION_REGIONS, LEGACY_VERSION_SOME_STABLE_MEMORY,
-        VERSION_GRAPH_COPY_NO_REGIONS, VERSION_GRAPH_COPY_REGIONS, VERSION_STABLE_HEAP_NO_REGIONS,
-        VERSION_STABLE_HEAP_REGIONS,
+        VERSION_GRAPH_COPY_NO_REGIONS, VERSION_GRAPH_COPY_REGIONS,
+        VERSION_GRAPH_COPY_V1_NO_REGIONS, VERSION_GRAPH_COPY_V1_REGIONS,
+        VERSION_STABLE_HEAP_NO_REGIONS, VERSION_STABLE_HEAP_REGIONS,
     },
     rts_trap_with,
     stable_mem::read_persistence_version,
@@ -158,6 +159,8 @@ unsafe fn use_enhanced_orthogonal_persistence() -> bool {
         VERSION_STABLE_HEAP_NO_REGIONS | VERSION_STABLE_HEAP_REGIONS => true,
         VERSION_GRAPH_COPY_NO_REGIONS
         | VERSION_GRAPH_COPY_REGIONS
+        | VERSION_GRAPH_COPY_V1_NO_REGIONS
+        | VERSION_GRAPH_COPY_V1_REGIONS
         | LEGACY_VERSION_NO_STABLE_MEMORY
         | LEGACY_VERSION_SOME_STABLE_MEMORY
         | LEGACY_VERSION_REGIONS => false,
@@ -250,6 +253,14 @@ unsafe fn update_stable_type<M: Memory>(
     (*metadata).stable_type.assign(mem, &new_type);
 }
 
+/// Restore the old stable type from graph-copy stabilization metadata into
+/// the freshly initialized persistent metadata so that `register_stable_type`
+/// can check compatibility when the migration chain runs after destabilization.
+pub unsafe fn restore_stable_type<M: Memory>(mem: &mut M, old_type: &TypeDescriptor) {
+    let metadata = PersistentMetadata::get();
+    (*metadata).stable_type.assign(mem, old_type);
+}
+
 /// Register the stable actor type on canister initialization and upgrade.
 /// The type is stored in the persistent metadata memory for later retrieval on canister upgrades.
 /// On an upgrade, the memory compatibility between the new and existing stable type is checked.

rts/motoko-rts/src/region.rs

@@ -16,6 +16,10 @@ pub(crate) const VERSION_GRAPH_COPY_NO_REGIONS: usize = 3;
 pub(crate) const VERSION_GRAPH_COPY_REGIONS: usize = 4;
 pub(crate) const VERSION_STABLE_HEAP_NO_REGIONS: usize = 5;
 pub(crate) const VERSION_STABLE_HEAP_REGIONS: usize = 6;
+// V1 graph-copy: adds a 16-byte extension block in front of the legacy 40-byte
+// last-page record carrying extra GC roots (dedup table, migrations list).
+pub(crate) const VERSION_GRAPH_COPY_V1_NO_REGIONS: usize = 7;
+pub(crate) const VERSION_GRAPH_COPY_V1_REGIONS: usize = 8;
 
 const _: () = assert!(meta_data::size::PAGE_IN_BYTES == crate::stable_mem::PAGE_SIZE);
 const _: () = assert!(meta_data::size::PAGES_IN_BLOCK <= u8::MAX as u32);

rts/motoko-rts/src/stabilization/deserialization.rs

@@ -6,7 +6,7 @@ use crate::{
     gc::incremental::array_slicing::slice_array,
     memory::Memory,
     stabilization::deserialization::scan_stack::STACK_EMPTY,
-    types::{FwdPtr, Tag, Value, TAG_ARRAY_SLICE_MIN, TAG_FWD_PTR},
+    types::{FwdPtr, Tag, Value, NULL_POINTER, TAG_ARRAY_SLICE_MIN, TAG_FWD_PTR},
     visitor::visit_pointer_fields,
 };
 
@@ -26,6 +26,9 @@ pub struct Deserialization {
     stable_root: Option<Value>,
     limit: ExecutionMonitor,
     clear_position: u64,
+    /// Heap addresses for the helper GC roots.
+    pub dedup_table_address: Value,
+    pub migrations_list_address: Value,
 }
 
 /// Helper type to pass serialization context instead of closures.
@@ -58,7 +61,13 @@ impl<'a, M: Memory> DeserializationContext<'a, M> {
 /// mechanism to avoid instruction limit exceeding.
 impl Deserialization {
     /// Start the deserialization, followed by a series of copy increments.
-    pub fn start<M: Memory>(mem: &mut M, stable_start: u64, stable_size: u64) -> Deserialization {
+    pub fn start<M: Memory>(
+        mem: &mut M,
+        stable_start: u64,
+        stable_size: u64,
+        dedup_table_address: StableValue,
+        migrations_list_address: StableValue,
+    ) -> Deserialization {
         let from_space = StableMemoryAccess::open(stable_start, stable_size);
         let scan_stack = unsafe { ScanStack::new(mem) };
         let limit = ExecutionMonitor::new();
@@ -70,8 +79,18 @@ impl Deserialization {
             stable_root: None,
             limit,
             clear_position: stable_start,
+            dedup_table_address: NULL_POINTER,
+            migrations_list_address: NULL_POINTER,
         };
-        deserialization.start(mem, StableValue::serialize(Value::from_ptr(0)));
+        let _ = deserialization.start(mem, StableValue::serialize(Value::from_ptr(0)));
+        // Load up the heap addresses of the GC helper roots.
+        if dedup_table_address != StableValue::from_raw(0) {
+            deserialization.dedup_table_address = deserialization.start(mem, dedup_table_address);
+        }
+        if migrations_list_address != StableValue::from_raw(0) {
+            deserialization.migrations_list_address =
+                deserialization.start(mem, migrations_list_address);
+        }
         deserialization
     }
 

rts/motoko-rts/src/stabilization/graph_copy.rs

@@ -23,8 +23,8 @@ pub trait GraphCopy<S: Copy, T: Copy, P: Copy + Default> {
     ///     copy_algorthm.copy_increment();
     /// }
     /// ```
-    fn start<M: Memory>(&mut self, mem: &mut M, root: S) {
-        self.evacuate(mem, root);
+    fn start<M: Memory>(&mut self, mem: &mut M, root: S) -> T {
+        self.evacuate(mem, root)
     }
 
     /// Determine whether the scanning algorithm is completed,

rts/motoko-rts/src/stabilization/ic.rs

@@ -1,4 +1,4 @@
-mod metadata;
+pub mod metadata;
 mod performance;
 
 use motoko_rts_macros::ic_mem_fn;
@@ -7,11 +7,13 @@ use crate::{
     gc::incremental::{is_gc_stopped, resume_gc, stop_gc},
     memory::Memory,
     persistence::{
-        compatibility::{memory_compatible, TypeDescriptor},
+        compatibility::TypeDescriptor, get_dedup_table_ptr, get_migration_functions_ptr,
+        restore_stable_type, set_dedup_table_ptr, set_migration_functions_ptr,
         set_upgrade_instructions,
     },
     rts_trap_with,
     stabilization::ic::metadata::StabilizationMetadata,
+    stabilization::serialization::SerializationRoots,
     stable_mem::{self, moc_stable_mem_set_size, PAGE_SIZE},
     types::Value,
 };
@@ -25,8 +27,8 @@ struct StabilizationState {
     old_candid_data: Value,
     old_type_offsets: Value,
     completed: bool,
-    serialization: Serialization,
-    instruction_meter: InstructionMeter,
+    pub serialization: Serialization,
+    pub instruction_meter: InstructionMeter,
 }
 
 impl StabilizationState {
@@ -72,7 +74,12 @@ pub unsafe fn start_graph_stabilization<M: Memory>(
     assert!(is_gc_stopped());
     let stable_memory_pages = stable_mem::size(); // Backup the virtual size.
     let serialized_data_start = stable_memory_pages * PAGE_SIZE;
-    let serialization = Serialization::start(mem, stable_actor, serialized_data_start);
+    let serialization_roots = SerializationRoots {
+        actor: stable_actor,
+        dedup_table: *get_dedup_table_ptr(),
+        migrations_list: *get_migration_functions_ptr(),
+    };
+    let serialization = Serialization::start(mem, serialization_roots, serialized_data_start);
     STABILIZATION_STATE = Some(StabilizationState::new(
         serialization,
         old_candid_data,
@@ -126,7 +133,11 @@ unsafe fn write_metadata() {
         type_descriptor,
     };
     state.instruction_meter.stop();
-    metadata.store(&mut state.instruction_meter);
+    metadata.store(
+        &mut state.instruction_meter,
+        state.serialization.dedup_table_address,
+        state.serialization.migrations_list_address,
+    );
 }
 
 struct DestabilizationState {
@@ -140,33 +151,17 @@ static mut DESTABILIZATION_STATE: Option<DestabilizationState> = None;
 
 /// Starts the graph-copy-based destabilization process.
 /// This requires that the deserialization is subsequently run and completed.
-/// Also checks whether the new program version is compatible to the stored state by comparing the type
-/// tables of both the old and the new program version.
-/// The check is identical to enhanced orthogonal persistence, except that the metadata is obtained from
-/// stable memory and not the persistent main memory.
-/// The parameters encode the type table of the new program version to which that data is to be upgraded.
-/// `new_candid_data`: A blob encoding the Candid type as a table.
-/// `new_type_offsets`: A blob encoding the type offsets in the Candid type table.
-///   Type index 0 represents the stable actor object to be serialized.
-/// Traps if the stable state is incompatible with the new program version and the upgrade is not
-/// possible.
+/// The old type descriptor from stable memory is restored into `PersistentMetadata`
+/// so that `register_stable_type` can check compatibility when the migration chain
+/// runs after destabilization (inside `Persistence.load` / `ICStableRead`).
 #[ic_mem_fn(ic_only)]
-pub unsafe fn start_graph_destabilization<M: Memory>(
-    mem: &mut M,
-    new_candid_data: Value,
-    new_type_offsets: Value,
-) {
+pub unsafe fn start_graph_destabilization<M: Memory>(mem: &mut M) {
     assert!(DESTABILIZATION_STATE.is_none());
 
     let mut instruction_meter = InstructionMeter::new();
     instruction_meter.start();
-    let mut new_type_descriptor = TypeDescriptor::new(new_candid_data, new_type_offsets);
-    let (metadata, statistics) = StabilizationMetadata::load(mem);
-    let mut old_type_descriptor = metadata.type_descriptor;
-    if !memory_compatible(mem, &mut old_type_descriptor, &mut new_type_descriptor) {
-        rts_trap_with("Memory-incompatible program upgrade");
-    }
-    // Restore the virtual size.
+    let (metadata, last_page_record) = StabilizationMetadata::load(mem);
+    restore_stable_type(mem, &metadata.type_descriptor);
     moc_stable_mem_set_size(metadata.serialized_data_start / PAGE_SIZE);
 
     // Stop the GC until the incremental graph destabilization has been completed.
@@ -176,11 +171,13 @@ pub unsafe fn start_graph_destabilization<M: Memory>(
         mem,
         metadata.serialized_data_start,
         metadata.serialized_data_length,
+        last_page_record.dedup_table_address,
+        last_page_record.migrations_list_address,
     );
     instruction_meter.stop();
     DESTABILIZATION_STATE = Some(DestabilizationState {
         deserialization,
-        stabilization_statistics: statistics,
+        stabilization_statistics: last_page_record.statistics,
         completed: false,
         instruction_meter,
     });
@@ -214,6 +211,12 @@ pub unsafe fn graph_destabilization_increment<M: Memory>(mem: &mut M) -> bool {
         state.instruction_meter.stop();
         if state.deserialization.is_completed() {
             record_upgrade_costs();
+
+            // We need to put back in the metadata pointing to the
+            // helper GC roots for the dedup table and migration list.
+            set_dedup_table_ptr(mem, state.deserialization.dedup_table_address);
+            set_migration_functions_ptr(mem, state.deserialization.migrations_list_address);
+
             state.completed = true;
             memory_sanity_check(mem);
         }