feat: append call stack for runInBackground

eggjs · Feb 8, 2023 · 3d82dd0 · 3d82dd0
1 parent 53c6f06
commit 3d82dd0
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 31 deletions.
diff --git a/plugin/tegg/app.ts b/plugin/tegg/app.ts
@@ -1,15 +1,12 @@
 import './lib/AppLoadUnit';
 import './lib/AppLoadUnitInstance';
 import './lib/EggCompatibleObject';
-import { Application, Context } from 'egg';
-import { BackgroundTaskHelper, PrototypeUtil } from '@eggjs/tegg';
-import { EggPrototype } from '@eggjs/tegg-metadata';
+import { Application } from 'egg';
 import { EggContextCompatibleHook } from './lib/EggContextCompatibleHook';
 import { CompatibleUtil } from './lib/CompatibleUtil';
 import { ModuleHandler } from './lib/ModuleHandler';
 import { EggContextHandler } from './lib/EggContextHandler';
-import { TEGG_CONTEXT } from '@eggjs/egg-module-common';
-import { TEggPluginContext } from './app/extend/context';
+import { hijackRunInBackground } from './lib/run_in_background';
 
 export default class App {
   private readonly app: Application;
@@ -32,32 +29,7 @@ export default class App {
   }
 
   async didLoad() {
-    const eggRunInBackground = this.app.context.runInBackground;
-    this.app.context.runInBackground = function runInBackground(this: TEggPluginContext, scope: (ctx: Context) => Promise<any>) {
-      if (!this[TEGG_CONTEXT]) {
-        return Reflect.apply(eggRunInBackground, this, [ scope ]);
-      }
-      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;
-      });
-    };
-
+    hijackRunInBackground(this.app);
     await this.app.moduleHandler.ready();
     this.compatibleHook = new EggContextCompatibleHook(this.app.moduleHandler);
     this.app.eggContextLifecycleUtil.registerLifecycle(this.compatibleHook);

