fix: can not get currentContext in error handler (#1758)

koajs · Apr 12, 2023 · e98b8d1 · e98b8d1
1 parent bec13ec
commit e98b8d1
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/__tests__/application/currentContext.js b/__tests__/application/currentContext.js
@@ -61,4 +61,50 @@ describe('app.currentContext', () => {
 
     await request(app.callback()).get('/').expect('ok')
   })
+
+  it('should get currentContext return context in error handler when asyncLocalStorage enable', async () => {
+    const app = new Koa({ asyncLocalStorage: true })
+
+    app.use(async () => {
+      throw new Error('error message')
+    })
+
+    const handleError = new Promise((resolve, reject) => {
+      app.on('error', (err, ctx) => {
+        try {
+          assert.strictEqual(err.message, 'error message')
+          assert.strictEqual(app.currentContext, ctx)
+          resolve()
+        } catch (e) {
+          reject(e)
+        }
+      })
+    })
+
+    await request(app.callback()).get('/').expect('Internal Server Error')
+    await handleError
+  })
+
+  it('should get currentContext return undefined in error handler when asyncLocalStorage disable', async () => {
+    const app = new Koa()
+
+    app.use(async () => {
+      throw new Error('error message')
+    })
+
+    const handleError = new Promise((resolve, reject) => {
+      app.on('error', (err, ctx) => {
+        try {
+          assert.strictEqual(err.message, 'error message')
+          assert.strictEqual(app.currentContext, undefined)
+          resolve()
+        } catch (e) {
+          reject(e)
+        }
+      })
+    })
+
+    await request(app.callback()).get('/').expect('Internal Server Error')
+    await handleError
+  })
 })
diff --git a/lib/application.js b/lib/application.js
@@ -81,7 +81,6 @@ module.exports = class Application extends Emitter {
       const { AsyncLocalStorage } = require('async_hooks')
       assert(AsyncLocalStorage, 'Requires node 12.17.0 or higher to enable asyncLocalStorage')
       this.ctxStorage = new AsyncLocalStorage()
-      this.use(this.createAsyncCtxStorageMiddleware())
     }
   }
 
@@ -160,7 +159,12 @@ module.exports = class Application extends Emitter {
 
     const handleRequest = (req, res) => {
       const ctx = this.createContext(req, res)
-      return this.handleRequest(ctx, fn)
+      if (!this.ctxStorage) {
+        return this.handleRequest(ctx, fn)
+      }
+      return this.ctxStorage.run(ctx, async () => {
+        return await this.handleRequest(ctx, fn)
+      })
     }
 
     return handleRequest