feat: support asyncLocalStorage (#1721)

fengmk2 · web-flow · commit a29312212839 · 2022-12-06T19:32:12.000+08:00
pick from #1455 pick from #1718
diff --git a/.github/workflows/node.js.yml b/.github/workflows/node.js.yml
@@ -13,7 +13,7 @@ jobs:
 
     strategy:
       matrix:
-        node-version: [12.x, 14.x, 16.x]
+        node-version: [12.x, 14.x, 16.x, 18.x]
 
     steps:
     - uses: actions/checkout@v2
diff --git a/__tests__/application/currentContext.js b/__tests__/application/currentContext.js
@@ -0,0 +1,64 @@
+
+'use strict';
+
+const request = require('supertest');
+const assert = require('assert');
+const Koa = require('../..');
+
+describe('app.currentContext', () => {
+  it('should throw error if AsyncLocalStorage not support', () => {
+    if (require('async_hooks').AsyncLocalStorage) return;
+    assert.throws(() => new Koa({ asyncLocalStorage: true }),
+      /Requires node 12\.17\.0 or higher to enable asyncLocalStorage/);
+  });
+
+  it('should get currentContext return context when asyncLocalStorage enable', async() => {
+    if (!require('async_hooks').AsyncLocalStorage) return;
+
+    const app = new Koa({ asyncLocalStorage: true });
+
+    app.use(async ctx => {
+      assert(ctx === app.currentContext);
+      await new Promise(resolve => {
+        setTimeout(() => {
+          assert(ctx === app.currentContext);
+          resolve();
+        }, 1);
+      });
+      await new Promise(resolve => {
+        assert(ctx === app.currentContext);
+        setImmediate(() => {
+          assert(ctx === app.currentContext);
+          resolve();
+        });
+      });
+      assert(ctx === app.currentContext);
+      app.currentContext.body = 'ok';
+    });
+
+    const requestServer = async() => {
+      assert(app.currentContext === undefined);
+      await request(app.callback()).get('/').expect('ok');
+      assert(app.currentContext === undefined);
+    };
+
+    await Promise.all([
+      requestServer(),
+      requestServer(),
+      requestServer(),
+      requestServer(),
+      requestServer()
+    ]);
+  });
+
+  it('should get currentContext return undefined when asyncLocalStorage disable', async() => {
+    const app = new Koa();
+
+    app.use(async ctx => {
+      assert(app.currentContext === undefined);
+      ctx.body = 'ok';
+    });
+
+    await request(app.callback()).get('/').expect('ok');
+  });
+});
diff --git a/__tests__/application/index.js b/__tests__/application/index.js
@@ -6,7 +6,8 @@ const assert = require('assert');
 const Koa = require('../..');
 
 describe('app', () => {
-  it('should handle socket errors', done => {
+  // ignore test on Node.js v18
+  (/^v18\./.test(process.version) ? it.skip : it)('should handle socket errors', done => {
     const app = new Koa();
 
     app.use((ctx, next) => {
diff --git a/lib/application.js b/lib/application.js
@@ -8,6 +8,7 @@
 const isGeneratorFunction = require('is-generator-function');
 const debug = require('debug')('koa:application');
 const onFinished = require('on-finished');
+const assert = require('assert');
 const response = require('./response');
 const compose = require('koa-compose');
 const context = require('./context');
@@ -64,6 +65,12 @@ module.exports = class Application extends Emitter {
     if (util.inspect.custom) {
       this[util.inspect.custom] = this.inspect;
     }
+    if (options.asyncLocalStorage) {
+      const { AsyncLocalStorage } = require('async_hooks');
+      assert(AsyncLocalStorage, 'Requires node 12.17.0 or higher to enable asyncLocalStorage');
+      this.ctxStorage = new AsyncLocalStorage();
+      this.use(createAsyncCtxStorage(this));
+    }
   }
 
   /**
@@ -153,6 +160,13 @@ module.exports = class Application extends Emitter {
     return handleRequest;
   }
 
+  /**
+   * return currnect contenxt from async local storage
+   */
+  get currentContext() {
+    if (this.ctxStorage) return this.ctxStorage.getStore();
+  }
+
   /**
    * Handle request in callback.
    *
@@ -222,6 +236,14 @@ module.exports = class Application extends Emitter {
   }
 };
 
+function createAsyncCtxStorage(app) {
+  return async function asyncCtxStorage(ctx, next) {
+    await app.ctxStorage.run(ctx, async() => {
+      return await next();
+    });
+  };
+}
+
 /**
  * Response helper.
  */
