feat: implement cork/uncork for eventbus (#60)

killagu · web-flow · commit 38114bd7ea3b · 2022-11-28T09:35:12.000+08:00
diff --git a/core/eventbus-decorator/README.md b/core/eventbus-decorator/README.md
@@ -24,6 +24,27 @@ class Foo {
 }
 ```
 
+### cork events
+
+Cache events in memory until uncork. 
+
+```ts
+class Foo {
+  @Inject()
+  private readonly eventBus: ContextEventBus;
+
+  bar() {
+    this.eventBus.cork();
+    // ...do something
+    this.eventBus.emit('hello', '01');
+    // ...do other things
+    
+    // emit all cached events
+    this.eventBus.uncork();
+  }
+}
+```
+
 ### handle event
 
 ```ts
diff --git a/core/eventbus-decorator/src/EventBus.ts b/core/eventbus-decorator/src/EventBus.ts
@@ -9,6 +9,15 @@ export type EventName = string | symbol;
  * use `emit` to emit a event
  */
 export interface EventBus extends Pick<TypedEventEmitter<Events>, 'emit'> {
+  cork(corkId: string);
+  uncork(corkId: string);
+}
+
+export const CORK_ID = Symbol.for('eventBus#corkId');
+
+export interface ContextEventBus extends EventBus {
+  cork();
+  uncork();
 }
 
 export type EventKeys = keyof Events;
diff --git a/core/eventbus-runtime/package.json b/core/eventbus-runtime/package.json
@@ -54,7 +54,8 @@
     "@eggjs/tegg-runtime": "^1.4.1",
     "coffee": "^5.4.0",
     "egg": "^2.29.3",
-    "mm": "^3.2.0"
+    "mm": "^3.2.0",
+    "mz-modules": "^2.1.0"
   },
   "publishConfig": {
     "access": "public"
diff --git a/core/eventbus-runtime/src/SingletonEventBus.ts b/core/eventbus-runtime/src/SingletonEventBus.ts
@@ -1,5 +1,5 @@
 import { AccessLevel, Inject, SingletonProto } from '@eggjs/core-decorator';
-import { EventBus, Events, EventWaiter, EventName } from '@eggjs/eventbus-decorator';
+import { EventBus, Events, EventWaiter, EventName, CORK_ID } from '@eggjs/eventbus-decorator';
 import { EggContext } from '@eggjs/tegg-runtime';
 import type { EggLogger } from 'egg';
 import { EventContextFactory } from './EventContextFactory';
@@ -9,10 +9,21 @@ import awaitEvent from 'await-event';
 import awaitFirst from 'await-first';
 
 // from typed-emitter
-type Arguments<T> = [ T ] extends [ (...args: infer U) => any ]
+type Array<T> = [ T ] extends [ (...args: infer U) => any ]
   ? U
   : [ T ] extends [ void ] ? [] : [ T ];
 
+export interface Event {
+  name: EventName;
+  args: Array<any>;
+  context?: EggContext;
+}
+
+export interface CorkEvents {
+  times: number;
+  events: Array<Event>;
+}
+
 @SingletonProto({
   // TODO 需要考虑支持别名
   // SingletonEventBus 同时实现了两个接口
@@ -31,6 +42,10 @@ export class SingletonEventBus implements EventBus, EventWaiter {
   @Inject()
   private readonly logger: EggLogger;
 
+  private corkIdSequence = 0;
+
+  private readonly corkedEvents = new Map<string /* corkId */, CorkEvents>();
+
   /**
    * only use for ensure event will happen
    */
@@ -39,38 +54,88 @@ export class SingletonEventBus implements EventBus, EventWaiter {
     return this;
   }
 
-  async await<E extends keyof Events>(event: E): Promise<Arguments<Events[E]>> {
+  async await<E extends keyof Events>(event: E): Promise<Array<Events[E]>> {
     return awaitEvent(this.emitter, event);
   }
 
-  awaitFirst<E extends keyof Events>(...e: Array<E>): Promise<{ event: EventName, args: Arguments<Events[E]> }> {
+  awaitFirst<E extends keyof Events>(...e: Array<E>): Promise<{ event: EventName, args: Array<Events[E]> }> {
     return awaitFirst(this.emitter, e);
   }
 
-  emit<E extends keyof Events>(event: E, ...args: Arguments<Events[E]>): boolean {
+  emit<E extends keyof Events>(event: E, ...args: Array<Events[E]>): boolean {
     const ctx = this.eventContextFactory.createContext();
     const hasListener = this.eventHandlerFactory.hasListeners(event);
     this.doEmit(ctx, event, args);
     return hasListener;
   }
 
-  emitWithContext<E extends keyof Events>(parentContext: EggContext, event: E, args: Arguments<Events[E]>): boolean {
-    const ctx = this.eventContextFactory.createContext(parentContext);
+  generateCorkId(): string {
+    return String(++this.corkIdSequence);
+  }
+
+  cork(corkId: string) {
+    let corkEvents = this.corkedEvents.get(corkId);
+    if (!corkEvents) {
+      corkEvents = {
+        times: 0,
+        events: [],
+      } as unknown as CorkEvents;
+      this.corkedEvents.set(corkId, corkEvents);
+    }
+    corkEvents!.times++;
+  }
+
+  uncork(corkId: string) {
+    const corkEvents = this.corkedEvents.get(corkId);
+    if (!corkEvents) {
+      throw new Error(`eventbus corkId ${corkId} not found`);
+    }
+    if (--corkEvents.times !== 0) {
+      return;
+    }
+    this.corkedEvents.delete(corkId);
+    for (const event of corkEvents.events) {
+      if (event.context) {
+        this.doEmitWithContext(event.context, event.name, event.args);
+      }
+    }
+  }
+
+  queueEvent(corkId: string, event: Event) {
+    const corkdEvents = this.corkedEvents.get(corkId);
+    if (!corkdEvents) {
+      throw new Error(`eventbus corkId ${corkId} not found`);
+    }
+    corkdEvents.events.push(event);
+  }
+
+  emitWithContext<E extends keyof Events>(parentContext: EggContext, event: E, args: Array<Events[E]>): boolean {
+    const corkId = parentContext.get(CORK_ID);
+    const hasListener = this.eventHandlerFactory.hasListeners(event);
+    if (corkId) {
+      this.queueEvent(corkId, { name: event, args, context: parentContext });
+      return hasListener;
+    }
+    return this.doEmitWithContext(parentContext, event, args);
+  }
+
+  private doEmitWithContext(parentContext: EggContext, event: EventName, args: Array<any>): boolean {
     const hasListener = this.eventHandlerFactory.hasListeners(event);
+    const ctx = this.eventContextFactory.createContext(parentContext);
     this.doEmit(ctx, event, args);
     return hasListener;
   }
 
-  doOnceEmit<E extends keyof Events>(event: E, args: Arguments<Events[E]>) {
+  private doOnceEmit(event: EventName, args: Array<any>) {
     try {
       this.emitter.emit(event, ...args);
     } catch (e) {
-      e.message = `[EventBus] process once event ${event} failed: ${e.message}`;
+      e.message = `[EventBus] process once event ${String(event)} failed: ${e.message}`;
       this.logger.error(e);
     }
   }
 
-  async doEmit<E extends keyof Events>(ctx: EggContext, event: E, args: Arguments<Events[E]>) {
+  private async doEmit(ctx: EggContext, event: EventName, args: Array<any>) {
     const lifecycle = {};
     if (ctx.init) {
       await ctx.init(lifecycle);
@@ -82,12 +147,12 @@ export class SingletonEventBus implements EventBus, EventWaiter {
           await Reflect.apply(handler.handle, handler, args);
         } catch (e) {
           // should wait all handlers done then destroy ctx
-          e.message = `[EventBus] process event ${event} failed: ${e.message}`;
+          e.message = `[EventBus] process event ${String(event)} failed: ${e.message}`;
           this.logger.error(e);
         }
       }));
     } catch (e) {
-      e.message = `[EventBus] process event ${event} failed: ${e.message}`;
+      e.message = `[EventBus] process event ${String(event)} failed: ${e.message}`;
       this.logger.error(e);
     } finally {
       if (ctx.destroy) {
diff --git a/core/eventbus-runtime/test/EventBus.test.ts b/core/eventbus-runtime/test/EventBus.test.ts
@@ -7,9 +7,10 @@ import { EggProtoImplClass, PrototypeUtil } from '@eggjs/core-decorator';
 import { LoaderFactory } from '../../loader';
 import { EggTestContext } from '../../test-util';
 import { EventContextFactory, EventHandlerFactory, SingletonEventBus } from '..';
-import { EventInfoUtil } from '@eggjs/eventbus-decorator';
+import { EventInfoUtil, CORK_ID } from '@eggjs/eventbus-decorator';
 import assert from 'assert';
 import { Timeout0Handler, Timeout100Handler, TimeoutProducer } from './fixtures/modules/event/MultiEvent';
+import sleep from 'mz-modules/sleep';
 
 describe('test/EventBus.test.ts', () => {
   async function getLoadUnitInstance(moduleDir: string): Promise<LoadUnitInstance> {
@@ -121,4 +122,70 @@ describe('test/EventBus.test.ts', () => {
     await timeoutEvent;
     assert(Timeout100Handler.called === true);
   });
+
+  it('cork should work', async () => {
+    const ctx = new EggTestContext();
+    const eventContextFactory = await getObject(EventContextFactory);
+    eventContextFactory.registerContextCreator(() => {
+      return ctx;
+    });
+    const eventHandlerFactory = await getObject(EventHandlerFactory);
+    eventHandlerFactory.registerHandler(
+      EventInfoUtil.getEventName(HelloHandler)!,
+      PrototypeUtil.getClazzProto(HelloHandler) as EggPrototype);
+
+    const eventBus = await getObject(SingletonEventBus);
+    const corkId = eventBus.generateCorkId();
+    ctx.set(CORK_ID, corkId);
+    eventBus.cork(corkId);
+
+    const helloHandler = await getObject(HelloHandler, ctx);
+    const helloEvent = eventBus.await('hello');
+    let eventTime: number;
+    mm(helloHandler, 'handle', async () => {
+      eventTime = Date.now();
+    });
+    eventBus.emitWithContext(ctx, 'hello', [ '01' ]);
+    const triggerTime = Date.now();
+
+    await sleep(100);
+    eventBus.uncork(corkId);
+    await helloEvent;
+    assert(eventTime >= triggerTime + 100);
+  });
+
+  it('multi cork should work', async () => {
+    const ctx = new EggTestContext();
+    const eventContextFactory = await getObject(EventContextFactory);
+    eventContextFactory.registerContextCreator(() => {
+      return ctx;
+    });
+    const eventHandlerFactory = await getObject(EventHandlerFactory);
+    eventHandlerFactory.registerHandler(
+      EventInfoUtil.getEventName(HelloHandler)!,
+      PrototypeUtil.getClazzProto(HelloHandler) as EggPrototype);
+
+    const eventBus = await getObject(SingletonEventBus);
+    const corkId = eventBus.generateCorkId();
+    ctx.set(CORK_ID, corkId);
+    eventBus.cork(corkId);
+    eventBus.cork(corkId);
+
+    const helloHandler = await getObject(HelloHandler, ctx);
+    const helloEvent = eventBus.await('hello');
+    let eventTime: number;
+    mm(helloHandler, 'handle', async () => {
+      eventTime = Date.now();
+    });
+    eventBus.emitWithContext(ctx, 'hello', [ '01' ]);
+    const triggerTime = Date.now();
+
+    await sleep(100);
+    eventBus.uncork(corkId);
+    await sleep(100);
+    eventBus.uncork(corkId);
+
+    await helloEvent;
+    assert(eventTime >= triggerTime + 200);
+  });
 });
diff --git a/plugin/eventbus/lib/EggContextEventBus.ts b/plugin/eventbus/lib/EggContextEventBus.ts
@@ -1,12 +1,14 @@
+import assert from 'assert';
 import { Context } from 'egg';
-import { EventBus, Events, PrototypeUtil } from '@eggjs/tegg';
+import { Events, PrototypeUtil, CORK_ID, ContextEventBus } from '@eggjs/tegg';
 import { SingletonEventBus } from '@eggjs/tegg-eventbus-runtime';
 import { EggPrototype } from '@eggjs/tegg-metadata';
 import { EggContext } from '@eggjs/tegg-runtime';
 
-export class EggContextEventBus implements EventBus {
+export class EggContextEventBus implements ContextEventBus {
   private readonly eventBus: SingletonEventBus;
   private readonly context: EggContext;
+  private corkId?: string;
 
   constructor(ctx: Context) {
     this.context = ctx.teggContext;
@@ -15,6 +17,19 @@ export class EggContextEventBus implements EventBus {
     this.eventBus = eggObject.obj as SingletonEventBus;
   }
 
+  cork() {
+    if (!this.corkId) {
+      this.corkId = this.eventBus.generateCorkId();
+      this.context.set(CORK_ID, this.corkId);
+    }
+    this.eventBus.cork(this.corkId);
+  }
+
+  uncork() {
+    assert(this.corkId, 'eventbus uncork without cork');
+    this.eventBus.uncork(this.corkId);
+  }
+
   emit<E extends keyof Events>(event: E, ...args: any): boolean {
     return this.eventBus.emitWithContext(this.context, event, args);
   }
diff --git a/plugin/eventbus/package.json b/plugin/eventbus/package.json
@@ -59,7 +59,8 @@
     "@eggjs/tegg-plugin": "^1.4.1",
     "await-event": "^2.1.0",
     "egg": "^2.29.4",
-    "egg-mock": "^4.0.1"
+    "egg-mock": "^4.0.1",
+    "mz-modules": "^2.1.0"
   },
   "publishConfig": {
     "access": "public"
diff --git a/plugin/eventbus/test/eventbus.test.ts b/plugin/eventbus/test/eventbus.test.ts
@@ -1,6 +1,7 @@
 import assert from 'assert';
 import path from 'path';
 import mm from 'egg-mock';
+import sleep from 'mz-modules/sleep';
 import { HelloService } from './fixtures/apps/event-app/app/event-module/HelloService';
 import { HelloLogger } from './fixtures/apps/event-app/app/event-module/HelloLogger';
 
@@ -43,4 +44,26 @@ describe('test/eventbus.test.ts', () => {
     await helloEvent;
     assert(msg === '01');
   });
+
+  it('cork/uncork should work', async () => {
+    ctx = await app.mockModuleContext();
+
+    const helloService = await ctx.getEggObject(HelloService);
+    let helloTime: number;
+    // helloLogger is in child context
+    mm(HelloLogger.prototype, 'handle', () => {
+      helloTime = Date.now();
+    });
+    helloService.cork();
+    const triggerTime = Date.now();
+    helloService.hello();
+
+    await sleep(100);
+    helloService.uncork();
+
+    const eventWaiter = await app.getEventWaiter();
+    const helloEvent = eventWaiter.await('helloEgg');
+    await helloEvent;
+    assert(helloTime >= triggerTime + 100);
+  });
 });
diff --git a/plugin/eventbus/test/fixtures/apps/event-app/app/event-module/HelloService.ts b/plugin/eventbus/test/fixtures/apps/event-app/app/event-module/HelloService.ts
@@ -1,4 +1,4 @@
-import { AccessLevel, ContextProto, Inject, EventBus } from '@eggjs/tegg';
+import { AccessLevel, ContextProto, Inject, ContextEventBus } from '@eggjs/tegg';
 
 declare module '@eggjs/tegg' {
   interface Events {
@@ -12,7 +12,15 @@ declare module '@eggjs/tegg' {
 })
 export class HelloService {
   @Inject()
-  private readonly eventBus: EventBus;
+  private readonly eventBus: ContextEventBus;
+
+  cork() {
+    this.eventBus.cork();
+  }
+
+  uncork() {
+    this.eventBus.uncork();
+  }
 
   hello() {
     this.eventBus.emit('helloEgg', '01');
diff --git a/plugin/schedule/lib/ScheduleSubscriberRegister.ts b/plugin/schedule/lib/ScheduleSubscriberRegister.ts
