diff --git a/packages/services/schema/__tests__/stitching.spec.ts b/packages/services/schema/__tests__/stitching.spec.ts new file mode 100644 index 00000000000..0c316ad5801 --- /dev/null +++ b/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; +} diff --git a/packages/services/schema/src/composition-scheduler.ts b/packages/services/schema/src/composition-scheduler.ts index b2ba433a301..67f43ea5acd 100644 --- a/packages/services/schema/src/composition-scheduler.ts +++ b/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);