diff --git a/plugin/tegg/lib/Utils.ts b/plugin/tegg/lib/Utils.ts
@@ -0,0 +1,38 @@
+function prepareObjectStackTrace(_, stack) {
+  return stack;
+}
+
+export function getCalleeFromStack(withLine: boolean, stackIndex?: number) {
+  stackIndex = stackIndex === undefined ? 2 : stackIndex;
+  const limit = Error.stackTraceLimit;
+  const prep = Error.prepareStackTrace;
+
+  Error.prepareStackTrace = prepareObjectStackTrace;
+  Error.stackTraceLimit = 5;
+
+  // capture the stack
+  const obj: any = {};
+  Error.captureStackTrace(obj);
+  let callSite = obj.stack[stackIndex];
+  let fileName;
+  /* istanbul ignore else */
+  if (callSite) {
+    // egg-mock will create a proxy
+    // https://github.com/eggjs/egg-mock/blob/master/lib/app.js#L174
+    fileName = callSite.getFileName();
+    /* istanbul ignore if */
+    if (fileName && fileName.endsWith('egg-mock/lib/app.js')) {
+      // TODO: add test
+      callSite = obj.stack[stackIndex + 1];
+      fileName = callSite.getFileName();
+    }
+  }
+
+  Error.prepareStackTrace = prep;
+  Error.stackTraceLimit = limit;
+
+  /* istanbul ignore if */
+  if (!callSite || !fileName) return '<anonymous>';
+  if (!withLine) return fileName;
+  return `${fileName}:${callSite.getLineNumber()}:${callSite.getColumnNumber()}`;
+}
diff --git a/plugin/tegg/lib/run_in_background.ts b/plugin/tegg/lib/run_in_background.ts
@@ -0,0 +1,59 @@
+import { Application, Context } from 'egg';
+import { BackgroundTaskHelper, PrototypeUtil } from '@eggjs/tegg';
+import { EggPrototype } from '@eggjs/tegg-metadata';
+import { TEGG_CONTEXT } from '@eggjs/egg-module-common';
+import { TEggPluginContext } from '../app/extend/context';
+import { getCalleeFromStack } from './Utils';
+
+export const LONG_STACK_DELIMITER = '\n --------------------\n';
+
+function addLongStackTrace(err: Error, causeError: Error) {
+  const callSiteStack = causeError.stack;
+  if (!callSiteStack || typeof callSiteStack !== 'string') {
+    return;
+  }
+  const index = callSiteStack.indexOf('\n');
+  if (index !== -1) {
+    err.stack += LONG_STACK_DELIMITER + callSiteStack.substring(index + 1);
+  }
+}
+
+export function hijackRunInBackground(app: Application) {
+  const eggRunInBackground = app.context.runInBackground;
+  app.context.runInBackground = function runInBackground(this: TEggPluginContext, scope: (ctx: Context) => Promise<any>) {
+    if (!this[TEGG_CONTEXT]) {
+      return Reflect.apply(eggRunInBackground, this, [ scope ]);
+    }
+    const caseError = new Error('cause');
+    let resolveBackgroundTask;
+    const backgroundTaskPromise = new Promise(resolve => {
+      resolveBackgroundTask = resolve;
+    });
+    const newScope = async () => {
+      try {
+        await scope(this);
+      } catch (e) {
+        addLongStackTrace(e, caseError);
+        throw e;
+      } finally {
+        resolveBackgroundTask();
+      }
+    };
+    const taskName = (scope as any)._name || scope.name || getCalleeFromStack(true, 2);
+    (scope as any)._name = taskName;
+    Object.defineProperty(newScope, 'name', {
+      value: taskName,
+      enumerable: false,
+      configurable: true,
+      writable: false,
+    });
+    Reflect.apply(eggRunInBackground, this, [ newScope ]);
+
+    const proto = PrototypeUtil.getClazzProto(BackgroundTaskHelper);
+    const eggObject = app.eggContainerFactory.getEggObject(proto as EggPrototype);
+    const backgroundTaskHelper = eggObject.obj as BackgroundTaskHelper;
+    backgroundTaskHelper.run(async () => {
+      await backgroundTaskPromise;
+    });
+  };
+}
diff --git a/plugin/tegg/test/app/extend/context.test.ts b/plugin/tegg/test/app/extend/context.test.ts
@@ -5,6 +5,7 @@ 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';
+import { LONG_STACK_DELIMITER } from '../../../lib/run_in_background';
 
 describe('test/app/extend/context.test.ts', () => {
   let app: Application;
@@ -77,5 +78,31 @@ describe('test/app/extend/context.test.ts', () => {
       });
       assert(backgroundIsDone);
     });
+
+    it('stack should be continuous', async () => {
+      let backgroundError;
+      app.on('error', e => {
+        backgroundError = e;
+      });
+      await app.mockModuleContextScope(async ctx => {
+        ctx.runInBackground(async () => {
+          throw new Error('background');
+        });
+        await sleep(1000);
+      });
+      const stack: string = backgroundError.stack;
+      // background
+      // at ~/plugin/tegg/test/app/extend/context.test.ts:88:17
+      // at ~/plugin/tegg/test/app/extend/context.test.ts:82:21 (~/plugin/tegg/lib/run_in_background.ts:34:15)
+      // at ~/node_modules/egg/app/extend/context.js:232:49
+      // --------------------
+      //   at Object.runInBackground (~/plugin/tegg/lib/run_in_background.ts:27:23)
+      // at ~/plugin/tegg/test/app/extend/context.test.ts:87:13
+      // at ~/plugin/tegg/app/extend/application.unittest.ts:49:22
+      // at async Proxy.mockContextScope (~/node_modules/egg-mock/app/extend/application.js:81:12)
+      // at async Context.<anonymous> (~/plugin/tegg/test/app/extend/context.test.ts:86:7)
+      assert(stack.includes(__filename));
+      assert(stack.includes(LONG_STACK_DELIMITER));
+    });
   });
 });