use timer to self kill canisters after TTL (#153)

* use timer to self kill canisters after TTL

* fix

* add comments

* add timer in postupgrade

* bump moc

* fix

* add tests

* fix

service/pool/Main.mo

@@ -29,13 +29,15 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
     stable var stablePool : [Types.CanisterInfo] = [];
     stable var stableMetadata : [(Principal, (Int, Bool))] = [];
     stable var stableChildren : [(Principal, [Principal])] = [];
+    stable var stableTimers : [Types.CanisterInfo] = [];
     stable var previousParam : ?Types.InitParams = null;
 
     system func preupgrade() {
-        let (tree, metadata, children) = pool.share();
+        let (tree, metadata, children, timers) = pool.share();
         stablePool := tree;
         stableMetadata := metadata;
         stableChildren := children;
+        stableTimers := timers;
         previousParam := ?params;
     };
 
@@ -46,6 +48,9 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
             };
         };
         pool.unshare(stablePool, stableMetadata, stableChildren);
+        for (info in stableTimers.vals()) {
+            updateTimer(info);
+        }
     };
 
     public query func getInitParams() : async Types.InitParams {
@@ -86,8 +91,9 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
                     Cycles.add topUpCycles;
                     await IC.deposit_cycles cid;
                 };
-                // Lazily cleanup the reused canister
-                await IC.uninstall_code cid;
+                if (Option.isSome(status.module_hash)) {
+                    await IC.uninstall_code cid;
+                };
                 switch (status.status) {
                     case (#stopped or #stopping) {
                         await IC.start_canister cid;
@@ -144,12 +150,25 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
             await IC.install_code newArgs;
             stats := Logs.updateStats(stats, #install);
             switch (pool.refresh(info, profiling)) {
-                case (?newInfo) newInfo;
+                case (?newInfo) {
+                     updateTimer(newInfo);
+                     newInfo;
+                 };
                 case null { throw Error.reject "Cannot find canister" };
             };
         };
     };
 
+    func updateTimer(info: Types.CanisterInfo) {
+        func job() : async () {
+            pool.removeTimer(info.id);
+            // It is important that the timer job checks for the timestamp first.
+            // This prevents late-runner jobs from deleting newly installed code.
+            await removeCode(info);
+        };
+        pool.updateTimer(info, job);
+    };
+
     public func callForward(info : Types.CanisterInfo, function : Text, args : Blob) : async Blob {
         if (pool.find info) {
             await InternetComputer.call(info.id, function, args);

service/pool/Types.mo

@@ -9,6 +9,7 @@ import Array "mo:base/Array";
 import List "mo:base/List";
 import Option "mo:base/Option";
 import Int "mo:base/Int";
+import Timer "mo:base/Timer";
 
 module {
     public type InitParams = {
@@ -51,9 +52,11 @@ module {
     public class CanisterPool(size: Nat, ttl: Nat, max_family_tree_size: Nat) {
         var len = 0;
         var tree = Splay.Splay<CanisterInfo>(canisterInfoCompare);
+        // Metadata is a replicate of splay tree, which allows lookup without timestamp. Internal use only.
         var metadata = TrieMap.TrieMap<Principal, (Int, Bool)>(Principal.equal, Principal.hash);
         var childrens = TrieMap.TrieMap<Principal, List.List<Principal>>(Principal.equal, Principal.hash);
         var parents = TrieMap.TrieMap<Principal, Principal>(Principal.equal, Principal.hash);
+        let timers = TrieMap.TrieMap<Principal, Timer.TimerId>(Principal.equal, Principal.hash);
 
         public type NewId = { #newId; #reuse:CanisterInfo; #outOfCapacity:Nat };
 
@@ -123,6 +126,24 @@ module {
             return true;
         };
 
+        public func updateTimer(info: CanisterInfo, job : () -> async ()) {
+            let elapsed = Time.now() - info.timestamp;
+            let duration = if (elapsed > ttl) { 0 } else { Int.abs(ttl - elapsed) };
+            let tid = Timer.setTimer(#nanoseconds duration, job);
+            switch (timers.replace(info.id, tid)) {
+            case null {};
+            case (?old_id) {
+                     // The old job can still run when it has expired, but the future
+                     // just started to run. To be safe, the job needs to check for timestamp.
+                     Timer.cancelTimer(old_id);
+                 };
+            };
+        };
+
+        public func removeTimer(cid: Principal) {
+            timers.delete cid;
+        };
+
         private func notExpired(info: CanisterInfo, now: Int) : Bool = (info.timestamp > now - ttl);
 
         // Return a list of canister IDs from which to uninstall code
@@ -140,17 +161,22 @@ module {
             result
         };
 
-        public func share() : ([CanisterInfo], [(Principal, (Int, Bool))], [(Principal, [Principal])]) {
+        public func share() : ([CanisterInfo], [(Principal, (Int, Bool))], [(Principal, [Principal])], [CanisterInfo]) {
             let stableInfos = Iter.toArray(tree.entries());
             let stableMetadata = Iter.toArray(metadata.entries());
-            let stableChildrens = 
+            let stableChildren = 
                 Iter.toArray(
                     Iter.map<(Principal, List.List<Principal>), (Principal, [Principal])>(
                         childrens.entries(),
                         func((parent, children)) = (parent, List.toArray(children))
                     )
                 );
-            (stableInfos, stableMetadata, stableChildrens)
+            let stableTimers = Iter.toArray(
+              Iter.filter<CanisterInfo>(
+                tree.entries(),
+                func (info) = Option.isSome(timers.get(info.id))
+              ));
+            (stableInfos, stableMetadata, stableChildren, stableTimers)
         };
 
         public func unshare(stableInfos: [CanisterInfo], stableMetadata: [(Principal, (Int, Bool))], stableChildrens : [(Principal, [Principal])]) {

service/pool/tests/actor_class/test.sh

@@ -53,9 +53,7 @@ let args = record { arg = blob ""; wasm_module = parent; mode = variant { instal
 call S.installCode(c1, args, false);
 let c1 = c1.id;
 
-call c1.makeChild(0);
-fail call c1.makeChild(1);
-assert _ ~= "Canister has been uninstalled";
+fail call c1.makeChild(0);
 call S.getCanisterId(nonce);
 call S.getCanisterId(nonce);
 

service/pool/tests/canisterPool.test.sh

@@ -2,6 +2,20 @@
 load "prelude.sh";
 
 let wasm = file("../../../.dfx/local/canisters/backend/backend.wasm");
+let empty_wasm = blob "\00asm\01\00\00\00";
+
+let init = opt record {
+  cycles_per_canister = 105_000_000_000 : nat;
+  max_num_canisters = 2 : nat;
+  nonce_time_to_live = 1 : nat;
+  canister_time_to_live = 5_000_000_000 : nat;
+  max_family_tree_size = 5 : nat;
+};
+let S = install(wasm, init, null);
+let nonce = record { timestamp = 1 : int; nonce = 1 : nat };
+let CID = call S.getCanisterId(nonce);
+call S.installCode(CID, record { arg = blob ""; wasm_module = empty_wasm; mode = variant { install }; canister_id = CID.id }, false);
+metadata(CID.id, "module_hash");
 
 // Immediately expire
 let init = opt record {
@@ -13,7 +27,6 @@ let init = opt record {
 };
 let S = install(wasm, init, null);
 
-let nonce = record { timestamp = 1 : int; nonce = 1 : nat };
 let c1 = call S.getCanisterId(nonce);
 c1;
 let c2 = call S.getCanisterId(nonce);
@@ -64,3 +77,7 @@ call ic.provisional_top_up_canister(
   },
 );
 call S.getCanisterId(nonce);
+
+// Enough time has passed that the timer has removed the canister code
+fail metadata(CID.id, "module_hash");
+assert _ ~= "unknown";