From 768c848f4eeb77d76062b126b18ebf6f43209eb5 Mon Sep 17 00:00:00 2001 From: Wu Jian Ping Date: Thu, 6 May 2021 15:02:26 +0800 Subject: [PATCH] #10020: asyncLocalStorage supporting --- .../asyncLocalStorageExample.js | 136 ++++++++++++++++++ examples/asyncLocalStorage/callContext.js | 19 +++ examples/asyncLocalStorage/package.json | 18 +++ lib/helpers/asyncLocalStorageWrapper.js | 34 +++++ lib/helpers/query/wrapThunk.js | 5 +- lib/model.js | 4 + 6 files changed, 215 insertions(+), 1 deletion(-) create mode 100644 examples/asyncLocalStorage/asyncLocalStorageExample.js create mode 100644 examples/asyncLocalStorage/callContext.js create mode 100644 examples/asyncLocalStorage/package.json create mode 100644 lib/helpers/asyncLocalStorageWrapper.js diff --git a/examples/asyncLocalStorage/asyncLocalStorageExample.js b/examples/asyncLocalStorage/asyncLocalStorageExample.js new file mode 100644 index 00000000000..97fc4672c51 --- /dev/null +++ b/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(); diff --git a/examples/asyncLocalStorage/callContext.js b/examples/asyncLocalStorage/callContext.js new file mode 100644 index 00000000000..da293d8a17b --- /dev/null +++ b/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; diff --git a/examples/asyncLocalStorage/package.json b/examples/asyncLocalStorage/package.json new file mode 100644 index 00000000000..4998881e41c --- /dev/null +++ b/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" +} diff --git a/lib/helpers/asyncLocalStorageWrapper.js b/lib/helpers/asyncLocalStorageWrapper.js new file mode 100644 index 00000000000..207913446ff --- /dev/null +++ b/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; +}; diff --git a/lib/helpers/query/wrapThunk.js b/lib/helpers/query/wrapThunk.js index 0005c330c44..d203a430409 100644 --- a/lib/helpers/query/wrapThunk.js +++ b/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); }; }; \ No newline at end of file diff --git a/lib/model.js b/lib/model.js index 3bb4f38ede7..0afe7ef109c 100644 --- a/lib/model.js +++ b/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 = {};