feat: make zones overrides and scoped async resources

Author: Viacheslav Bereza

src/index.ts

@@ -35,7 +35,7 @@ export { autorun, reaction, sync, when } from './reaction';
 
 export { contextHook } from './contextHook';
 export { event, type Event } from './event';
-export { service, isolate, isolated, initAsyncHooksZonator } from './service';
+export { service, isolate, isolated, scoped, ScopedAsyncResource, initAsyncHooksZonator } from './service';
 export { hook, type HookWithParams, type HookWithoutParams } from './hook';
 export { untracked } from './untracked';
 

src/service/ServiceHandler.ts

@@ -27,6 +27,8 @@ export const getGlobalServicesRegistry = () => {
   return globalServicesRegistryByZones.get(zoneId)!;
 };
 
+type Configure<T extends object> = (instance: T & ServiceInternal<T>) => void;
+
 /**
  * Registry of all currently active services
  * Necessary for destroy all machanism
@@ -43,6 +45,17 @@ class ServiceState<T extends object> {
    * Set of destroy listeners
    */
   public unsubs = unsubscriber();
+
+  constructor(
+    /**
+     * Service class can be overrided in child packages
+     */
+    public FinalClass: Ctor<T>,
+    /**
+     * List of configures functions
+     */
+    public configures: Configure<T>[] = [],
+  ) {}
 }
 
 /**
@@ -55,14 +68,9 @@ export class ServiceHandler<T extends object> {
   private stateByZones = new Map<ZoneId, ServiceState<T>>();
 
   /**
-   * Service class can be overrided in child packages
-   */
-  private FinalClass: Ctor<T>;
-
-  /**
-   * List of configures functions
+   * Root zone state
    */
-  private configures: ((instance: T & ServiceInternal<T>) => void)[] = [];
+  private rootZoneState: ServiceState<T>;
 
   /**
    * Proxy of service
@@ -73,7 +81,8 @@ export class ServiceHandler<T extends object> {
    * Service handler constructor
    */
   constructor(Class: Ctor<T>) {
-    this.FinalClass = Class;
+    this.rootZoneState = new ServiceState<T>(Class);
+    this.stateByZones.set(0, this.rootZoneState);
     this.proxy = this.createProxy();
   }
 
