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

gxkl · web-flow · commit e98b8d191837 · 2023-04-12T17:47:09.000+08:00
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

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,6 @@ module.exports = class Application extends Emitter {`
`81`	`81`	`const { AsyncLocalStorage } = require('async_hooks')`
`82`	`82`	`assert(AsyncLocalStorage, 'Requires node 12.17.0 or higher to enable asyncLocalStorage')`
`83`	`83`	`this.ctxStorage = new AsyncLocalStorage()`
`84`		`- this.use(this.createAsyncCtxStorageMiddleware())`
`85`	`84`	`}`
`86`	`85`	`}`
`87`	`86`
`@@ -160,7 +159,12 @@ module.exports = class Application extends Emitter {`
`160`	`159`
`161`	`160`	`const handleRequest = (req, res) => {`
`162`	`161`	`const ctx = this.createContext(req, res)`
`163`		`- return this.handleRequest(ctx, fn)`
	`162`	`+ if (!this.ctxStorage) {`
	`163`	`+ return this.handleRequest(ctx, fn)`
	`164`	`+ }`
	`165`	`+ return this.ctxStorage.run(ctx, async () => {`
	`166`	`+ return await this.handleRequest(ctx, fn)`
	`167`	`+ })`
`164`	`168`	`}`
`165`	`169`
`166`	`170`	`return handleRequest`