feat: ensure max state memory size (#1458)

* feat: ensure max state memory

* tests: add tests for state size limit

* tests: fix flaky test on dataclip loading

* chore: remove log

* chore: update changeset

* versions

* reduce size of state objects

Down from 20% of allow memory to 15%

* log state sizes

* typo

---------

Co-authored-by: Joe Clark <jclark@openfn.org>

Author: doc-han

integration-tests/worker/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @openfn/integration-tests-worker
 
+## 1.1.0
+
+### Minor Changes
+
+- 8ebc086: Add support for state memory limit via worker
+
+### Patch Changes
+
+- Updated dependencies [8ebc086]
+  - @openfn/engine-multi@1.12.0
+  - @openfn/ws-worker@1.27.0
+  - @openfn/lightning-mock@2.4.20
+
 ## 1.0.96
 
 ### Patch Changes

integration-tests/worker/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@openfn/integration-tests-worker",
   "private": true,
-  "version": "1.0.96",
+  "version": "1.1.0",
   "description": "Lightning WOrker integration tests",
   "author": "Open Function Group <admin@openfn.org>",
   "license": "ISC",

integration-tests/worker/test/exit-reasons.test.ts

@@ -226,6 +226,92 @@ test.serial('kill: oom (large, kill vm)', async (t) => {
   t.is(error_message, 'Run exceeded maximum memory usage');
 });
 
+test.serial('kill: state exceeds the configured state limit', async (t) => {
+  const attempt = {
+    id: crypto.randomUUID(),
+    jobs: [
+      {
+        adaptor: '@openfn/language-common@latest',
+        // ~2mb string for a 1mb limit
+        body: `fn((s) => {
+          s.data = new Array(2 * 1024 * 1024).fill('a').join('');
+          return s;
+        })`,
+      },
+    ],
+    options: {
+      state_limit_mb: 1,
+    },
+  };
+
+  const result = await run(attempt);
+
+  const { reason, error_type, error_message } = result;
+  t.is(reason, 'kill');
+  t.is(error_type, 'StateTooLargeError');
+  t.regex(error_message, /State exceeds the limit of 1mb/);
+});
+
+test.serial(
+  'kill: state limit is enforced between jobs (downstream job does not run)',
+  async (t) => {
+    const jobOne = {
+      id: crypto.randomUUID(),
+      adaptor: '@openfn/language-common@latest',
+      // ~2mb state, over the 1mb limit set below
+      body: `fn((s) => {
+        s.data = new Array(2 * 1024 * 1024).fill('a').join('');
+        return s;
+      })`,
+    };
+
+    // not expected to run because the first job is expected to trigger state size crash
+    const jobTwo = {
+      id: crypto.randomUUID(),
+      adaptor: '@openfn/language-common@latest',
+      body: `fn(() => ({ data: 'ok' }))`,
+    };
+
+    const attempt = {
+      id: crypto.randomUUID(),
+      jobs: [jobOne, jobTwo],
+      edges: [
+        {
+          id: crypto.randomUUID(),
+          source_job_id: jobOne.id,
+          target_job_id: jobTwo.id,
+          condition: 'always',
+        },
+      ],
+      options: {
+        state_limit_mb: 1,
+      },
+    };
+
+    const startedJobs: string[] = [];
+    const unsubscribe = lightning.onSocketEvent(
+      'step:start',
+      attempt.id,
+      (evt) => {
+        if (evt.runId === attempt.id) {
+          startedJobs.push(evt.payload.job_id);
+        }
+      },
+      false
+    );
+
+    const result = await run(attempt);
+    unsubscribe();
+
+    const { reason, error_type, error_message } = result;
+    t.is(reason, 'kill');
+    t.is(error_type, 'StateTooLargeError');
+    t.regex(error_message, /State exceeds the limit of 1mb/);
+
+    t.deepEqual(startedJobs, [jobOne.id]);
+  }
+);
+
 test.serial('crash: process.exit() triggered by postgres', async (t) => {
   const attempt = {
     id: crypto.randomUUID(),

packages/engine-multi/CHANGELOG.md

@@ -1,5 +1,17 @@
 # engine-multi
 
+## 1.12.0
+
+### Minor Changes
+
+- 8ebc086: Add support for state memory limit via worker
+- reduce default size of state objects
+
+### Patch Changes
+
+- Updated dependencies [8ebc086]
+  - @openfn/lexicon@2.3.0
+
 ## 1.11.4
 
 ### Patch Changes

packages/engine-multi/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openfn/engine-multi",
-  "version": "1.11.4",
+  "version": "1.12.0",
   "description": "Multi-process runtime engine",
   "main": "dist/index.js",
   "type": "module",

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

@@ -50,7 +50,7 @@ const execute = async (context: ExecutionContext) => {
       // This must be fairly high to prevent crashes
       stateLimitMb:
         options.stateLimitMb ??
-        Math.max((options.memoryLimitMb ?? 1000) * 0.25),
+        Math.max(50, options.memoryLimitMb ?? 1000 * 0.15),
     } as RunOptions;
 
     logger.debug(

packages/engine-multi/src/engine.ts

@@ -74,6 +74,7 @@ export type EngineOptions = {
   logger: Logger;
   maxWorkers?: number;
   memoryLimitMb?: number;
+  stateLimitMb?: number;
   payloadLimitMb?: number;
   logPayloadLimitMb?: number;
   repoDir: string;
@@ -169,7 +170,7 @@ const createEngine = async (
       callWorker,
       options: {
         ...options,
-        stateLimitMb: opts.stateLimitMb,
+        stateLimitMb: opts.stateLimitMb ?? options.stateLimitMb,
         sanitize: opts.sanitize,
         resolvers: opts.resolvers,
         runTimeoutMs: opts.runTimeoutMs ?? defaultTimeout,

packages/lexicon/CHANGELOG.md

@@ -1,5 +1,11 @@
 # lexicon
 
+## 2.3.0
+
+### Minor Changes
+
+- 8ebc086: Add support for state memory limit via worker
+
 ## 2.2.1
 
 ### Patch Changes

packages/lexicon/lightning.d.ts

@@ -59,6 +59,7 @@ export type LightningPlanOptions = {
   output_dataclips?: boolean;
 
   run_memory_limit_mb?: number;
+  state_limit_mb?: number;
   payload_limit_mb?: number;
   log_payload_limit_mb?: number;
   job_log_level?: LogLevel;

packages/lexicon/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openfn/lexicon",
-  "version": "2.2.1",
+  "version": "2.3.0",
   "description": "Central repo of names and type definitions",
   "author": "Open Function Group <admin@openfn.org>",
   "license": "ISC",