Automattic#10020: asyncLocalStorage supporting

examples/asyncLocalStorage/asyncLocalStorageExample.js

@@ -0,0 +1,136 @@
+'use strict';
+
+const mongoose = require('../..');
+const { MongoMemoryServer } = require('mongodb-memory-server');
+const uuid = require('uuid').v4;
+const _ = require('lodash');
+const callContext = require('./callContext');
+
+const pluginSave = (schema) => {
+  schema.pre(['save'], function() {
+    const store = callContext.get();
+
+    if (this.name !== store.name) {
+      console.error('[static-hooks] [pre] [save]', this.name, store.name);
+    } else {
+      console.log('[OK] [static-hooks] [pre] [save]');
+    }
+  });
+
+  schema.post(['save'], function() {
+    const store = callContext.get();
+
+    if (this.name !== store.name) {
+      console.error('[ERROR] [static-hooks] [post] [save]', this.name, store.name);
+    } else {
+      console.log('[OK] [static-hooks] [post] [save]');
+    }
+  });
+};
+
+const pluginOther = (schema) => {
+  schema.pre(['find', 'findOne', 'count', 'aggregate'], function() {
+    const store = callContext.get();
+
+    if (this._conditions.name !== store.name) {
+      console.error(`[ERROR] [static-hooks] [pre] [${this.op}]`, this._conditions.name, store.name);
+    } else {
+      console.log(`[OK] [static-hooks] [pre] [${this.op}]`);
+    }
+  });
+
+  schema.post(['find', 'findOne', 'count', 'aggregate'], function() {
+    const store = callContext.get();
+    if (this._conditions.name !== store.name) {
+      console.error(`[ERROR] [static-hooks] [post] [${this.op}]`, this._conditions.name, store.name);
+    } else {
+      console.log(`[OK] [static-hooks] [post] [${this.op}]`);
+    }
+  });
+};
+
+mongoose.plugin(pluginSave);
+mongoose.plugin(pluginOther);
+
+let createCounter = 0;
+let findCallbackCounter = 0;
+let findPromiseCounter = 0;
+
+(async() => {
+  const mongod = new MongoMemoryServer();
+  const uri = await mongod.getUri();
+
+  await mongoose.connect(uri, {
+    useNewUrlParser: true,
+    useUnifiedTopology: true
+  });
+
+  const userSchema = new mongoose.Schema({ name: String });
+  const UserModel = mongoose.model('UserModel', userSchema);
+
+  const names = [];
+
+  // prepare data
+  await new Promise((resolve, reject) => {
+    for (let i = 0; i < 50; ++i) {
+      setTimeout(async() => {
+        const name = uuid();
+        names.push(name);
+        callContext.enter({ name });
+
+        const user = new UserModel({ name });
+        try {
+          await user.save();
+        } catch (err) {
+          reject(err);
+        }
+
+        createCounter++;
+
+        if (createCounter === 50) {
+          resolve();
+        }
+      }, _.random(10, 50));
+    }
+  });
+
+  for (let i = 0; i < 50; ++i) {
+    setTimeout(async() => {
+      const name = names[_.random(0, names.length - 1)];
+      callContext.enter({ name });
+      // for testing callback
+      UserModel.find({ name }, (err, data) => {
+        ++findCallbackCounter;
+        data = data[0];
+        const store = callContext.get();
+        if (data.name !== store.name) {
+          console.error(`[ERROR] ${findCallbackCounter}: post-find-in-callback`, data.name, store.name);
+        } else {
+          console.log(`[OK] ${findCallbackCounter}: post-find-in-callback`);
+        }
+      });
+
+      // for tesing promise
+      let data = await UserModel.find({ name }).exec();
+      ++findPromiseCounter;
+
+      data = data[0];
+      const store = callContext.get();
+      if (data.name !== store.name) {
+        console.error(`[ERROR] ${findPromiseCounter}: post-find-in-promise`, data.name, store.name);
+      } else {
+        console.log(`[OK] ${findPromiseCounter}: post-find-in-promise`);
+      }
+    }, _.random(10, 50));
+  }
+})();
+
+const exit = () => {
+  if (createCounter === 50 && findCallbackCounter === 50 && findPromiseCounter === 50) {
+    process.exit(0);
+  } else {
+    setTimeout(exit, 100);
+  }
+};
+
+exit();

