Merge pull request #12933 from Budibase/cleanup-isolates

Clean up isolates when a request is finished.

packages/backend-core/src/context/types.ts

@@ -1,5 +1,4 @@
 import { IdentityContext, VM } from "@budibase/types"
-import { ExecutionTimeTracker } from "../timers"
 
 // keep this out of Budibase types, don't want to expose context info
 export type ContextMap = {
@@ -10,6 +9,6 @@ export type ContextMap = {
   isScim?: boolean
   automationId?: string
   isMigrating?: boolean
-  jsExecutionTracker?: ExecutionTimeTracker
   vm?: VM
+  cleanup?: (() => void | Promise<void>)[]
 }

packages/backend-core/src/timers/timers.ts

@@ -20,41 +20,3 @@ export function cleanup() {
   }
   intervals = []
 }
-
-export class ExecutionTimeoutError extends Error {
-  public readonly name = "ExecutionTimeoutError"
-}
-
-export class ExecutionTimeTracker {
-  static withLimit(limitMs: number) {
-    return new ExecutionTimeTracker(limitMs)
-  }
-
-  constructor(readonly limitMs: number) {}
-
-  private totalTimeMs = 0
-
-  track<T>(f: () => T): T {
-    this.checkLimit()
-    const start = process.hrtime.bigint()
-    try {
-      return f()
-    } finally {
-      const end = process.hrtime.bigint()
-      this.totalTimeMs += Number(end - start) / 1e6
-      this.checkLimit()
-    }
-  }
-
-  get elapsedMS() {
-    return this.totalTimeMs
-  }
-
-  checkLimit() {
-    if (this.totalTimeMs > this.limitMs) {
-      throw new ExecutionTimeoutError(
-        `Execution time limit of ${this.limitMs}ms exceeded: ${this.totalTimeMs}ms`
-      )
-    }
-  }
-}

packages/server/src/api/index.ts

@@ -1,6 +1,7 @@
 import Router from "@koa/router"
 import { auth, middleware, env as envCore } from "@budibase/backend-core"
 import currentApp from "../middleware/currentapp"
+import cleanup from "../middleware/cleanup"
 import zlib from "zlib"
 import { mainRoutes, staticRoutes, publicRoutes } from "./routes"
 import { middleware as pro } from "@budibase/pro"
@@ -62,6 +63,8 @@ if (apiEnabled()) {
     .use(auth.auditLog)
     // @ts-ignore
     .use(migrations)
+    // @ts-ignore
+    .use(cleanup)
 
   // authenticated routes
   for (let route of mainRoutes) {

packages/server/src/jsRunner/index.ts

@@ -8,6 +8,7 @@ import {
 import { context, logging } from "@budibase/backend-core"
 import tracer from "dd-trace"
 import { IsolatedVM } from "./vm"
+import type { VM } from "@budibase/types"
 
 export function init() {
   setJSRunner((js: string, ctx: Record<string, any>) => {
@@ -15,18 +16,23 @@ export function init() {
       try {
         const bbCtx = context.getCurrentContext()
 
-        const vm = bbCtx?.vm
-          ? bbCtx.vm
-          : new IsolatedVM({
-              memoryLimit: env.JS_RUNNER_MEMORY_LIMIT,
-              invocationTimeout: env.JS_PER_INVOCATION_TIMEOUT_MS,
-              isolateAccumulatedTimeout: env.JS_PER_REQUEST_TIMEOUT_MS,
-            }).withHelpers()
+        const vm =
+          bbCtx?.vm ||
+          new IsolatedVM({
+            memoryLimit: env.JS_RUNNER_MEMORY_LIMIT,
+            invocationTimeout: env.JS_PER_INVOCATION_TIMEOUT_MS,
+            isolateAccumulatedTimeout: env.JS_PER_REQUEST_TIMEOUT_MS,
+          }).withHelpers()
 
-        if (bbCtx) {
-          // If we have a context, we want to persist it to reuse the isolate
+        if (bbCtx && !bbCtx.vm) {
           bbCtx.vm = vm
+          bbCtx.cleanup = bbCtx.cleanup || []
+          bbCtx.cleanup.push(() => vm.close())
         }
+
+        // Because we can't pass functions into an Isolate, we remove them from
+        // the passed context and rely on the withHelpers() method to add them
+        // back in.
         const { helpers, ...rest } = ctx
         return vm.withContext(rest, () => vm.execute(js))
       } catch (error: any) {

packages/server/src/jsRunner/vm/isolated-vm.ts

@@ -195,6 +195,11 @@ export class IsolatedVM implements VM {
     return result[this.runResultKey]
   }
 
+  close(): void {
+    this.vm.release()
+    this.isolate.dispose()
+  }
+
   private registerCallbacks(functions: Record<string, any>) {
     const libId = crypto.randomUUID().replace(/-/g, "")
 

packages/server/src/middleware/cleanup.ts

@@ -0,0 +1,33 @@
+import { Ctx } from "@budibase/types"
+import { context } from "@budibase/backend-core"
+import { tracer } from "dd-trace"
+
+export default async (ctx: Ctx, next: any) => {
+  const resp = await next()
+
+  const current = context.getCurrentContext()
+  if (!current || !current.cleanup) {
+    return resp
+  }
+
+  let errors = []
+  for (let fn of current.cleanup) {
+    try {
+      await tracer.trace("cleanup", async span => {
+        await fn()
+      })
+    } catch (e) {
+      // We catch errors here to ensure we at least attempt to run all cleanup
+      // functions. We'll throw the first error we encounter after all cleanup
+      // functions have been run.
+      errors.push(e)
+    }
+  }
+  delete current.cleanup
+
+  if (errors.length > 0) {
+    throw errors[0]
+  }
+
+  return resp
+}

packages/types/src/sdk/vm.ts

@@ -1,4 +1,5 @@
 export interface VM {
   execute(code: string): any
   withContext<T>(context: Record<string, any>, executeWithContext: () => T): T
+  close(): void
 }