@@ -190,7 +199,7 @@ export class ServiceHandler<T extends object> {
     const zoneId = getZoneId();
     let state = this.stateByZones.get(zoneId);
     if (!state) {
-      state = new ServiceState();
+      state = new ServiceState<T>(this.rootZoneState.FinalClass, [...this.rootZoneState.configures]);
       this.stateByZones.set(zoneId, state);
     }
     return state;
@@ -210,6 +219,27 @@ export class ServiceHandler<T extends object> {
     this.getStateByCurrentZone().rawInstance = value;
   }
 
+  /**
+   * Zoned get final class
+   */
+  private get finalClass(): Ctor<T> {
+    return this.getStateByCurrentZone().FinalClass;
+  }
+
+  /**
+   * Zoned set final class
+   */
+  private set finalClass(value: Ctor<T>) {
+    this.getStateByCurrentZone().FinalClass = value;
+  }
+
+  /**
+   * Zoned get configures
+   */
+  private get configures(): ((instance: T & ServiceInternal<T>) => void)[] {
+    return this.getStateByCurrentZone().configures;
+  }
+
   /**
    * Zoned get unsubs
    */
@@ -232,7 +262,7 @@ export class ServiceHandler<T extends object> {
   public ensureInstantiate() {
     if (!this.rawInstance) {
       untracked(() => {
-        this.rawInstance = new this.FinalClass();
+        this.rawInstance = new this.finalClass();
         this.register();
         this.runConfigures();
         this.runInit();
@@ -247,7 +277,7 @@ export class ServiceHandler<T extends object> {
     if (this.rawInstance) {
       throw new Error('You should override service before its instantiate');
     }
-    this.FinalClass = OverridedClass;
+    this.finalClass = OverridedClass;
   }
 
   /**

src/service/index.ts

@@ -1,3 +1,4 @@
 export { service } from './service';
 export { isolate, isolated } from './isolate';
+export { ScopedAsyncResource, scoped } from './resource';
 export { initAsyncHooksZonator } from './zones';

src/service/isolate.test.ts

@@ -2,7 +2,7 @@ import * as fs from 'node:fs';
 import { beforeAll, expect, it, vi } from 'vitest';
 import { service, un } from '..';
 import { getZoneId } from './zones';
-import { initAsyncHooksZonator, isolate } from '.';
+import { initAsyncHooksZonator, isolate, ScopedAsyncResource } from '.';
 
 beforeAll(async () => {
   await initAsyncHooksZonator();
@@ -130,6 +130,76 @@ it('isolate works', async () => {
   expect(a.value).toBe(10);
 });
 
+it('ScopedAsyncResource works', async () => {
+  const destroy_spy = vi.fn();
+  class A {
+    value = 0;
+
+    init() {
+      un(() => destroy_spy(this.value));
+    }
+  }
+
+  const a = service(A);
+
+  a.value = 10;
+
+  const zone_id_0 = getZoneId();
+
+  const asyncResource = new ScopedAsyncResource('test');
+
+  let storedResolve: (value: void) => void;
+  const promise = new Promise<void>(resolve => (storedResolve = resolve));
+
+  asyncResource.runInAsyncScope(() => {
+    expect(a.value).toBe(0);
+    a.value = 11;
+
+    const zone_id_1 = getZoneId();
+    expect(zone_id_1).not.toBe(zone_id_0);
+
+    // Test that nested async operation not change current zone
+    process.nextTick(() =>
+      (async () => {
+        const zone_id_async_1 = getZoneId();
+        expect(zone_id_async_1).toBe(zone_id_1);
+        expect(a.value).toBe(11);
+
+        a.value = 12;
+        expect(a.value).toBe(12);
+      })().then(() => {
+        storedResolve();
+      }),
+    );
+
+    expect(a.value).toBe(11);
+    expect(destroy_spy).not.toBeCalled();
+  });
+
+  await promise;
+
+  asyncResource.emitDestroy();
+  expect(destroy_spy).toBeCalledWith(12);
+  destroy_spy.mockClear();
+
+  expect(zone_id_0).toBe(getZoneId());
+  expect(a.value).toBe(10);
+
+  // Next serial isolated resource
+  const asyncResource2 = new ScopedAsyncResource('test2');
+  asyncResource2.runInAsyncScope(() => {
+    expect(a.value).toBe(0);
+    a.value = 9;
+    expect(destroy_spy).not.toBeCalled();
+  });
+
+  asyncResource2.emitDestroy();
+  expect(destroy_spy).toBeCalledWith(9);
+  destroy_spy.mockClear();
+
+  expect(a.value).toBe(10);
+});
+
 vi.mock('node:fs');
 
 it('isolate log not destroyed nested resource', async () => {

src/service/resource.ts

@@ -0,0 +1,14 @@
+import { service } from './service';
+import { createIsolatedAsyncResourceClass } from './zones';
+
+export type ScopedAsyncResource = InstanceType<typeof ScopedAsyncResource>;
+export const ScopedAsyncResource = createIsolatedAsyncResourceClass(() => service.destroy());
+
+export function scoped<This extends ScopedAsyncResource, A extends unknown[] = unknown[], R = unknown>(
+  target: (this: This, ...args: A) => R,
+  _context: ClassMethodDecoratorContext<This, (this: This, ...args: A) => R>,
+) {
+  return function (this: This, ...args: A) {
+    return this.runInAsyncScope(target, this, ...args);
+  };
+}

src/service/service.test.ts

@@ -2,6 +2,7 @@ import { expect, it, vi } from 'vitest';
 import { Initable, signal, un, wrapSignal } from '..';
 import { service } from './service';
 import { getGlobalServicesRegistry } from './ServiceHandler';
+import { initAsyncHooksZonator, isolate } from '.';
 
 it('service works', () => {
   const create_spy = vi.fn();
@@ -247,3 +248,59 @@ it('destroy all should works', () => {
 
   expect(getGlobalServicesRegistry().size).toBe(0);
 });
+
+it('service isolation', async () => {
+  await initAsyncHooksZonator();
+
+  class A {
+    a = 10;
+    test() {
+      return this.a;
+    }
+  }
+
+  const a = service(A);
+
+  expect(a.test()).toBe(10);
+
+  await isolate(async () => {
+    class B extends A {
+      b = 5;
+      test() {
+        return this.a + this.b;
+      }
+    }
+
+    service.override(a, B);
+
+    expect(a.test()).toBe(15);
+  });
+
+  expect(a.test()).toBe(10);
+
+  class D {
+    a = 10;
+    test() {
+      return this.a;
+    }
+  }
+
+  const d = service(D);
+
+  class C extends D {
+    c = 2;
+    test() {
+      return this.a - this.c;
+    }
+  }
+
+  service.override(d, C);
+
+  expect(d.test()).toBe(8);
+
+  await isolate(async () => {
+    expect(d.test()).toBe(8);
+  });
+
+  expect(d.test()).toBe(8);
+});

src/service/zones.ts

@@ -4,14 +4,17 @@ import { NamedLogger } from '../logger/logger';
 const logger = new NamedLogger('zones');
 
 type AsyncHooks = typeof AsyncHooksModule;
+type AsyncResource = AsyncHooksModule.AsyncResource;
 
 export type ZoneId = number;
 
 const isNode = () => typeof window === 'undefined';
 
 let nodeAsyncHooksModulePromise: Promise<AsyncHooks> | undefined;
+let asyncHooks!: AsyncHooks;
+
 export const asyncHooksZonator = async () => {
-  const asyncHooks = await (nodeAsyncHooksModulePromise ??= import('node:async_hooks'));
+  asyncHooks ??= await (nodeAsyncHooksModulePromise ??= import('node:async_hooks'));
   const zoneIndex = new Map<number, ZoneId>();
   let zoneId: ZoneId = 0;
   asyncHooks
@@ -20,21 +23,30 @@ export const asyncHooksZonator = async () => {
         zoneIndex.set(asyncId, zoneIndex.get(triggerAsyncId) ?? 0);
       },
       before(asyncId) {
+        // Called only before start new async context execution
         zoneId = zoneIndex.get(asyncId) ?? 0; // root zone
       },
+      after() {
+        // Called only after current async context execution backwards one level up
+        zoneId = zoneIndex.get(asyncHooks.triggerAsyncId()) ?? 0; // root zone
+      },
       destroy(asyncId) {
         zoneIndex.delete(asyncId);
       },
     })
     .enable();
   return {
     getZoneId: () => zoneId,
-    ensureZone: () => {
-      zoneId = asyncHooks.executionAsyncId();
-      zoneIndex.set(zoneId, zoneId);
+    ensureZone: (overrideZoneId?: ZoneId) => {
+      if (overrideZoneId === undefined) {
+        zoneId = asyncHooks.executionAsyncId();
+        zoneIndex.set(zoneId, zoneId);
+      } else {
+        zoneIndex.set(overrideZoneId, overrideZoneId);
+      }
     },
-    createCheckDestroy: () => {
-      const keepZoneId = zoneId;
+    createCheckDestroy: (overrideZoneId?: ZoneId) => {
+      const keepZoneId = overrideZoneId ?? zoneId;
       return () => {
         if (zoneIndex.has(keepZoneId)) {
           logger.warn('Isolation destroyed but zone has not fulfilled async resources');
@@ -57,11 +69,11 @@ export const initAsyncHooksZonator = async () => {
 
 export interface Zonator {
   getZoneId(): ZoneId;
-  ensureZone(): void; // Function for set current zone_id when new isolate starts
-  createCheckDestroy(): () => void;
+  ensureZone(overrideZoneId?: ZoneId): void; // Function for set current zone_id when new isolate starts
+  createCheckDestroy(overrideZoneId?: ZoneId): () => void;
 }
 
-let zonator: Zonator | undefined;
+let zonator!: Zonator;
 
 export const getZoneId = () => zonator?.getZoneId() ?? 0;
 
@@ -86,3 +98,36 @@ export const createIsolate =
     zonePromise.finally(() => checkZoneDestroy());
     return zonePromise;
   };
+
+export const createIsolatedAsyncResourceClass = (destroy: () => void) => {
+  return class IsolatedAsyncResource {
+    private resource: AsyncResource;
+    private zoneId: ZoneId;
+    private checkZoneDestroy: () => void;
+
+    runInAsyncScope<This, A extends unknown[], R = unknown>(
+      target: (this: This, ...args: A) => R,
+      thisArg?: This,
+      ...args: A
+    ): R {
+      return this.resource.runInAsyncScope(target, thisArg, ...args);
+    }
+
+    constructor(name = 'IsolatedAsyncResource') {
+      if (!zonator) {
+        throw new Error('initAsyncHooksZonator should be called for isolate usage');
+      }
+      this.resource = new asyncHooks.AsyncResource(name, { requireManualDestroy: true });
+      this.zoneId = this.resource.asyncId();
+      zonator.ensureZone(this.zoneId);
+      this.checkZoneDestroy = zonator.createCheckDestroy(this.zoneId);
+    }
+
+    emitDestroy(): this {
+      this.runInAsyncScope(destroy);
+      this.resource.emitDestroy();
+      this.checkZoneDestroy();
+      return this;
+    }
+  };
+};