examples/asyncLocalStorage/callContext.js

@@ -0,0 +1,19 @@
+'use strict';
+
+const { AsyncLocalStorage } = require('async_hooks');
+const asyncLocalStorage = new AsyncLocalStorage();
+
+const enter = (contextData) => {
+  asyncLocalStorage.enterWith(contextData);
+};
+
+const get = (defaultValue) => {
+  let obj = asyncLocalStorage.getStore();
+  if (!obj && defaultValue) {
+    obj = defaultValue;
+  }
+  return obj;
+};
+
+module.exports.enter = enter;
+module.exports.get = get;

examples/asyncLocalStorage/package.json

@@ -0,0 +1,18 @@
+{
+  "name": "async-local-storage-example",
+  "private": "true",
+  "version": "0.0.0",
+  "description": "for tesing asyncLocalStorage",
+  "main": "asyncLocalStorageExample",
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1"
+  },
+  "dependencies": {
+    "lodash": "^4.17.21",
+    "mongodb-memory-server": "^6.9.6",
+    "uuid": "^8.3.2"
+  },
+  "repository": "",
+  "author": "",
+  "license": "BSD"
+}

lib/helpers/asyncLocalStorageWrapper.js

@@ -0,0 +1,34 @@
+'use strict';
+
+let AsyncResource;
+let executionAsyncId;
+let isSupported = false;
+
+try {
+  const asyncHooks = require('async_hooks');
+  if (
+    typeof asyncHooks.AsyncResource.prototype.runInAsyncScope === 'function'
+  ) {
+    AsyncResource = asyncHooks.AsyncResource;
+    executionAsyncId = asyncHooks.executionAsyncId;
+    isSupported = true;
+  }
+} catch (e) {
+  console.log('async_hooks does not support');
+}
+
+module.exports.wrap = function(callback) {
+  if (isSupported && typeof callback === 'function') {
+    const asyncResource = new AsyncResource('mongoose', executionAsyncId());
+    return function() {
+      try {
+        // asyncResource.runInAsyncScope(callback, this, ...arguments);
+        const params = [callback, this].concat(Array.from(arguments));
+        asyncResource.runInAsyncScope.apply(asyncResource, params);
+      } finally {
+        asyncResource.emitDestroy();
+      }
+    };
+  }
+  return callback;
+};

lib/helpers/query/wrapThunk.js

@@ -9,10 +9,13 @@
  * This function defines common behavior for all query thunks.
  */
 
+const asyncLocalStorageWrapper = require('../../helpers/asyncLocalStorageWrapper');
+
 module.exports = function wrapThunk(fn) {
   return function _wrappedThunk(cb) {
     ++this._executionCount;
-
+    // wrap callback with AsyncResource
+    cb = asyncLocalStorageWrapper.wrap(cb);
     fn.call(this, cb);
   };
 };

lib/model.js

@@ -48,6 +48,7 @@ const parallelLimit = require('./helpers/parallelLimit');
 const removeDeselectedForeignField = require('./helpers/populate/removeDeselectedForeignField');
 const util = require('util');
 const utils = require('./utils');
+const asyncLocalStorageWrapper = require('./helpers/asyncLocalStorageWrapper');
 
 const VERSION_WHERE = 1;
 const VERSION_INC = 2;
@@ -225,6 +226,9 @@ function _applyCustomWhere(doc, where) {
  */
 
 Model.prototype.$__handleSave = function(options, callback) {
+  // wrap callback with AsyncResource
+  callback = asyncLocalStorageWrapper.wrap(callback);
+
   const _this = this;
   let saveOptions = {};