fix: wait egg background task done before destroy tegg ctx (#87)

eggjs · Jan 29, 2023 · deea4d8 · deea4d8
1 parent 6b31708
commit deea4d8
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 2 deletions.
diff --git a/plugin/tegg/app.ts b/plugin/tegg/app.ts
@@ -1,7 +1,9 @@
 import './lib/AppLoadUnit';
 import './lib/AppLoadUnitInstance';
 import './lib/EggCompatibleObject';
-import { Application } from 'egg';
+import { Application, Context } from 'egg';
+import { BackgroundTaskHelper, PrototypeUtil } from '@eggjs/tegg';
+import { EggPrototype } from '@eggjs/tegg-metadata';
 import { EggContextCompatibleHook } from './lib/EggContextCompatibleHook';
 import { CompatibleUtil } from './lib/CompatibleUtil';
 import { ModuleHandler } from './lib/ModuleHandler';
@@ -28,6 +30,29 @@ export default class App {
   }
 
   async didLoad() {
+    const eggRunInBackground = this.app.context.runInBackground;
+    this.app.context.runInBackground = function runInBackground(this: Context, scope: (ctx: Context) => Promise<any>) {
+      let resolveBackgroundTask;
+      const backgroundTaskPromise = new Promise(resolve => {
+        resolveBackgroundTask = resolve;
+      });
+      const newScope = async () => {
+        try {
+          await scope(this);
+        } finally {
+          resolveBackgroundTask();
+        }
+      };
+      Reflect.apply(eggRunInBackground, this, [ newScope ]);
+
+      const proto = PrototypeUtil.getClazzProto(BackgroundTaskHelper);
+      const eggObject = this.app.eggContainerFactory.getEggObject(proto as EggPrototype);
+      const backgroundTaskHelper = eggObject.obj as BackgroundTaskHelper;
+      backgroundTaskHelper.run(async () => {
+        await backgroundTaskPromise;
+      });
+    };
+
     await this.app.moduleHandler.ready();
     this.compatibleHook = new EggContextCompatibleHook(this.app.moduleHandler);
     this.app.eggContextLifecycleUtil.registerLifecycle(this.compatibleHook);

diff --git a/plugin/tegg/lib/EggContextCompatibleHook.ts b/plugin/tegg/lib/EggContextCompatibleHook.ts
@@ -1,4 +1,4 @@
-import { LifecycleHook, ObjectInitType } from '@eggjs/tegg';
+import { BackgroundTaskHelper, LifecycleHook, ObjectInitType, PrototypeUtil } from '@eggjs/tegg';
 import { EggContainerFactory, EggContext, EggContextLifecycleContext } from '@eggjs/tegg-runtime';
 import { EggPrototype } from '@eggjs/tegg-metadata';
 import { ModuleHandler } from './ModuleHandler';
@@ -38,6 +38,13 @@ export class EggContextCompatibleHook implements LifecycleHook<EggContextLifecyc
       await Promise.all(this.initProtoList.map(async proto => {
         await EggContainerFactory.getOrCreateEggObject(proto);
       }));
+    } else {
+      // Use for ctx.runInBackground.
+      // BackgroundTaskHelper should get by sync,
+      // or tegg context may be destroyed before background task run.
+      // So create it in preCreate.
+      const protoObj = PrototypeUtil.getClazzProto(BackgroundTaskHelper);
+      await EggContainerFactory.getOrCreateEggObject(protoObj as EggPrototype);
     }
   }
 }
diff --git a/plugin/tegg/test/app/extend/context.test.ts b/plugin/tegg/test/app/extend/context.test.ts
@@ -2,6 +2,7 @@ import assert from 'assert';
 import path from 'path';
 import mm from 'egg-mock';
 import { Application } from 'egg';
+import sleep from 'mz-modules/sleep';
 import AppService from '../../fixtures/apps/egg-app/modules/multi-module-service/AppService';
 import PersistenceService from '../../fixtures/apps/egg-app/modules/multi-module-repo/PersistenceService';
 
@@ -50,4 +51,17 @@ describe('test/app/extend/context.test.ts', () => {
       });
     });
   });
+
+  describe('runInBackground', () => {
+    it('should notify background task helper', async () => {
+      let backgroundIsDone = false;
+      await app.mockModuleContextScope(async ctx => {
+        ctx.runInBackground(async () => {
+          await sleep(100);
+          backgroundIsDone = true;
+        });
+      });
+      assert(backgroundIsDone);
+    });
+  });
 });