Re-Insert Memory Reserve for Queries and Upgrades (#4158)

With the incremental GC, programs can scale to the full 4GB memory space. As a consequence, update calls can allocate the full memory space, such that even a simple query, a simple composite query, or a simple upgrade logic can no longer succeed. This could also happen with classical non-copying GCs if deterministic time slicing is sufficiently extended. This PR prevents update calls (including canister initialization, heartbeats, and timers) from allocating the full memory by leaving a reserve for queries, composite queries, and canister upgrades. During queries, composite queries, and canister upgrades, garbage collection is suspended, such that the reserve is available to the mutator code. Callbacks of composite queries can also use the memory reserve. The current allocation limit for upgrade calls is 3.75 GB and applies to all GCs. This gives a reserve of 224 MB for the incremental GC as the last 32MB partition is unallocated in the current design. The scheduling heuristics of the incremental and generational GC needs to consider the reduced capacity for determining memory shortage. This PR can be viewed as a temporary measure until a memory reserve is implemented in the IC runtime system. The memory reserve may **not** be sufficient for a complex canister upgrade logic or large amount of stable heap data. Moreover, the upgrade logic may exceed the instruction limit. Thorough upgrade testing is required for canisters in any case.
dfinity · Sep 27, 2023 · c4c98d7 · c4c98d7
1 parent 64bad1a
commit c4c98d7
Show file tree

Hide file tree

Showing 17 changed files with 233 additions and 22 deletions.
diff --git a/rts/motoko-rts/src/gc/incremental.rs b/rts/motoko-rts/src/gc/incremental.rs
@@ -70,9 +70,10 @@ unsafe fn incremental_gc<M: Memory>(mem: &mut M) {
 #[cfg(feature = "ic")]
 unsafe fn should_start() -> bool {
     use self::partitioned_heap::PARTITION_SIZE;
-    use crate::memory::ic;
+    use crate::memory::{ic, MEMORY_RESERVE};
 
-    const CRITICAL_HEAP_LIMIT: Bytes<u32> = Bytes(u32::MAX - 768 * 1024 * 1024);
+    const CRITICAL_HEAP_LIMIT: Bytes<u32> =
+        Bytes(u32::MAX - 768 * 1024 * 1024 - MEMORY_RESERVE as u32);
     const CRITICAL_GROWTH_THRESHOLD: f64 = 0.01;
     const NORMAL_GROWTH_THRESHOLD: f64 = 0.65;
 

diff --git a/rts/motoko-rts/src/memory.rs b/rts/motoko-rts/src/memory.rs
@@ -7,6 +7,11 @@ use crate::types::*;
 
 use motoko_rts_macros::ic_mem_fn;
 
+// Memory reserve in bytes ensured during update and initialization calls.
+// For use by queries and upgrade calls.
+#[cfg(feature = "ic")]
+pub(crate) const MEMORY_RESERVE: usize = 256 * 1024 * 1024;
+
 /// A trait for heap allocation. RTS functions allocate in heap via this trait.
 ///
 /// To be able to link the RTS with moc-generated code, we implement wrappers around allocating

diff --git a/rts/motoko-rts/src/memory/ic.rs b/rts/motoko-rts/src/memory/ic.rs
@@ -2,10 +2,15 @@
 
 use super::Memory;
 use crate::constants::WASM_PAGE_SIZE;
+use crate::memory::MEMORY_RESERVE;
 use crate::rts_trap_with;
 use crate::types::{Bytes, Value, Words};
 use core::arch::wasm32;
 
+extern "C" {
+    fn keep_memory_reserve() -> bool;
+}
+
 /// Provides a `Memory` implementation, to be used in functions compiled for IC or WASI. The
 /// `Memory` implementation allocates in Wasm heap with Wasm `memory.grow` instruction.
 pub struct IcMemory;
@@ -20,9 +25,16 @@ impl Memory for IcMemory {
     /// with the slight exception of not allocating the extra page for address 0xFFFF_0000.
     #[inline(never)]
     unsafe fn grow_memory(&mut self, ptr: u64) {
-        debug_assert_eq!(0xFFFF_0000, usize::MAX - WASM_PAGE_SIZE.as_usize() + 1);
-        if ptr > 0xFFFF_0000 {
-            // spare the last wasm memory page
+        const LAST_PAGE_LIMIT: usize = 0xFFFF_0000;
+        debug_assert_eq!(LAST_PAGE_LIMIT, usize::MAX - WASM_PAGE_SIZE.as_usize() + 1);
+        let limit = if keep_memory_reserve() {
+            // Spare a memory reserve during update and initialization calls for use by queries and upgrades.
+            usize::MAX - MEMORY_RESERVE + 1
+        } else {
+            // Spare the last Wasm memory page on queries and upgrades to support the Rust call stack boundary checks.
+            LAST_PAGE_LIMIT
+        };
+        if ptr > limit as u64 {
             rts_trap_with("Cannot grow memory")
         };
         let page_size = u64::from(WASM_PAGE_SIZE.as_u32());

diff --git a/src/codegen/compile.ml b/src/codegen/compile.ml
@@ -4199,7 +4199,7 @@ module Lifecycle = struct
     | InPreUpgrade -> [Idle]
     | PostPreUpgrade -> [InPreUpgrade]
     | InPostUpgrade -> [InInit]
-    | InComposite -> [Idle]
+    | InComposite -> [Idle; InComposite]
 
   let get env =
     compile_unboxed_const (ptr ()) ^^
@@ -4226,6 +4226,10 @@ module Lifecycle = struct
       set env new_state
     )
 
+  let is_in env state =
+    get env ^^
+    compile_eq_const (int_of_state state)
+
 end (* Lifecycle *)
 
 
@@ -5335,7 +5339,23 @@ module RTS_Exports = struct
       edesc = nr (FuncExport (nr rts_trap_fi))
     });
 
-  let ic0_stable64_write_fi =
+    (* Keep a memory reserve when in update or init state. 
+       This reserve can be used by queries, composite queries, and upgrades. *)
+    let keep_memory_reserve_fi = E.add_fun env "keep_memory_reserve" (
+      Func.of_body env [] [I32Type] (fun env ->
+        Lifecycle.get env ^^
+        compile_eq_const Lifecycle.(int_of_state InUpdate) ^^
+        Lifecycle.get env ^^
+        compile_eq_const Lifecycle.(int_of_state InInit) ^^
+        G.i (Binary (Wasm.Values.I32 I32Op.Or))
+      )
+    ) in
+    E.add_export env (nr {
+      name = Lib.Utf8.decode "keep_memory_reserve";
+      edesc = nr (FuncExport (nr keep_memory_reserve_fi))
+    });
+
+    let ic0_stable64_write_fi =
       if E.mode env = Flags.WASIMode then
         E.add_fun env "ic0_stable64_write" (
             Func.of_body env ["to", I64Type; "from", I64Type; "len", I64Type] []
@@ -8150,9 +8170,25 @@ module FuncDec = struct
       | Type.Shared Type.Query ->
         Lifecycle.trans env Lifecycle.PostQuery
       | Type.Shared Type.Composite ->
-        Lifecycle.trans env Lifecycle.Idle
+        (* Stay in composite query state such that callbacks of 
+        composite queries can also use the memory reserve. 
+        The state is isolated since memory changes of queries 
+        are rolled back by the IC runtime system. *)
+        Lifecycle.trans env Lifecycle.InComposite
       | _ -> assert false
 
+  let callback_start env =
+    Lifecycle.is_in env Lifecycle.InComposite ^^
+    G.if0
+      (G.nop)
+      (message_start env (Type.Shared Type.Write))
+
+  let callback_cleanup env =
+    Lifecycle.is_in env Lifecycle.InComposite ^^
+    G.if0
+      (G.nop)
+      (message_cleanup env (Type.Shared Type.Write))
+
   let compile_const_message outer_env outer_ae sort control args mk_body ret_tys at : E.func_with_names =
     let ae0 = VarEnv.mk_fun_ae outer_ae in
     Func.of_body outer_env [] [] (fun env -> G.with_region at (
@@ -8309,7 +8345,7 @@ module FuncDec = struct
            (fun env -> compile_unboxed_const 0l)))
     in
     Func.define_built_in env reply_name ["env", I32Type] [] (fun env ->
-        message_start env (Type.Shared Type.Write) ^^
+        callback_start env ^^
         (* Look up continuation *)
         let (set_closure, get_closure) = new_local env "closure" in
         G.i (LocalGet (nr 0l)) ^^
@@ -8325,12 +8361,12 @@ module FuncDec = struct
         get_closure ^^
         Closure.call_closure env arity 0 ^^
 
-        message_cleanup env (Type.Shared Type.Write)
+        callback_cleanup env
       );
 
     let reject_name = "@reject_callback" in
     Func.define_built_in env reject_name ["env", I32Type] [] (fun env ->
-        message_start env (Type.Shared Type.Write) ^^
+        callback_start env ^^
         (* Look up continuation *)
         let (set_closure, get_closure) = new_local env "closure" in
         G.i (LocalGet (nr 0l)) ^^
@@ -8347,7 +8383,7 @@ module FuncDec = struct
         get_closure ^^
         Closure.call_closure env 1 0 ^^
 
-        message_cleanup env (Type.Shared Type.Write)
+        callback_cleanup env
       );
 
     (* result is a function that accepts a list of closure getters, from which

diff --git a/test/bench/ok/alloc.drun-run-opt.ok b/test/bench/ok/alloc.drun-run-opt.ok
@@ -1,8 +1,8 @@
 ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_764_731_706)
+debug.print: (+335_544_320, 4_764_731_816)
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_764_731_523)
+debug.print: (+335_544_320, 4_764_731_633)
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_764_731_535)
+debug.print: (+335_544_320, 4_764_731_645)
 ingress Completed: Reply: 0x4449444c0000
diff --git a/test/bench/ok/alloc.drun-run.ok b/test/bench/ok/alloc.drun-run.ok
@@ -1,8 +1,8 @@
 ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_898_949_603)
+debug.print: (+335_544_320, 4_898_949_723)
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_898_949_410)
+debug.print: (+335_544_320, 4_898_949_530)
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (+335_544_320, 4_898_949_423)
+debug.print: (+335_544_320, 4_898_949_543)
 ingress Completed: Reply: 0x4449444c0000
