add better error handling to heartbeat, inspect_message, and pre_upgr…

…ade for the class-based syntax
demergent-labs · Jul 11, 2024 · 5415d8f · 5415d8f
1 parent 3e555d8
commit 5415d8f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/src/lib/stable/heartbeat.ts b/src/lib/stable/heartbeat.ts
@@ -1,5 +1,7 @@
+import { handleUncaughtError } from './error';
+
 export function heartbeat<T>(
-    _target: object,
+    target: object,
     propertyKey: string | symbol,
     descriptor: TypedPropertyDescriptor<T>
 ): TypedPropertyDescriptor<T> | void {
@@ -11,7 +13,13 @@ export function heartbeat<T>(
     };
 
     globalThis._azleCanisterMethods.callbacks[index.toString()] =
-        descriptor.value as any;
+        async (): Promise<void> => {
+            try {
+                await (descriptor.value as any).bind(target)();
+            } catch (error) {
+                handleUncaughtError(error);
+            }
+        };
 
     return descriptor;
 }

diff --git a/src/lib/stable/inspect_message.ts b/src/lib/stable/inspect_message.ts
@@ -4,8 +4,10 @@
 // TODO but it will break all of your other methods
 // TODO so do we just leave params out?
 
+import { handleUncaughtError } from './error';
+
 export function inspectMessage<T>(
-    _target: object,
+    target: object,
     propertyKey: string | symbol,
     descriptor: TypedPropertyDescriptor<T>
 ): TypedPropertyDescriptor<T> | void {
@@ -16,8 +18,13 @@ export function inspectMessage<T>(
         index
     };
 
-    globalThis._azleCanisterMethods.callbacks[index.toString()] =
-        descriptor.value as any;
+    globalThis._azleCanisterMethods.callbacks[index.toString()] = (): void => {
+        try {
+            (descriptor.value as any).bind(target)();
+        } catch (error) {
+            handleUncaughtError(error);
+        }
+    };
 
     return descriptor;
 }
diff --git a/src/lib/stable/pre_upgrade.ts b/src/lib/stable/pre_upgrade.ts
@@ -1,5 +1,7 @@
+import { handleUncaughtError } from './error';
+
 export function preUpgrade<T>(
-    _target: object,
+    target: object,
     propertyKey: string | symbol,
     descriptor: TypedPropertyDescriptor<T>
 ): TypedPropertyDescriptor<T> | void {
@@ -10,8 +12,13 @@ export function preUpgrade<T>(
         index
     };
 
-    globalThis._azleCanisterMethods.callbacks[index.toString()] =
-        descriptor.value as any;
+    globalThis._azleCanisterMethods.callbacks[index.toString()] = (): void => {
+        try {
+            (descriptor.value as any).bind(target)();
+        } catch (error) {
+            handleUncaughtError(error);
+        }
+    };
 
     return descriptor;
 }