AsyncLocalStorage supporting in callback & promise & plugin

examples/asyncLocalStorage/asyncLocalStorageExample.js

@@ -0,0 +1,238 @@
+
+'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 contextData = callContext.get();
+
+    // verify asyncLocalStorage
+    if (this.name !== contextData.name) {
+      console.error("[static-hooks] [pre] [save]", this.name, contextData.name);
+    } else {
+      console.log("[OK] [static-hooks] [pre] [save]");
+    }
+  });
+
+  schema.post(["save"], function () {
+    const contextData = callContext.get();
+
+    // verify asyncLocalStorage
+    if (this.name !== contextData.name) {
+      console.error(
+        "[ERROR] [static-hooks] [post] [save]",
+        this.name,
+        contextData.name
+      );
+    } else {
+      console.log("[OK] [static-hooks] [post] [save]");
+    }
+  });
+};
+
+const pluginQuery = (schema) => {
+  schema.pre(["find", "findOne", "count", "countDocuments"], function () {
+    const contextData = callContext.get();
+
+    // verify asyncLocalStorage
+    if (this._conditions.name !== contextData.name) {
+      console.error(
+        `[ERROR] [static-hooks] [pre] [${this.op}]`,
+        this._conditions.name,
+        contextData.name
+      );
+    } else {
+      console.log(`[OK] [static-hooks] [pre] [${this.op}]`);
+    }
+  });
+
+  schema.post(["find", "findOne", "count", "countDocuments"], function () {
+    const contextData = callContext.get();
+
+    // verify asyncLocalStorage
+    if (this._conditions.name !== contextData.name) {
+      console.error(
+        `[ERROR] [static-hooks] [post] [${this.op}]`,
+        this._conditions.name,
+        contextData.name
+      );
+    } else {
+      console.log(`[OK] [static-hooks] [post] [${this.op}]`);
+    }
+  });
+};
+
+const pluginAggregate = (schema) => {
+  schema.pre(["aggregate"], function () {
+    // Special Case: aggregate should keep store
+    const contextData = callContext.get();
+    this.__asyncLocalStore = contextData;
+
+    const name = this._pipeline[0].$match.name;
+
+    // verify asyncLocalStorage
+    if (name !== contextData.name) {
+      console.error(
+        "[ERROR] [static-hooks] [pre] [aggregate]",
+        name,
+        contextData.name
+      );
+    } else {
+      console.log("[OK] [static-hooks] [pre] [aggregate]");
+    }
+  });
+
+  schema.post(["aggregate"], function () {
+    const contextData = this.__asyncLocalStore;
+    const name = this._pipeline[0].$match.name;
+
+    // verify asyncLocalStorage
+    if (name !== contextData.name) {
+      console.error(
+        "[ERROR] [static-hooks] [post] [aggregate]",
+        name,
+        contextData.name
+      );
+    } else {
+      console.log("[OK] [static-hooks] [post] [aggregate]");
+    }
+  });
+};
+
+mongoose.plugin(pluginSave);
+mongoose.plugin(pluginQuery);
+mongoose.plugin(pluginAggregate);
+
+let createCounter = 0;
+let findCallbackCounter = 0;
+let findPromiseCounter = 0;
+let aggregateCounter = 0;
+let countCounter = 0;
+
+const docCount = 50;
+
+const start = 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(async (resolve, reject) => {
+    for (let i = 0; i < docCount; ++i) {
+      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 === docCount) {
+        resolve();
+      }
+    }
+  });
+
+  for (let i = 0; i < docCount; ++i) {
+    setTimeout(async () => {
+      const name = names[i];
+      callContext.enter({ name });
+
+      // for testing callback
+      UserModel.find({ name }, (err, data) => {
+        ++findCallbackCounter;
+        data = data[0];
+        const contextData = callContext.get();
+
+        // verify asyncLocalStorage
+        if (data.name !== contextData.name) {
+          console.error(
+            `[ERROR] ${findCallbackCounter}: post-find-in-callback`,
+            data.name,
+            contextData.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 contextData = callContext.get();
+
+      // verify asyncLocalStorage
+      if (data.name !== contextData.name) {
+        console.error(
+          `[ERROR] ${findPromiseCounter}: post-find-in-promise`,
+          data.name,
+          contextData.name
+        );
+      } else {
+        console.log(`[OK] ${findPromiseCounter}: post-find-in-promise`);
+      }
+
+      // aggregate
+      UserModel.aggregate([{ $match: { name: name } }], (err, data) => {
+        const contextData = callContext.get();
+        data = data[0];
+
+        // verify asyncLocalStorage
+        if (data.name !== contextData.name) {
+          console.error(
+            `[ERROR] ${findCallbackCounter}: post-aggregate-in-callback`,
+            data.name,
+            contextData.name
+          );
+        } else {
+          console.log(
+            `[OK] ${findCallbackCounter}: post-aggregate-in-callback`
+          );
+        }
+        ++aggregateCounter;
+      });
+
+      await UserModel.countDocuments({ name }).exec();
+      ++countCounter;
+    }, _.random(10, 50));
+  }
+
+  const exit = () => {
+    if (
+      createCounter === docCount &&
+      findCallbackCounter === docCount &&
+      findPromiseCounter === docCount &&
+      aggregateCounter === docCount &&
+      countCounter === docCount
+    ) {
+      process.exit(0);
+    } else {
+      setTimeout(exit, 1000);
+    }
+  };
+
+  exit();
+};
+
+module.exports.start = start;

examples/asyncLocalStorage/callContext.js

@@ -0,0 +1,20 @@
+
+'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/index.js

@@ -0,0 +1,6 @@
+
+'use strict';
+
+const { start } = require('./asyncLocalStorageExample');
+
+start();

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,15 @@
  * 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 = {};
 
@@ -4838,6 +4842,9 @@ Model.$handleCallbackError = function(callback) {
     throw new MongooseError('Callback must be a function, got ' + callback);
   }
 
+  // wrap callback with AsyncResource
+  callback = asyncLocalStorageWrapper.wrap(callback);
+
   const _this = this;
   return function() {
     process.nextTick(() => {