refactor(ajv): rewrite validateSchema to leverage internal ajv cache

jujaga · jujaga · commit 88170f52ba50 · 2026-03-09T10:58:44.000-07:00
This change slightly improves memory performance by generating fewer AJV
objects and using smarter logic to funnel repeated schema to the same
pre-compiled definition.

Signed-off-by: Jeremy Ho &lt;jujaga@gmail.com&gt;
diff --git a/src/validators/validator.ts b/src/validators/validator.ts
@@ -1,15 +1,14 @@
-import { Ajv } from 'ajv';
-
 import { integrityValidators } from './integrity/index.ts';
-import { createAjvInstance, getPiesSchemaUri, loadSchema, pies } from './schema/index.ts';
+import { createAjvInstance, ensureSchemaId, getPiesSchemaUri, loadSchema, pies } from './schema/index.ts';
 import { getLogger } from '../utils/index.ts';
 
 import type { AnySchemaObject, AnyValidateFunction, ErrorObject } from 'ajv/dist/core.js';
 import type { IntegrityDictionary, IntegrityResult } from '../types/index.d.ts';
 
+const ajv = createAjvInstance();
+
 const log = getLogger(import.meta.filename);
-const stringSchemaCache: Record<string, Ajv> = {};
-const objectSchemaCache = new WeakMap<AnySchemaObject, Promise<AnyValidateFunction<unknown>>>();
+const inFlightCompilations = new Map<string, Promise<AnyValidateFunction<unknown>>>();
 
 // Only pre-cache schemas in production to avoid bombarding Github API in development
 if (process.env.NODE_ENV === 'production') await preCachePiesSchema();
