Test memory limits for composition worker

packages/services/schema/__tests__/stitching.spec.ts

@@ -0,0 +1,174 @@
+import { env } from 'node:process';
+import { Worker } from 'node:worker_threads';
+import type { CompositionEvent, CompositionResultEvent } from '@hive/schema/composition-worker';
+
+// Helper function to wrap Worker execution in a Promise
+function runWorkerWithLimits(
+  workerPath: string,
+  workerData: CompositionEvent['data'],
+  limits: { stackMb?: number; heapMb: number; codeRangeSizeMb?: number },
+) {
+  return new Promise((resolve, reject) => {
+    const worker = new Worker(workerPath, {
+      name: 'composition-worker-test',
+      resourceLimits: {
+        stackSizeMb: limits.stackMb,
+        maxOldGenerationSizeMb: limits.heapMb,
+        codeRangeSizeMb: limits.codeRangeSizeMb,
+      },
+      env: {
+        ENCRYPTION_SECRET: env.ENCRYPTION_SECRET,
+        REDIS_HOST: env.REDIS_HOST,
+        REDIS_PORT: env.REDIS_PORT ?? '6379',
+      },
+    });
+
+    worker.on('message', data => {
+      if (data.event === 'error') {
+        reject(data.err);
+      }
+      if (data.event === 'compositionResult') {
+        resolve(data);
+      }
+    });
+    worker.on('error', err => reject(err));
+    worker.on('exit', code => {
+      if (code !== 0) reject(new Error(`Worker stopped with exit code ${code}`));
+    });
+
+    worker.postMessage({
+      event: 'composition',
+      id: '1',
+      data: workerData,
+      taskId: '1',
+      requestId: '1',
+    });
+  });
+}
+
+describe('stitching', () => {
+  test('worker memory is exceeded when the schema size to worker size ratio drops below 150x', async () => {
+    await expect(
+      checkWorkerMemoryLimit({
+        schemaSizeInMb: 5,
+        schemaMbToWorkerMbRatio: 150,
+      }),
+    ).resolves.toBeTruthy();
+
+    await expect(
+      checkWorkerMemoryLimit({
+        schemaSizeInMb: 5,
+        schemaMbToWorkerMbRatio: 147,
+      }),
+    ).rejects.toThrow();
+  }, 20000);
+});
+
+/**
+ * Note that this limit is approximate because workers use some memory to hold code. And subtle differences in the schema can
+ * cause different memory usage.
+ */
+const checkWorkerMemoryLimit = async (opts: {
+  // The size of the schema to generate for this test.
+  // This is useful as an input in case there are edge cases that involve different sizes of schemas
+  schemaSizeInMb: number;
+  // This is a multiplier for how much heap memory is required in the worker to handle a schema
+  // of a certain size (maxMemoryMb)
+  schemaMbToWorkerMbRatio: number;
+}) => {
+  const charactersPerMb = 1048576;
+  const schema = generateGraphQLSchemaBySize(charactersPerMb * opts.schemaSizeInMb);
+  const workerFilePath = require.resolve('@hive/schema/dist/composition-worker-main.js');
+  const heavyInputData: CompositionEvent['data'] = {
+    type: 'stitching',
+    args: {
+      schemas: [
+        {
+          raw: schema,
+          source: 'test',
+        },
+      ],
+    },
+  };
+
+  try {
+    const result = (await runWorkerWithLimits(workerFilePath, heavyInputData, {
+      heapMb: opts.schemaSizeInMb * opts.schemaMbToWorkerMbRatio,
+    })) as CompositionResultEvent;
+
+    // Assert that it completed without hitting engine caps
+    expect(result.event).toBe('compositionResult');
+    expect(typeof result.data.result.sdl).toBe('string');
+    return true;
+  } catch (err: any) {
+    // If the thread runs completely out of memory, Node may forcefully crash the thread
+    throw new Error(`The thread crashed due to exceeding resource boundaries: ${err.message}`);
+  }
+};
+
+/**
+ * Creates an approximation of a graphql schema that is a specific number of characters in length.
+ */
+function generateGraphQLSchemaBySize(targetChars: number): string {
+  const MAX_FIELDS_PER_TYPE = 40;
+  const SCALAR_TYPES = ['String', 'Int', 'Boolean', 'Float', 'ID'];
+
+  let schema = 'type Query {\n  root: Type0!\n}\n\n';
+
+  const baselineMinLength = schema.length;
+  if (targetChars < baselineMinLength) {
+    throw new Error(
+      `Target size is too small. Minimum required length is ${baselineMinLength} characters.`,
+    );
+  }
+
+  let typeIndex = 0;
+  let keepGenerating = true;
+
+  while (keepGenerating) {
+    let currentType = `type Type${typeIndex} {\n  id: ID!\n`;
+
+    // Check if adding this type shell would overshoot the limit
+    if (schema.length + currentType.length + 2 >= targetChars) {
+      break;
+    }
+
+    // Populate the type with various scalar fields
+    for (let fieldIndex = 0; fieldIndex < MAX_FIELDS_PER_TYPE; fieldIndex++) {
+      const scalarType = SCALAR_TYPES[fieldIndex % SCALAR_TYPES.length];
+      const nextField = `  field${fieldIndex}: ${scalarType}!\n`;
+
+      // Stop appending fields if it would overshoot the target limit once closed
+      if (schema.length + currentType.length + nextField.length + 2 > targetChars) {
+        keepGenerating = false;
+        break;
+      }
+      currentType += nextField;
+    }
+
+    // Add an edge relationship to the *next* type to create graph depth
+    const relationField = `  related: Type${typeIndex + 1}\n`;
+    if (
+      keepGenerating &&
+      schema.length + currentType.length + relationField.length + 2 <= targetChars
+    ) {
+      currentType += relationField;
+    } else {
+      keepGenerating = false;
+    }
+
+    currentType += '}\n\n';
+    schema += currentType;
+    typeIndex++;
+  }
+
+  // Exactly pad any remaining bytes using GraphQL comments
+  if (schema.length < targetChars) {
+    schema += '#';
+    while (schema.length < targetChars) {
+      schema += 'x';
+    }
+  }
+
+  return schema;
+}

packages/services/schema/src/composition-scheduler.ts

@@ -115,6 +115,11 @@ export class CompositionScheduler {
         extra: {
           requestId: workerState?.args.requestId ?? '',
           compositionType: workerState?.args.data.type,
+          schemaCount: workerState?.args.data.args.schemas.length,
+          schemaTotalSize: workerState?.args.data.args.schemas.reduce(
+            (size, s) => size + s.raw.length,
+            0,
+          ),
         },
       });
       recreate(error);