Fix memory leaks in clear_timer

demergent-labs · Sep 11, 2023 · 0690ef3 · 0690ef3
1 parent a6fb236
commit 0690ef3
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 39 deletions.
diff --git a/src/lib_new/globals.ts b/src/lib_new/globals.ts
@@ -1,17 +1,25 @@
 import { ic } from './ic';
 import { Buffer } from 'buffer';
 
-declare var globalThis: {
-    TextDecoder: any;
-    TextEncoder: any;
+export declare var globalThis: {
     _azleCandidInitParams: any[];
     _azleCandidMethods: any[];
     _azleCandidTypes: any[];
+    Buffer: BufferConstructor;
     console: any;
+    crypto: {
+        getRandomValues: () => Uint8Array;
+    };
+    icTimers: {
+        [key: string]: string;
+    };
+    TextDecoder: any;
+    TextEncoder: any;
 };
 
 globalThis.TextDecoder = require('text-encoding').TextDecoder;
 globalThis.TextEncoder = require('text-encoding').TextEncoder;
+globalThis.icTimers ||= {};
 
 globalThis.console = {
     ...globalThis.console,

diff --git a/src/lib_new/ic.ts b/src/lib_new/ic.ts
@@ -582,12 +582,17 @@ export const ic: Ic = globalThis._azleIc
               return IDL.decode([IDL.Nat64], canisterVersionCandidBytes)[0];
           },
           clearTimer: (timerId: nat64) => {
-              // TODO: We need to delete the callback from the global scope as well
-              const timerIdCandidBytes = new Uint8Array(
-                  IDL.encode([IDL.Nat64], [timerId])
-              ).buffer;
+              const encode = (value: nat64) => {
+                  return new Uint8Array(IDL.encode([IDL.Nat64], [value]))
+                      .buffer;
+              };
+
+              globalThis._azleIc.clearTimer(encode(timerId));
+
+              const timerCallbackId = globalThis.icTimers[timerId.toString()];
 
-              return globalThis._azleIc.clearTimer(timerIdCandidBytes);
+              delete globalThis.icTimers[timerId.toString()];
+              delete globalThis[timerCallbackId];
           },
           dataCertificate: () => {
               const rawRustValue: ArrayBuffer | undefined =
@@ -711,58 +716,63 @@ export const ic: Ic = globalThis._azleIc
               return globalThis._azleIc.setCertifiedData(dataBytes);
           },
           setTimer: (delay: nat64, callback: () => void | Promise<void>) => {
+              const encode = (value: nat64) => {
+                  return new Uint8Array(IDL.encode([IDL.Nat64], [value]))
+                      .buffer;
+              };
+
+              const decode = (value: ArrayBufferLike) => {
+                  return BigInt(IDL.decode([IDL.Nat64], value)[0] as number);
+              };
+
               const timerCallbackId = `_timer_${v4()}`;
 
+              const timerId = decode(
+                  globalThis._azleIc.setTimer(encode(delay), timerCallbackId)
+              );
+
+              globalThis.icTimers[timerId.toString()] = timerCallbackId;
+
               globalThis[timerCallbackId] = () => {
                   try {
                       callback();
                   } finally {
+                      delete globalThis.icTimers[timerId.toString()];
                       delete globalThis[timerCallbackId];
                   }
               };
 
-              const delayCandidBytes = new Uint8Array(
-                  IDL.encode([IDL.Nat64], [delay])
-              ).buffer;
-
-              try {
-                  const timerIdCandidBytes = globalThis._azleIc.setTimer(
-                      delayCandidBytes,
-                      timerCallbackId
-                  );
-
-                  return IDL.decode([IDL.Nat64], timerIdCandidBytes)[0];
-              } catch (error) {
-                  delete globalThis[timerCallbackId];
-                  throw error;
-              }
+              return timerId;
           },
           setTimerInterval: (
               interval: nat64,
               callback: () => void | Promise<void>
           ) => {
+              const encode = (value: nat64) => {
+                  return new Uint8Array(IDL.encode([IDL.Nat64], [value]))
+                      .buffer;
+              };
+
+              const decode = (value: ArrayBufferLike) => {
+                  return BigInt(IDL.decode([IDL.Nat64], value)[0] as number);
+              };
+
               const timerCallbackId = `_interval_timer_${v4()}`;
 
+              const timerId = decode(
+                  globalThis._azleIc.setTimerInterval(
+                      encode(interval),
+                      timerCallbackId
+                  )
+              );
+
+              globalThis.icTimers[timerId.toString()] = timerCallbackId;
+
               // We don't delete this even if the callback throws because
               // it still needs to be here for the next tick
               globalThis[timerCallbackId] = callback;
 
-              const intervalCandidBytes = new Uint8Array(
-                  IDL.encode([IDL.Nat64], [interval])
-              ).buffer;
-
-              try {
-                  const timerIdCandidBytes =
-                      globalThis._azleIc.setTimerInterval(
-                          intervalCandidBytes,
-                          timerCallbackId
-                      );
-
-                  return IDL.decode([IDL.Nat64], timerIdCandidBytes)[0];
-              } catch (error) {
-                  delete globalThis[timerCallbackId];
-                  throw error;
-              }
+              return timerId;
           },
           stableBytes: () => {
               return new Uint8Array(globalThis._azleIc.stableBytes());