Limit state size in runtime (#1209)

* runtime: add function to ensure that state size does not exceed a particular limit

* add new stateLimit_mb option

* engine: support state limit mb as an option

I don't think this is the right option in the right place - I'll check into that later

* refactor rutime option

* tidying up

* fix failing test in mock

* more generous timeout on flaky test

* add debug log

* changeset

* tweak memory limit

Drop the payload size thing - runtime memory will always be higher

* refine options

* runtime: type safety

Author: josephjclark

.changeset/ready-items-shout.md

@@ -0,0 +1,9 @@
+---
+'@openfn/engine-multi': minor
+'@openfn/ws-worker': minor
+'@openfn/runtime': minor
+---
+
+Measure the size of state objects at the end of each step, and throw if they exceed a limit
+
+In the Worker, this limit is set to 25% of the available runtime memory.

packages/engine-multi/src/api/execute.ts

@@ -46,8 +46,17 @@ const execute = async (context: ExecutionContext) => {
       repoDir: options.repoDir,
       profile: context.options.profile,
       profilePollInteval: context.options.profilePollInterval,
+      // work out the max size of the state object at the end of each step
+      // This must be fairly high to prevent crashes
+      stateLimitMb:
+        options.stateLimitMb ??
+        Math.max((options.memoryLimitMb ?? 1000) * 0.25),
     } as RunOptions;
 
+    logger.debug(
+      `${state.plan.id} setting runtime state limit to ${runOptions.stateLimitMb}mb`
+    );
+
     // Construct the payload limits object
     const payloadLimits: PayloadLimits = {};
     if (options.payloadLimitMb !== undefined) {

packages/engine-multi/src/engine.ts

@@ -167,6 +167,7 @@ const createEngine = async (
       callWorker,
       options: {
         ...options,
+        stateLimitMb: opts.stateLimitMb,
         sanitize: opts.sanitize,
         resolvers: opts.resolvers,
         runTimeoutMs: opts.runTimeoutMs ?? defaultTimeout,

packages/engine-multi/src/types.ts

@@ -45,6 +45,7 @@ export type ExecutionContextConstructor = {
 };
 
 export type ExecuteOptions = {
+  stateLimitMb?: number;
   payloadLimitMb?: number;
   logPayloadLimitMb?: number;
   memoryLimitMb?: number;
@@ -72,7 +73,7 @@ export interface RuntimeEngine {
   execute(
     plan: ExecutionPlan,
     input: State,
-    options?: Partial<EngineOptions>
+    options?: ExecuteOptions
   ): Pick<EventEmitter, 'on' | 'off' | 'once'>;
 
   destroy(instant?: boolean): Promise<void>;

packages/engine-multi/src/worker/thread/run.ts

@@ -20,6 +20,7 @@ export type RunOptions = {
   jobLogLevel?: LogLevel;
   profile?: boolean;
   profilePollInterval?: number;
+  stateLimitMb?: number;
 };
 
 const eventMap = {
@@ -39,6 +40,7 @@ register({
       jobLogLevel,
       profile,
       profilePollInterval,
+      stateLimitMb,
     } = runOptions;
     const { logger, jobLogger, adaptorLogger } = createLoggers(
       plan.id!,
@@ -54,9 +56,8 @@ register({
     console = adaptorLogger;
 
     // Leave console.debug for local debugging
-    // This goes to stdout but not the adapator logger
+    // This goes to stdout but not the adaptor logger
     console.debug = debug;
-
     // TODO I would like to pull these options out of here
     const options = {
       // disable the runtime's own timeout
@@ -72,6 +73,7 @@ register({
       profile,
       profilePollInterval,
       statePropsToRemove,
+      stateLimitMb,
       callbacks: {
         // TODO: this won't actually work across the worker boundary
         // For now I am preloading credentials

packages/engine-multi/test/errors.test.ts

@@ -165,6 +165,35 @@ test.serial.skip('vm oom error', (t) => {
   });
 });
 
+test.serial('state object too big', async (t) => {
+  return new Promise((done) => {
+    const plan = {
+      id: 'x',
+      workflow: {
+        steps: [
+          {
+            expression: `export default [(s) => {
+              s.data = new Array(1024 * 1024).fill("x").join("");
+              return s;
+            }]`,
+          },
+        ],
+      },
+    };
+
+    const options = {
+      stateLimitMb: 0.1,
+    };
+
+    engine.execute(plan, {}, options).on(WORKFLOW_ERROR, (evt) => {
+      t.is(evt.type, 'StateTooLargeError');
+      t.is(evt.severity, 'kill');
+      t.is(evt.message, 'State exceeds the limit of 0.1mb');
+      done();
+    });
+  });
+});
+
 test.serial('execution error from async code', (t) => {
   return new Promise((done) => {
     const plan = {

packages/runtime/package.json

@@ -48,6 +48,7 @@
   "dependencies": {
     "@openfn/logger": "workspace:*",
     "fast-safe-stringify": "^2.1.1",
+    "json-stream-stringify": "^3.1.6",
     "semver": "^7.7.2",
     "source-map": "^0.7.6"
   }

packages/runtime/src/errors.ts

@@ -91,6 +91,15 @@ export const extractStackTrace = (e: Error) => {
   }
 };
 
+export class StateTooLargeError extends Error {
+  name = 'StateTooLargeError';
+  severity = 'kill';
+  constructor(limit_mb: number) {
+    super();
+    this.message = `State exceeds the limit of ${limit_mb}mb`;
+  }
+}
+
 // Abstract error supertype
 export class RTError extends Error {
   source = 'runtime';

packages/runtime/src/execute/step.ts

@@ -20,6 +20,7 @@ import {
 import { isNullState } from '../util/null-state';
 import sourcemapErrors from '../util/sourcemap-errors';
 import createProfiler from '../util/profile-memory';
+import ensureStateSize from '../util/ensure-state-size';
 
 const loadCredentials = async (
   job: Job,
@@ -79,13 +80,21 @@ const calculateNext = (job: CompiledStep, result: any, logger: Logger) => {
 
 // TODO this is suboptimal and may be slow on large objects
 // (especially as the result get stringified again downstream)
-const prepareFinalState = (
+const prepareFinalState = async (
   state: any,
   logger: Logger,
-  statePropsToRemove?: string[]
+  statePropsToRemove?: string[],
+  stateLimit_mb?: number
 ) => {
   if (isNullState(state)) return undefined;
   if (state) {
+    try {
+      await ensureStateSize(state, stateLimit_mb);
+    } catch (e) {
+      logger.error('Critical error processing state:');
+      throw e;
+    }
+
     if (!statePropsToRemove) {
       // As a strict default, remove the configuration key
       // tbh this should happen higher up in the stack but it causes havoc in unit testing
@@ -190,7 +199,12 @@ const executeStep = async (
         const duration = logger.timer(timerId);
         logger.error(`${jobName} aborted with error (${duration})`);
 
-        state = prepareFinalState(state, logger, ctx.opts.statePropsToRemove);
+        state = await prepareFinalState(
+          state,
+          logger,
+          ctx.opts.statePropsToRemove,
+          ctx.opts.stateLimitMb
+        );
         // Whatever the final state was, save that as the initial state to the next thing
         result = state;
 
@@ -219,7 +233,12 @@ const executeStep = async (
     if (!didError) {
       const humanDuration = logger.timer(timerId);
       logger.success(`${jobName} completed in ${humanDuration}`);
-      result = prepareFinalState(result, logger, ctx.opts.statePropsToRemove);
+      result = await prepareFinalState(
+        result,
+        logger,
+        ctx.opts.statePropsToRemove,
+        ctx.opts.stateLimitMb
+      );
 
       // Take a memory snapshot
       // IMPORTANT: this runs _after_ the state object has been serialized

packages/runtime/src/runtime.ts

@@ -38,6 +38,8 @@ export type Options = {
 
   /** Optional name for the expression (if passed a string) */
   defaultStepId?: string;
+
+  stateLimitMb?: number;
 };
 
 type RawOptions = Omit<Options, 'linker'> & {