fix: union model in promise (#107)

src/declare.ts

@@ -131,7 +131,7 @@ export type MethodName = string | number | symbol;
 export enum HookType {
   Before = 'Before',
   After = 'After',
-  Around = 'Around ',
+  Around = 'Around',
   AfterReturning = 'AfterReturning',
   AfterThrowing = 'AfterThrowing',
 }
@@ -259,3 +259,12 @@ export interface IHookOptions {
   // Whether to wait for the hook (if the return value of the hook is a promise)
   await?: boolean;
 }
+
+export interface IAroundHookOptions {
+  /**
+   * @deprecated AroundHook will always await the promise, it act as the union model.
+   *
+   * Whether to wait for the hook (if the return value of the hook is a promise)
+   */
+  await?: boolean;
+}

src/decorator.ts

@@ -12,6 +12,7 @@ import {
   IHookOptions,
   IAfterReturningAspectHookFunction,
   IAfterThrowingAspectHookFunction,
+  IAroundHookOptions,
 } from './declare';
 import { markAsAspect, markAsHook } from './helper';
 
@@ -239,7 +240,7 @@ export function After<ThisType = any, Args extends any[] = any, Result = any>(
 export function Around<ThisType = any, Args extends any[] = any, Result = any>(
   token: Token,
   method: MethodName,
-  options: IHookOptions = {},
+  options: IAroundHookOptions = {},
 ) {
   return <T extends Record<K, IAroundAspectHookFunction<ThisType, Args, Result>>, K extends MethodName>(
     target: T,

src/helper/compose.ts

@@ -0,0 +1,54 @@
+interface CallStack {
+  depth: number;
+}
+
+export type Optional<T extends object, K extends keyof T = keyof T> = Omit<T, K> & Partial<Pick<T, K>>;
+
+export type Context<T> = {
+  proceed(): Promise<void> | void;
+} & T;
+
+export type PureContext<T> = Optional<Context<T>, 'proceed'>;
+
+export interface Composed<C> {
+  (ctx: PureContext<C>): Promise<void> | void;
+}
+
+export interface Middleware<C> {
+  (ctx: Context<C>): Promise<void> | void;
+}
+
+function dispatch<C>(
+  middlewareList: Middleware<C>[],
+  idx: number,
+  stack: CallStack,
+  ctx: PureContext<C>,
+): Promise<void> | void {
+  if (idx <= stack.depth) {
+    throw new Error('ctx.proceed() called multiple times');
+  }
+
+  stack.depth = idx;
+
+  let maybePromise: Promise<void> | void;
+
+  if (idx < middlewareList.length) {
+    const middleware = middlewareList[idx];
+
+    maybePromise = middleware({
+      ...ctx,
+      proceed: () => dispatch(middlewareList, idx + 1, stack, ctx),
+    } as Context<C>);
+  } else if (ctx.proceed) {
+    maybePromise = ctx.proceed();
+  }
+
+  return maybePromise;
+}
+
+export default function compose<C>(middlewareList: Middleware<C>[]): Composed<C> {
+  return (ctx) => {
+    const stack: CallStack = { depth: -1 };
+    return dispatch<C>(middlewareList, 0, stack, ctx);
+  };
+}

src/helper/hook-helper.ts

@@ -22,6 +22,7 @@ import {
   IAfterThrowingJoinPoint,
   IInstanceHooks,
 } from '../declare';
+import compose, { Middleware } from './compose';
 
 export const HOOKED_SYMBOL = Symbol('COMMON_DI_HOOKED');
 
@@ -93,14 +94,14 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
   });
 
   return function (this: any, ...args: Args) {
-    let promise: Promise<any> | undefined;
+    let promise: Promise<any> | undefined | void;
     let ret: Result = undefined as any;
     let error: Error | undefined;
     const self = this;
     const originalArgs: Args = args;
 
     try {
-      runAroundHooks();
+      promise = runAroundHooks();
 
       if (promise) {
         // If there is one hook that is asynchronous, convert all of them to asynchronous.
@@ -123,8 +124,6 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
         let p: Promise<Result> | undefined;
         if (promise) {
           p = promise;
-        } else if (isPromiseLike(ret)) {
-          p = ret;
         }
 
         if (p) {
@@ -144,8 +143,13 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
       }
     }
 
-    function runAroundHooks(): Promise<void> | undefined {
-      let i = 0;
+    function runAroundHooks(): Promise<void> | void {
+      const hooks = aroundHooks.map((v) => {
+        const fn = v.hook as Middleware<IAroundJoinPoint<ThisType, Args, Result>>;
+        return fn;
+      });
+      const composed = compose<IAroundJoinPoint<ThisType, Args, Result>>(hooks);
+
       const aroundJoinPoint: IAroundJoinPoint<ThisType, Args, Result> = {
         getArgs: () => {
           return args;
@@ -162,35 +166,21 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
         getThis: () => {
           return self;
         },
-        proceed: () => {
-          i++;
-          promise = runAroundHookAtIndex(i);
-          return promise;
-        },
         setArgs: (_args: Args) => {
           args = _args;
         },
         setResult: (_ret: Result) => {
           ret = _ret;
         },
-      };
-
-      function runAroundHookAtIndex(index: number): Promise<void> | undefined {
-        const aroundHook = aroundHooks[index];
-        if (!aroundHook) {
-          // the innermost layer
-          wrapped();
-          if (promise) {
-            // it seems that asynchronous operations are being performed within the "before" and "after" hooks.
-            return promise;
+        proceed: () => {
+          const maybePromise = wrapped();
+          if (maybePromise && isPromiseLike(maybePromise)) {
+            return maybePromise.then(() => Promise.resolve());
           }
-        } else {
-          promise = runOneHook(aroundHook, aroundJoinPoint, promise);
-          return promise;
-        }
-      }
+        },
+      };
 
-      return runAroundHookAtIndex(0);
+      return composed(aroundJoinPoint);
     }
 
     function runBeforeHooks(): Promise<void> | undefined {
@@ -317,7 +307,7 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
       });
     }
 
-    function wrapped(): Result {
+    function wrapped(): Result | void | Promise<Result | void> {
       promise = runBeforeHooks();
 
       if (promise) {
@@ -340,9 +330,9 @@ export function createHookedFunction<ThisType, Args extends any[], Result>(
       if (promise) {
         return promise.then(() => {
           return ret;
-        }) as any; // 挂载了async hook方法
+        });
       } else {
-        return ret as any;
+        return ret;
       }
     }
   };
@@ -403,7 +393,6 @@ function runOneHook<
   P extends IJoinPoint,
 >(hook: T, joinPoint: P, promise: Promise<any> | undefined): Promise<void> | undefined {
   if (hook.awaitPromise) {
-    // 如果要求await hook，如果之前有promise，直接用，不然创建Promise给下一个使用
     promise = promise || Promise.resolve();
     promise = promise.then(() => {
       return hook.hook(joinPoint);

test/aspect.test.ts

@@ -0,0 +1,127 @@
+import { Injectable, Aspect, Around, IAroundJoinPoint, Injector } from '../src';
+
+describe('aspect', () => {
+  jest.setTimeout(50 * 1000);
+
+  it('test around hook: union model1', async () => {
+    function delay(value: number, time: number): Promise<number> {
+      return new Promise((resolve) => {
+        setTimeout(() => {
+          resolve(value);
+        }, time);
+      });
+    }
+
+    @Injectable()
+    class TestClass {
+      async add(a: number, b: number): Promise<number> {
+        const data = await delay(a + b, 1000);
+        console.log('TestClass add result', data);
+        return data;
+      }
+    }
+
+    @Aspect()
+    @Injectable()
+    class TestAspect {
+      @Around<TestClass, [number, number], number>(TestClass, 'add', { await: true })
+      async interceptAdd(joinPoint: IAroundJoinPoint<TestClass, [number, number], number>) {
+        expect(joinPoint.getMethodName()).toBe('add');
+        expect(joinPoint.getOriginalArgs()).toBeInstanceOf(Array);
+        expect(joinPoint.getThis()).toBeInstanceOf(TestClass);
+        await joinPoint.proceed();
+        const result = await joinPoint.getResult();
+        console.log('TestAspect', result);
+      }
+    }
+
+    @Aspect()
+    @Injectable()
+    class TestAspect2 {
+      @Around<TestClass, [number, number], number>(TestClass, 'add', { await: true })
+      async interceptAdd(joinPoint: IAroundJoinPoint<TestClass, [number, number], number>) {
+        const other = await delay(10, 1000);
+        console.log('TestAspect2 async', other);
+        expect(joinPoint.getMethodName()).toBe('add');
+        expect(joinPoint.getOriginalArgs()).toBeInstanceOf(Array);
+        expect(joinPoint.getThis()).toBeInstanceOf(TestClass);
+        await joinPoint.proceed();
+        const result = await joinPoint.getResult();
+        console.log('TestAspect2', result);
+      }
+    }
+
+    const injector = new Injector();
+    injector.addProviders(TestClass);
+    injector.addProviders(TestAspect);
+    injector.addProviders(TestAspect2);
+
+    const testClass = injector.get(TestClass);
+
+    const result = await testClass.add(1, 2);
+    console.log('TestClass invoke result', result);
+
+    expect(result).toBe(3);
+  });
+
+  it('test union model: union model2', async () => {
+    function delay(value: number, time: number): Promise<number> {
+      return new Promise((resolve) => {
+        setTimeout(() => {
+          resolve(value);
+        }, time);
+      });
+    }
+
+    @Injectable()
+    class TestClass {
+      async add(a: number, b: number): Promise<number> {
+        const data = await delay(a + b, 1000);
+        console.log('TestClass add result', data);
+        return data;
+      }
+    }
+
+    @Aspect()
+    @Injectable()
+    class TestAspect {
+      @Around<TestClass, [number, number], number>(TestClass, 'add', { await: true })
+      async interceptAdd(joinPoint: IAroundJoinPoint<TestClass, [number, number], number>) {
+        expect(joinPoint.getMethodName()).toBe('add');
+        expect(joinPoint.getOriginalArgs()).toBeInstanceOf(Array);
+        expect(joinPoint.getThis()).toBeInstanceOf(TestClass);
+        joinPoint.proceed();
+        const result = await joinPoint.getResult();
+        console.log('TestAspect', result);
+      }
+    }
+
+    @Aspect()
+    @Injectable()
+    class TestAspect2 {
+      @Around<TestClass, [number, number], number>(TestClass, 'add', { await: true })
+      async interceptAdd(joinPoint: IAroundJoinPoint<TestClass, [number, number], number>) {
+        const other = await delay(10, 1000);
+        console.log('TestAspect2 async', other);
+        expect(joinPoint.getMethodName()).toBe('add');
+        expect(joinPoint.getOriginalArgs()).toBeInstanceOf(Array);
+        expect(joinPoint.getThis()).toBeInstanceOf(TestClass);
+        joinPoint.proceed();
+        const result = await joinPoint.getResult();
+        console.log('TestAspect2', result);
+      }
+    }
+
+    const injector = new Injector();
+    injector.addProviders(TestClass);
+    injector.addProviders(TestAspect);
+    injector.addProviders(TestAspect2);
+
+    const testClass = injector.get(TestClass);
+
+    const result = await testClass.add(1, 2);
+    console.log('TestClass invoke result', result);
+
+    expect(result).toBe(3);
+  });
+});