Make timeout errors more understandable inside workers

src/worker/vm/imports.ts

@@ -1,4 +1,4 @@
-import { BigQuery } from '@google-cloud/bigquery'
+import * as bigquery from '@google-cloud/bigquery'
 import * as contrib from '@posthog/plugin-contrib'
 import * as scaffold from '@posthog/plugin-scaffold'
 import * as AWS from 'aws-sdk'
@@ -19,7 +19,7 @@ export const imports = {
     'generic-pool': genericPool,
     'node-fetch': fetch,
     'snowflake-sdk': snowflake,
-    '@google-cloud/bigquery': { BigQuery },
+    '@google-cloud/bigquery': bigquery,
     '@posthog/plugin-scaffold': scaffold,
     '@posthog/plugin-contrib': contrib,
     'aws-sdk': AWS,

src/worker/vm/transforms/promise-timeout.ts

@@ -31,7 +31,12 @@ export const promiseTimeout: PluginGen = () => ({ types: t }) => ({
             exit(path: any) {
                 const { node } = path
                 if (node && !node[REPLACED]) {
-                    const newAwait = t.awaitExpression(t.callExpression(t.identifier('__asyncGuard'), [node.argument]))
+                    const newAwait = t.awaitExpression(
+                        t.callExpression(t.identifier('__asyncGuard'), [
+                            node.argument,
+                            node.argument.callee || node.argument,
+                        ])
+                    )
                     ;(newAwait as any)[REPLACED] = true
                     path.replaceWith(newAwait)
                 }

src/worker/vm/vm.ts

@@ -14,6 +14,16 @@ import { imports } from './imports'
 import { transformCode } from './transforms'
 import { upgradeExportEvents } from './upgrades/export-events'
 
+export class TimeoutError extends Error {
+    name = 'TimeoutError'
+    caller?: string = undefined
+
+    constructor(message: string, caller?: string) {
+        super(message)
+        this.caller = caller
+    }
+}
+
 export async function createPluginConfigVM(
     hub: Hub,
     pluginConfig: PluginConfig, // NB! might have team_id = 0
@@ -46,7 +56,7 @@ export async function createPluginConfigVM(
     // Creating this outside the vm (so not in a babel plugin for example)
     // because `setTimeout` is not available inside the vm... and we don't want to
     // make it available for now, as it makes it easier to create malicious code
-    const asyncGuard = async (promise: () => Promise<any>) => {
+    const asyncGuard = async (promise: Promise<any>, name?: string) => {
         const timeout = hub.TASK_TIMEOUT
         return await Promise.race([
             promise,
@@ -55,7 +65,7 @@ export async function createPluginConfigVM(
                     const message = `Script execution timed out after promise waited for ${timeout} second${
                         timeout === 1 ? '' : 's'
                     }`
-                    reject(new Error(message))
+                    reject(new TimeoutError(message, `${name}`))
                 }, timeout * 1000)
             ),
         ])
@@ -122,8 +132,8 @@ export async function createPluginConfigVM(
 
             // export helpers
             function __getExported (key) { return exportDestinations.find(a => a[key])?.[key] };
-            function __asyncFunctionGuard (func) {
-                return func ? function __innerAsyncGuard${pluginConfigIdentifier}(...args) { return __asyncGuard(func(...args)) } : func
+            function __asyncFunctionGuard (func, name) {
+                return func ? function __innerAsyncGuard${pluginConfigIdentifier}(...args) { return __asyncGuard(func(...args), name) } : func
             };
 
             // inject the meta object + shareable 'global' to the end of each exported function
@@ -149,12 +159,12 @@ export async function createPluginConfigVM(
 
             // export various functions
             const __methods = {
-                setupPlugin: __asyncFunctionGuard(__bindMeta('setupPlugin')),
-                teardownPlugin: __asyncFunctionGuard(__bindMeta('teardownPlugin')),
-                exportEvents: __asyncFunctionGuard(__bindMeta('exportEvents')),
-                onEvent: __asyncFunctionGuard(__bindMeta('onEvent')),
-                onSnapshot: __asyncFunctionGuard(__bindMeta('onSnapshot')),
-                processEvent: __asyncFunctionGuard(__bindMeta('processEvent')),
+                setupPlugin: __asyncFunctionGuard(__bindMeta('setupPlugin'), 'setupPlugin'),
+                teardownPlugin: __asyncFunctionGuard(__bindMeta('teardownPlugin'), 'teardownPlugin'),
+                exportEvents: __asyncFunctionGuard(__bindMeta('exportEvents'), 'exportEvents'),
+                onEvent: __asyncFunctionGuard(__bindMeta('onEvent'), 'onEvent'),
+                onSnapshot: __asyncFunctionGuard(__bindMeta('onSnapshot'), 'onSnapshot'),
+                processEvent: __asyncFunctionGuard(__bindMeta('processEvent'), 'processEvent'),
             };
 
             const __tasks = {

tests/postgres/vm.timeout.test.ts

@@ -1,6 +1,7 @@
 import { Hub } from '../../src/types'
 import { createHub } from '../../src/utils/db/hub'
-import { createPluginConfigVM } from '../../src/worker/vm/vm'
+import { delay } from '../../src/utils/utils'
+import { createPluginConfigVM, TimeoutError } from '../../src/worker/vm/vm'
 import { pluginConfig39 } from '../helpers/plugins'
 import { resetTestDatabase } from '../helpers/sql'
 
@@ -224,14 +225,18 @@ describe('vm timeout tests', () => {
         const vm = await createPluginConfigVM(hub, pluginConfig39, indexJs)
         const date = new Date()
         let errorMessage = undefined
+        let caller = undefined
         try {
             await vm.methods.processEvent!({ ...defaultEvent })
         } catch (e) {
+            expect(e).toBeInstanceOf(TimeoutError)
             errorMessage = e.message
+            caller = e.caller
         }
         expect(new Date().valueOf() - date.valueOf()).toBeGreaterThanOrEqual(1000)
         expect(new Date().valueOf() - date.valueOf()).toBeLessThan(4000)
         expect(errorMessage!).toEqual('Script execution timed out after promise waited for 1 second')
+        expect(caller).toEqual('processEvent')
     })
 
     test('small promises and overriding async guard', async () => {

tests/transforms.test.ts

@@ -30,7 +30,47 @@ describe('transformCode', () => {
             "use strict";
 
             async function x() {
-              await __asyncGuard(console.log());
+              await __asyncGuard(console.log(), console.log);
+            }
+        `)
+    })
+
+    it('attaches caller information to awaits', () => {
+        const rawCode = code`
+            async function x() {
+              await anotherAsyncFunction('arg1', 'arg2')
+            }
+        `
+
+        const transformedCode = transformCode(rawCode, hub)
+
+        expect(transformedCode).toStrictEqual(code`
+            "use strict";
+
+            async function x() {
+              await __asyncGuard(anotherAsyncFunction('arg1', 'arg2'), anotherAsyncFunction);
+            }
+        `)
+    })
+
+    it('attaches caller information to awaits for anonymous functions', () => {
+        const rawCode = code`
+            async function x() {
+              await (async () => {console.log()})
+            }
+        `
+
+        const transformedCode = transformCode(rawCode, hub)
+
+        expect(transformedCode).toStrictEqual(code`
+            "use strict";
+
+            async function x() {
+              await __asyncGuard(async () => {
+                console.log();
+              }, async () => {
+                console.log();
+              });
             }
         `)
     })