diff --git a/test/bench/ok/heap-32.drun-run-opt.ok b/test/bench/ok/heap-32.drun-run-opt.ok
@@ -1,5 +1,5 @@
 ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (50_227, +37_021_304, 767_984_255)
-debug.print: (50_070, +40_056_236, 830_437_406)
+debug.print: (50_227, +37_021_304, 767_984_266)
+debug.print: (50_070, +40_056_236, 830_437_417)
 ingress Completed: Reply: 0x4449444c0000
diff --git a/test/bench/ok/heap-32.drun-run.ok b/test/bench/ok/heap-32.drun-run.ok
@@ -1,5 +1,5 @@
 ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
 ingress Completed: Reply: 0x4449444c0000
-debug.print: (50_227, +37_021_304, 815_894_722)
-debug.print: (50_070, +40_056_236, 880_136_109)
+debug.print: (50_227, +37_021_304, 815_894_734)
+debug.print: (50_070, +40_056_236, 880_136_121)
 ingress Completed: Reply: 0x4449444c0000
diff --git a/test/run-drun-non-ci/memory-reserve-composite.drun b/test/run-drun-non-ci/memory-reserve-composite.drun
@@ -0,0 +1,6 @@
+# INCREMENTAL-GC-ONLY
+# SKIP ic-ref-run
+install $ID memory-reserve-composite/memory-reserve-composite.mo ""
+ingress $ID prepare1 "DIDL\x00\x00"
+ingress $ID prepare2 "DIDL\x00\x00"
+query $ID allocateInCompositeQuery "DIDL\x00\x00"
diff --git a/test/run-drun-non-ci/memory-reserve-composite/memory-reserve-composite.mo b/test/run-drun-non-ci/memory-reserve-composite/memory-reserve-composite.mo
@@ -0,0 +1,39 @@
+import Prim "mo:⛔";
+actor {
+    stable var stableData = Prim.Array_tabulate<Nat>(1024 * 1024, func(index) { index });
+    var array0 : [var Nat] = [var];
+    var array1 : [var Nat] = [var];
+    var array2 : [var Nat] = [var];
+    var array3 : [var Nat] = [var];
+    Prim.debugPrint("Initialized " # debug_show (Prim.rts_memory_size()));
+
+    public func prepare1() : async () {
+        array0 := Prim.Array_init<Nat>(256 * 1024 * 1024, 0); // 1GB
+        array1 := Prim.Array_init<Nat>(256 * 1024 * 1024, 1); // 2GB
+        Prim.debugPrint("Prepared1 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public func prepare2() : async () {
+        array2 := Prim.Array_init<Nat>(256 * 1024 * 1024, 2); // 3GB
+        array3 := Prim.Array_init<Nat>(150 * 1024 * 1024, 3); // around 3.75GB
+        Prim.debugPrint("Prepared2 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public composite query func allocateInCompositeQuery() : async () {
+        ignore Prim.Array_init<Nat>(50 * 1024 * 1024, 4);
+        Prim.debugPrint("Composite query call " # debug_show (Prim.rts_memory_size()));
+        assert (Prim.rts_memory_size() > 3840 * 1024 * 1024);
+        await nestedQuery();
+        ignore Prim.Array_init<Nat>(5 * 1024 * 1024, 4);
+        Prim.debugPrint("Composite query callback " # debug_show (Prim.rts_memory_size()));
+        assert (Prim.rts_memory_size() > 3840 * 1024 * 1024);
+    };
+
+    public query func nestedQuery() : async () {
+        Prim.debugPrint("Nested query " # debug_show (Prim.rts_memory_size()));
+    };
+};
+
+//SKIP run
+//SKIP run-ir
+//SKIP run-low
diff --git a/test/run-drun-non-ci/memory-reserve-query.drun b/test/run-drun-non-ci/memory-reserve-query.drun
@@ -0,0 +1,6 @@
+# INCREMENTAL-GC-ONLY
+# SKIP ic-ref-run
+install $ID memory-reserve-query/memory-reserve-query.mo ""
+ingress $ID prepare1 "DIDL\x00\x00"
+ingress $ID prepare2 "DIDL\x00\x00"
+query $ID allocateInQuery "DIDL\x00\x00"
diff --git a/test/run-drun-non-ci/memory-reserve-query/memory-reserve-query.mo b/test/run-drun-non-ci/memory-reserve-query/memory-reserve-query.mo
@@ -0,0 +1,31 @@
+import Prim "mo:⛔";
+actor {
+    stable var stableData = Prim.Array_tabulate<Nat>(1024 * 1024, func(index) { index });
+    var array0 : [var Nat] = [var];
+    var array1 : [var Nat] = [var];
+    var array2 : [var Nat] = [var];
+    var array3 : [var Nat] = [var];
+    Prim.debugPrint("Initialized " # debug_show (Prim.rts_memory_size()));
+
+    public func prepare1() : async () {
+        array0 := Prim.Array_init<Nat>(256 * 1024 * 1024, 0); // 1GB
+        array1 := Prim.Array_init<Nat>(256 * 1024 * 1024, 1); // 2GB
+        Prim.debugPrint("Prepared1 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public func prepare2() : async () {
+        array2 := Prim.Array_init<Nat>(256 * 1024 * 1024, 2); // 3GB
+        array3 := Prim.Array_init<Nat>(150 * 1024 * 1024, 3); // around 3.75GB
+        Prim.debugPrint("Prepared2 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public query func allocateInQuery() : async () {
+        ignore Prim.Array_init<Nat>(50 * 1024 * 1024, 4);
+        Prim.debugPrint("Query call " # debug_show (Prim.rts_memory_size()));
+        assert (Prim.rts_memory_size() > 3840 * 1024 * 1024);
+    };
+};
+
+//SKIP run
+//SKIP run-ir
+//SKIP run-low
diff --git a/test/run-drun-non-ci/memory-reserve-update.drun b/test/run-drun-non-ci/memory-reserve-update.drun
@@ -0,0 +1,6 @@
+# INCREMENTAL-GC-ONLY
+# SKIP ic-ref-run
+install $ID memory-reserve-update/memory-reserve-update.mo ""
+ingress $ID prepare1 "DIDL\x00\x00"
+ingress $ID prepare2 "DIDL\x00\x00"
+ingress $ID allocateInUpdate "DIDL\x00\x00"
diff --git a/test/run-drun-non-ci/memory-reserve-update/memory-reserve-update.mo b/test/run-drun-non-ci/memory-reserve-update/memory-reserve-update.mo
@@ -0,0 +1,30 @@
+import Prim "mo:⛔";
+actor {
+    stable var stableData = Prim.Array_tabulate<Nat>(1024 * 1024, func(index) { index });
+    var array0 : [var Nat] = [var];
+    var array1 : [var Nat] = [var];
+    var array2 : [var Nat] = [var];
+    var array3 : [var Nat] = [var];
+    Prim.debugPrint("Initialized " # debug_show (Prim.rts_memory_size()));
+
+    public func prepare1() : async () {
+        array0 := Prim.Array_init<Nat>(256 * 1024 * 1024, 0); // 1GB
+        array1 := Prim.Array_init<Nat>(256 * 1024 * 1024, 1); // 2GB
+        Prim.debugPrint("Prepared1 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public func prepare2() : async () {
+        array2 := Prim.Array_init<Nat>(256 * 1024 * 1024, 2); // 3GB
+        array3 := Prim.Array_init<Nat>(150 * 1024 * 1024, 3); // around 3.75GB
+        Prim.debugPrint("Prepared2 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public func allocateInUpdate() : async () {
+        Prim.debugPrint("Update call " # debug_show (Prim.rts_memory_size()));
+        ignore Prim.Array_init<Nat>(50 * 1024 * 1024, 4);
+    };
+};
+
+//SKIP run
+//SKIP run-ir
+//SKIP run-low
diff --git a/test/run-drun-non-ci/memory-reserve-upgrade.drun b/test/run-drun-non-ci/memory-reserve-upgrade.drun
@@ -0,0 +1,6 @@
+# INCREMENTAL-GC-ONLY
+# SKIP ic-ref-run
+install $ID memory-reserve-upgrade/memory-reserve-upgrade.mo ""
+ingress $ID prepare1 "DIDL\x00\x00"
+ingress $ID prepare2 "DIDL\x00\x00"
+upgrade $ID memory-reserve-upgrade/memory-reserve-upgrade.mo ""
diff --git a/test/run-drun-non-ci/memory-reserve-upgrade/memory-reserve-upgrade.mo b/test/run-drun-non-ci/memory-reserve-upgrade/memory-reserve-upgrade.mo
@@ -0,0 +1,25 @@
+import Prim "mo:⛔";
+actor {
+    stable var stableData = Prim.Array_tabulate<Nat>(1024 * 1024, func(index) { index });
+    var array0 : [var Nat] = [var];
+    var array1 : [var Nat] = [var];
+    var array2 : [var Nat] = [var];
+    var array3 : [var Nat] = [var];
+    Prim.debugPrint("Initialized " # debug_show (Prim.rts_memory_size()));
+
+    public func prepare1() : async () {
+        array0 := Prim.Array_init<Nat>(256 * 1024 * 1024, 0); // 1GB
+        array1 := Prim.Array_init<Nat>(256 * 1024 * 1024, 1); // 2GB
+        Prim.debugPrint("Prepared1 " # debug_show (Prim.rts_memory_size()));
+    };
+
+    public func prepare2() : async () {
+        array2 := Prim.Array_init<Nat>(256 * 1024 * 1024, 2); // 3GB
+        array3 := Prim.Array_init<Nat>(150 * 1024 * 1024, 3); // around 3.75GB
+        Prim.debugPrint("Prepared2 " # debug_show (Prim.rts_memory_size()));
+    };
+};
+
+//SKIP run
+//SKIP run-ir
+//SKIP run-low
diff --git a/test/run.sh b/test/run.sh
@@ -435,6 +435,14 @@ do
         continue
       fi
 
+      if grep -q "# *INCREMENTAL-GC-ONLY" $(basename $file)
+      then
+        if [[ $EXTRA_MOC_ARGS != *"--incremental-gc"* ]]
+        then 
+          continue
+        fi
+      fi
+
       have_var_name="HAVE_${runner//-/_}"
       if [ ${!have_var_name} != yes ]
       then