diff --git a/service/pool/Main.mo b/service/pool/Main.mo index da61dc1d..3cad1468 100644 --- a/service/pool/Main.mo +++ b/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); diff --git a/service/pool/Types.mo b/service/pool/Types.mo index 54fbf070..44008709 100644 --- a/service/pool/Types.mo +++ b/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(canisterInfoCompare); + // Metadata is a replicate of splay tree, which allows lookup without timestamp. Internal use only. var metadata = TrieMap.TrieMap(Principal.equal, Principal.hash); var childrens = TrieMap.TrieMap>(Principal.equal, Principal.hash); var parents = TrieMap.TrieMap(Principal.equal, Principal.hash); + let timers = TrieMap.TrieMap(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])>( childrens.entries(), func((parent, children)) = (parent, List.toArray(children)) ) ); - (stableInfos, stableMetadata, stableChildrens) + let stableTimers = Iter.toArray( + Iter.filter( + 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])]) { diff --git a/service/pool/tests/actor_class/test.sh b/service/pool/tests/actor_class/test.sh index f81f3c29..70cd0f8a 100644 --- a/service/pool/tests/actor_class/test.sh +++ b/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); diff --git a/service/pool/tests/canisterPool.test.sh b/service/pool/tests/canisterPool.test.sh index 9ca6993a..fc0cba60 100644 --- a/service/pool/tests/canisterPool.test.sh +++ b/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";