@@ -47,46 +46,51 @@ export function validateIntegrity<K extends keyof IntegrityDictionary>(
 }
 
 /**
- * Validates data against a JSON schema using asynchronous compilation and caching.
+ * Validates data against a JSON schema using async compilation, deterministic hashing, and request deduplication.
  * @remarks
- * This function optimizes performance by:
- * 1. Caching Ajv instances for string-identified schemas.
- * 2. Caching compilation promises for schema objects to prevent redundant processing.
+ * This function optimizes performance and reliability through three layers:
+ * 1. **Persistent Cache**: Uses Ajv's internal registry for schemas with a known `$id` or URI.
+ * 2. **Deterministic Hashing**: Anonymous schema objects are automatically assigned a
+ * stable `$id` based on their content, preventing redundant compilations of identical structures.
+ * 3. **Concurrency Lock**: Uses an `in-flight` promise map to ensure that multiple
+ * simultaneous requests for the same new schema trigger only one compilation/load task.
  * @param schema - The validation blueprint. Can be a string identifier (URI/ID) or a schema object.
- * @param data - The payload to validate. Treated as `unknown` to ensure type-safe handling.
+ * Objects without an `$id` will have one auto-generated.
+ * @param data - The payload to validate. Treated as `unknown` for type safety.
  * @returns A promise resolving to a validation result:
  * - `valid`: true if the data satisfies the schema.
  * - `errors`: An array of `ErrorObject` if validation fails, otherwise undefined.
- * @throws {Error} Will throw if the schema cannot be loaded or if compilation fails.
+ * @throws {Error} Will throw if `loadSchema` fails to fetch a remote definition or
+ * if the schema contains syntax errors.
  */
 export async function validateSchema(
   schema: AnySchemaObject | string,
   data: unknown
 ): Promise<{ valid: boolean; errors?: ErrorObject[] }> {
-  let validate: AnyValidateFunction<unknown>;
-  if (typeof schema === 'string') {
-    const cached = schema in stringSchemaCache;
-    log.verbose('validateSchema', { cached, schema });
+  const isString = typeof schema === 'string';
+  const definition = isString ? null : ensureSchemaId(schema);
+  const schemaId = isString ? schema : definition!.$id!;
 
-    if (cached) {
-      validate = stringSchemaCache[schema].getSchema(schema)!;
-    } else {
-      const ajv = createAjvInstance();
+  let validate = ajv.getSchema(schemaId);
 
-      const def = await loadSchema(schema);
-      validate = await ajv.compileAsync(def);
-      stringSchemaCache[schema] = ajv;
+  if (!validate) {
+    // Check if another request is already compiling this
+    let promise = inFlightCompilations.get(schemaId);
+    if (!promise) {
+      promise = (async (): Promise<AnyValidateFunction<unknown>> => {
+        try {
+          const finalDef = definition ?? (await loadSchema(schemaId));
+          return await ajv.compileAsync(finalDef);
+        } finally {
+          inFlightCompilations.delete(schemaId);
+        }
+      })();
+      inFlightCompilations.set(schemaId, promise);
     }
-  } else {
-    let validatePromise = objectSchemaCache.get(schema);
-    if (!validatePromise) {
-      const ajv = createAjvInstance();
-      validatePromise = ajv.compileAsync(schema);
-      objectSchemaCache.set(schema, validatePromise);
-    }
-    validate = await validatePromise;
+
+    validate = await promise;
   }
 
   const valid = !!validate(data);
-  return { valid: valid, errors: validate.errors ?? undefined };
+  return { valid: valid, errors: validate.errors?.slice() };
 }
diff --git a/tests/unit/validators/validator.test.ts b/tests/unit/validators/validator.test.ts
@@ -1,19 +1,36 @@
-import { createAjvInstance, loadSchema, getPiesSchemaUri } from '../../../src/validators/schema/index.ts';
 import * as integrity from '../../../src/validators/integrity/index.ts';
-import { validateSchema, validateIntegrity, preCachePiesSchema } from '../../../src/validators/validator.ts';
+import {
+  createAjvInstance,
+  ensureSchemaId,
+  getPiesSchemaUri,
+  loadSchema
+} from '../../../src/validators/schema/index.ts';
+import * as validator from '../../../src/validators/validator.ts';
 
 import type { Mock } from 'vitest';
 
+const { mockAjv, mockValidate } = vi.hoisted(() => {
+  const mv = vi.fn();
+  return {
+    mockAjv: {
+      compileAsync: vi.fn().mockResolvedValue(mv),
+      getSchema: vi.fn()
+    },
+    mockValidate: mv
+  };
+});
+
 vi.mock('../../../src/validators/schema/index.ts', () => ({
-  createAjvInstance: vi.fn(),
-  loadSchema: vi.fn(),
+  createAjvInstance: vi.fn().mockReturnValue(mockAjv),
+  ensureSchemaId: vi.fn(
+    <T>(schema: T extends object ? T : never) =>
+      ({ ...schema, $id: `test-id-${Object.keys(schema).join('-')}` }) as T & { $id: string }
+  ),
   getPiesSchemaUri: vi.fn(),
+  loadSchema: vi.fn(),
   pies: {
     spec: {
-      message: {
-        A: 'schema-a',
-        B: 'schema-b'
-      }
+      message: { A: 'schema-a', B: 'schema-b' }
     }
   }
 }));
@@ -25,20 +42,21 @@ vi.mock('../../../src/validators/integrity/index.ts', () => ({
 }));
 
 describe('preCachePiesSchema', () => {
-  it('should iterate through all pies message schemas and trigger validation', async () => {
+  it('should iterate through all pies message schemas and trigger compilation', async () => {
+    (loadSchema as Mock).mockResolvedValue({ type: 'object' });
+    (mockAjv.getSchema as Mock).mockReturnValue(undefined); // Force compilation path
+    (mockAjv.compileAsync as Mock).mockResolvedValue(mockValidate.mockReturnValue(true));
     (getPiesSchemaUri as Mock).mockImplementation((val) => `uri-${val}`);
 
-    const mockAjv = {
-      compileAsync: vi.fn().mockResolvedValue(vi.fn().mockReturnValue(true))
-    };
-    (createAjvInstance as Mock).mockReturnValue(mockAjv);
-    (loadSchema as Mock).mockResolvedValue({});
-
-    await preCachePiesSchema();
+    await validator.preCachePiesSchema();
 
     expect(getPiesSchemaUri).toHaveBeenCalledWith('schema-a');
     expect(getPiesSchemaUri).toHaveBeenCalledWith('schema-b');
-    expect(loadSchema).toHaveBeenCalledTimes(2);
+
+    // Verify that the underlying compilation mechanism was triggered for each schema.
+    expect(loadSchema).toHaveBeenCalledWith('uri-schema-a');
+    expect(loadSchema).toHaveBeenCalledWith('uri-schema-b');
+    expect(mockAjv.compileAsync).toHaveBeenCalledTimes(2);
   });
 });
 
@@ -48,7 +66,7 @@ describe('validateIntegrity', () => {
     (integrity.integrityValidators.recordLinkage as Mock).mockReturnValue(mockResult);
 
     const data = { some: 'data' };
-    const result = validateIntegrity('recordLinkage', data);
+    const result = validator.validateIntegrity('recordLinkage', data);
 
     expect(integrity.integrityValidators.recordLinkage).toHaveBeenCalledWith(data);
     expect(result).toBe(mockResult);
@@ -59,65 +77,97 @@ describe('validateIntegrity', () => {
     const data = {};
 
     expect(() => {
-      validateIntegrity(invalidType, data);
+      validator.validateIntegrity(invalidType, data);
     }).toThrow();
   });
 });
 
 describe('validateSchema', () => {
+  beforeEach(() => {
+    (createAjvInstance as Mock).mockReturnValue(mockAjv);
+    (mockAjv.compileAsync as Mock).mockResolvedValue(mockValidate);
+  });
+
   it('should compile and cache a string-based schema', async () => {
     const mockUri = 'https://schema.com/test.json';
     const mockSchemaDef = { type: 'object' };
-    const mockValidate = vi.fn().mockReturnValue(true);
-    const mockAjv = {
-      compileAsync: vi.fn().mockResolvedValue(mockValidate),
-      getSchema: vi.fn().mockReturnValue(mockValidate)
-    };
 
     (loadSchema as Mock).mockResolvedValue(mockSchemaDef);
-    (createAjvInstance as Mock).mockReturnValue(mockAjv);
+    (mockAjv.getSchema as Mock).mockReturnValueOnce(undefined).mockReturnValue(mockValidate);
+    mockValidate.mockReturnValue(true);
 
     // First call: should compile
-    const result1 = await validateSchema(mockUri, { foo: 'bar' });
-
+    const result1 = await validator.validateSchema(mockUri, { foo: 'bar' });
     expect(loadSchema).toHaveBeenCalledWith(mockUri);
     expect(mockAjv.compileAsync).toHaveBeenCalledWith(mockSchemaDef);
     expect(result1.valid).toBe(true);
 
     // Second call: should use cache (mockAjv.getSchema)
-    await validateSchema(mockUri, { foo: 'bar' });
+    const result2 = await validator.validateSchema(mockUri, { foo: 'bar' });
     expect(mockAjv.getSchema).toHaveBeenCalledWith(mockUri);
+    expect(result2.valid).toBe(true);
+    expect(mockAjv.compileAsync).toHaveBeenCalledTimes(1);
   });
 
   it('should handle validation errors correctly', async () => {
     const mockSchemaObj = { type: 'number' };
     const mockErrors = [{ message: 'should be number' }];
-    const mockValidate = Object.assign(vi.fn().mockReturnValue(false), { errors: mockErrors });
+    const failingValidator = Object.assign(vi.fn().mockReturnValue(false), { errors: mockErrors });
 
-    const mockAjv = {
-      compileAsync: vi.fn().mockResolvedValue(mockValidate)
-    };
-    (createAjvInstance as Mock).mockReturnValue(mockAjv);
+    (mockAjv.getSchema as Mock).mockReturnValue(undefined);
+    (mockAjv.compileAsync as Mock).mockResolvedValue(failingValidator);
 
-    const result = await validateSchema(mockSchemaObj, 'not-a-number');
+    const result = await validator.validateSchema(mockSchemaObj, 'not-a-number');
 
     expect(result.valid).toBe(false);
     expect(result.errors).toEqual(mockErrors);
   });
 
-  it('should use WeakMap cache for object-based schemas', async () => {
+  it('should use cache for object-based schemas', async () => {
     const mockSchemaObj = { type: 'string' };
-    const mockValidate = vi.fn().mockReturnValue(true);
-    const mockAjv = {
-      compileAsync: vi.fn().mockResolvedValue(mockValidate)
-    };
-    (createAjvInstance as Mock).mockReturnValue(mockAjv);
+    (mockAjv.getSchema as Mock).mockReturnValueOnce(undefined).mockReturnValue(mockValidate);
+    mockValidate.mockReturnValue(true);
 
     // Call twice with the same object reference
-    await validateSchema(mockSchemaObj, 'test');
-    await validateSchema(mockSchemaObj, 'test');
+    await validator.validateSchema(mockSchemaObj, 'test');
+    await validator.validateSchema(mockSchemaObj, 'test');
+
+    // compileAsync should only be called once for this object
+    expect(mockAjv.compileAsync).toHaveBeenCalledTimes(1);
+    expect(ensureSchemaId).toHaveBeenCalledTimes(2);
+  });
+
+  it('should handle concurrent requests for the same schema with a single compilation', async () => {
+    const mockUri = 'https://schema.com/concurrent.json';
+    const mockSchemaDef = { type: 'string' };
+    (loadSchema as Mock).mockResolvedValue(mockSchemaDef);
+    (mockAjv.getSchema as Mock).mockReturnValue(undefined);
+    mockValidate.mockReturnValue(true);
+
+    // Call validateSchema multiple times without awaiting
+    const [result1, result2] = await Promise.all([
+      validator.validateSchema(mockUri, 'test'),
+      validator.validateSchema(mockUri, 'test')
+    ]);
+    expect(result1.valid).toBe(true);
+    expect(result2.valid).toBe(true);
+
+    // Check that loadSchema and compileAsync were only called once
+    expect(loadSchema).toHaveBeenCalledTimes(1);
+    expect(mockAjv.compileAsync).toHaveBeenCalledTimes(1);
+  });
 
-    // Ajv instance should only be created once for this object
-    expect(createAjvInstance).toHaveBeenCalledTimes(1);
+  it('should propagate errors from schema loading and cleanup in-flight map', async () => {
+    const mockUri = 'https://schema.com/failing.json';
+    const error = new Error('Failed to load schema');
+    (loadSchema as Mock).mockRejectedValue(error);
+    (mockAjv.getSchema as Mock).mockReturnValue(undefined);
+
+    await expect(validator.validateSchema(mockUri, 'test')).rejects.toThrow(error);
+    expect(loadSchema).toHaveBeenCalledTimes(1);
+
+    // Try again, it should try to load again since the in-flight promise was removed
+    await expect(validator.validateSchema(mockUri, 'test')).rejects.toThrow(error);
+    expect(loadSchema).toHaveBeenCalledTimes(2);
   });
 });
