Post-hooks do not work with AsyncLocalStorage

[jonathan-wilkinson]

Issue: Bug
Behaviour:

When running a Mongoose query inside of an AsyncLocalStorage context the correct context is not picked up inside post-type hooks. But it is correctly picked up inside pre-type hooks.

Reproduction

I couldn't load this reproduction into REPL (it didn't like installing mongoose memory server). So I've copied and pasted below.

This script creates two models: Model and TwinModel. After Model has saved (post-save hook) it saves aTwinModel that duplicates its data. Both models use a post-save hook to log the ids of their respective documents into an AsyncLocalStorage context.

The script then runs two functions within that same AsyncLocalStorage context. It then logs the result of that context. The expectation is that each context should have references to the two models (one Model and one TwinModel created within those contexts).

// package.json
{
  "dependencies": {
    "mongodb-memory-server": "^6.9.6",
    "mongoose": "^5.12.0"
  }
}

// index.js
// @ts-check
const mongoose = require("mongoose")
const { MongoMemoryServer} = require("mongodb-memory-server");
const { AsyncLocalStorage } = require("async_hooks");

(async () => {
    /**
     * SETUP CTX
     */
    const CTX = new AsyncLocalStorage();

    /**
     * SETUP DATABASE
     */
    const mongod = new MongoMemoryServer();
    const uri = await mongod.getUri();

    /**
     * CONNECT MONGOOSE
     */
    await mongoose.connect(uri, { useNewUrlParser: true, useUnifiedTopology: true });

    /**
     * CREATE SCHEMAS
     */
    const schema = new mongoose.Schema({ name: String });
    const twinSchema = new mongoose.Schema({ twins: String });

    // Hook that pushes ID of doc to `CTX`
    const pushToCtx = function () {
        const name =this.name || this.twins;
        const type = this.name ? 'Obj' : 'Twin';
        const id = this.id.slice(-4);
        console.log(`Pushing ${type} ${name} ${id} to CTX`);

        const store = CTX.getStore();
        store.ids = (store.ids || []).concat(id);
    };

    schema.post('save', function () {
        return new TwinModel({ twins: this.name }).save();
    });
    schema.post('save', pushToCtx);
    twinSchema.post('save', pushToCtx);
    
    
    const Model = mongoose.model('Model', schema);
    const TwinModel = mongoose.model('TwinModel', twinSchema);

    // TEST

    await CTX.run({ ctxId: 'a'}, async () => {
        const a = new Model({ name: 'a' });
        await a.save();    

        const store = CTX.getStore();
        console.log("A", store);
    })

    await CTX.run({ ctxId: 'b'}, async () => {
        const b = new Model({ name: 'b' });
        await b.save();    

        const store = CTX.getStore();
        console.log("B", store);
    });



    process.exit()
})()

Expected Behaviour

When I run a Mongoose query inside of an AsyncLocalStorage context I can retrieve the correct context from inside any of the downstream function calls including any Mongoose hooks.

Versions

Node: v12.21.0 / v14.8.0
Mongoose: 5.12.0

[IslandRhythms]

const mongoose = require('mongoose');
const { MongoMemoryServer} = require("mongodb-memory-server");
const { AsyncLocalStorage } = require("async_hooks");

(async () => {
    /**
     * SETUP CTX
     */
    const CTX = new AsyncLocalStorage();

    /**
     * SETUP DATABASE
     */
    const mongod = new MongoMemoryServer();
    const uri = await mongod.getUri();

    /**
     * CONNECT MONGOOSE
     */
    await mongoose.connect(uri, { useNewUrlParser: true, useUnifiedTopology: true });

    /**
     * CREATE SCHEMAS
     */
    const schema = new mongoose.Schema({ name: String });
    const twinSchema = new mongoose.Schema({ twins: String });

    // Hook that pushes ID of doc to `CTX`
    const pushToCtx = function () {
        const name =this.name || this.twins;
        const type = this.name ? 'Obj' : 'Twin';
        const id = this.id.slice(-4);
        console.log(`Pushing ${type} ${name} ${id} to CTX`);

        const store = CTX.getStore();
        store.ids = (store.ids || []).concat(id);
    };

    schema.post('save', function () {
        console.log('Hello');
        return new TwinModel({ twins: this.name }).save();
    });
    schema.pre('save', function() {
        console.log('Hi');
    });
    schema.pre('save', pushToCtx);
    twinSchema.pre('save', pushToCtx);
    schema.post('save', pushToCtx);
    twinSchema.post('save', pushToCtx);
    
    
    const Model = mongoose.model('Model', schema);
    const TwinModel = mongoose.model('TwinModel', twinSchema);

    // TEST

    await CTX.run({ ctxId: 'a'}, async () => {
        const a = new Model({ name: 'a' });
        await a.save();    

        const store = CTX.getStore();
        console.log("A", store);
    })

    await CTX.run({ ctxId: 'b'}, async () => {
        const b = new Model({ name: 'b' });
        await b.save();    

        const store = CTX.getStore();
        console.log("B", store);
    });
    await mongoose.connection.dropDatabase();
    process.exit()
})()

Output:
Hi
Pushing Obj a 9868 to CTX
Hello
Pushing Twin a 9869 to CTX
Pushing Twin a 9869 to CTX
Pushing Obj a 9868 to CTX
A { ctxId: 'a', ids: [ '9868', '9869', '9869', '9868' ] }
Hi
Pushing Obj b 986a to CTX
Hello
Pushing Twin b 986b to CTX
Pushing Twin b 986b to CTX
Pushing Obj b 986a to CTX
B { ctxId: 'b', ids: [ '986a' ] }

I am on Node 15.7.0
Please provide what your output with your script is.

[jonathan-wilkinson]

@IslandRhythms

This is the output of my script:

Pushing Twin a 80a7 to CTX
Pushing Obj a 80a6 to CTX
A { ctxId: 'a', ids: [ '80a7', '80a6' ] }
Pushing Twin b 80a9 to CTX
Pushing Obj b 80a8 to CTX
B { ctxId: 'b' }

The first two objects are pushed into the correct context . But the second two are not.

But if we rerun this with post-hooks changed to pre hooks we get this:

Pushing Twin a 10eb to CTX
Pushing Obj a 10ea to CTX
A { ctxId: 'a', ids: [ '10eb', '10ea' ] }
Pushing Twin b 10ed to CTX
Pushing Obj b 10ec to CTX
B { ctxId: 'b', ids: [ '10ed', '10ec' ] }

With pre-hooks all objects are pushed into the correct context.

[jonathan-wilkinson]

I've been trying to dig into this myself and I think I've found the source of the problem: AsyncLocalStorage gets lost when using callbacks on the mongo-driver.

To test, if you change https://github.com/Automattic/mongoose/blob/master/lib/model.js#L274 to use promises (as below), it works.

    this[modelCollectionSymbol]
      .insertOne(obj, saveOptions)
      .then(ret => {
        callback(null, ret);
      })
      .catch(err => {
        if (err) {
          _setIsNew(_this, true);

          callback(err, null);
        }
      });

It looks like a bug has been raised but rejected as "won't fix" on the MongoDB driver: https://jira.mongodb.org/browse/NODE-3010 .

If there's no objections to making the change, I'm happy to make a PR.

[IslandRhythms]

@vkarpov15 I will have to defer to you on this.

[vkarpov15]

We'll have to look into how async-local-storage works. We considered converting to using promises under the hood for a similar issue with a different npm module in #7292 , but decided against it since we found an easy fix. We'll see if we can do something similar here.

[jonathan-wilkinson]

@vkarpov15

I took a look at that linked issue. The same problem happens using the express-http-context library/workaround (it relies on the same async_hooks module under the hood). It just wasn't visible from the tests in that issue because it only uses pre hooks (i.e. should execute before the context is lost by the mongo lib) and only runs the request once.

I can’t imagine that there will be an easy fix here as the context is getting lost in the internals of the mongo lib.

One other option would be to explicitly bind the callback to the correct context. But that seems like the worst of all options.

    const AsyncResource = require("async_hooks”).AsyncResource; // || some shim ;
    const asyncResource = new AsyncResource('save');

    this[modelCollectionSymbol].insertOne(obj, saveOptions, function(err, ret) {
      asyncResource.runInAsyncScope(function() {  
        if (err) {
          _setIsNew(_this, true);

          callback(err, null);
          return;
        }

        callback(null, ret);

      }, this);
    });

[andreialecu]

I've been struggling with this today.

Trying to add instrumentation via a plugin, and I need to pass some state to the post hook but it gets lost.

@jonathan-wilkinson have you been able to find any workaround that doesn't involve patching mongoose?

Edit: Ah, just thought of this now, and it solves my problem:

The pre hook works, so I can set something like this._context = storage.getStore(); and then I can read it from this._context in the post hook.

[royalrover]

the same question, fixed use the method above @andreialecu

[Uzlopak]

Maybe connected to mongoosejs/kareem#26

[vkarpov15]

@royalrover have you tried #7292 (comment) ?

[vkarpov15]

Making a note to try this out again given recent changes in the MongoDB node driver. MongoDB node driver dropped callback support, so this may be fixed

[vkarpov15]

Looks like this is fixed. When running the following script against Mongoose 7:

const mongoose = require('mongoose');
const { AsyncLocalStorage } = require("async_hooks");

(async () => {
    /**
     * SETUP CTX
     */
    const CTX = new AsyncLocalStorage();

    /**
     * SETUP DATABASE
     */
    const uri = 'mongodb://127.0.0.1:27017/mongoose_test';

    /**
     * CONNECT MONGOOSE
     */
    await mongoose.connect(uri, { useNewUrlParser: true, useUnifiedTopology: true });

    /**
     * CREATE SCHEMAS
     */
    const schema = new mongoose.Schema({ name: String });
    const twinSchema = new mongoose.Schema({ twins: String });

    // Hook that pushes ID of doc to `CTX`
    const pushToCtx = function () {
        const name =this.name || this.twins;
        const type = this.name ? 'Obj' : 'Twin';
        const id = this.id.slice(-4);
        console.log(`Pushing ${type} ${name} ${id} to CTX`);

        const store = CTX.getStore();
        store.ids = (store.ids || []).concat(id);
    };

    schema.post('save', function () {
        console.log('Hello');
        return new TwinModel({ twins: this.name }).save();
    });
    schema.pre('save', function() {
        console.log('Hi');
    });
    schema.pre('save', pushToCtx);
    twinSchema.pre('save', pushToCtx);
    schema.post('save', pushToCtx);
    twinSchema.post('save', pushToCtx);


    const Model = mongoose.model('Model', schema);
    const TwinModel = mongoose.model('TwinModel', twinSchema);

    // TEST

    await CTX.run({ ctxId: 'a'}, async () => {
        const a = new Model({ name: 'a' });
        await a.save();

        const store = CTX.getStore();
        console.log("A", store);
    })

    await CTX.run({ ctxId: 'b'}, async () => {
        const b = new Model({ name: 'b' });
        await b.save();

        const store = CTX.getStore();
        console.log("B", store);
    });
    await mongoose.connection.dropDatabase();
    process.exit()
})()

I get the following output:

Hi
Pushing Obj a 9390 to CTX
Hello
Pushing Twin a 9392 to CTX
Pushing Twin a 9392 to CTX
Pushing Obj a 9390 to CTX
A { ctxId: 'a', ids: [ '9390', '9392', '9392', '9390' ] }
Hi
Pushing Obj b 9394 to CTX
Hello
Pushing Twin b 9396 to CTX
Pushing Twin b 9396 to CTX
Pushing Obj b 9394 to CTX
B { ctxId: 'b', ids: [ '9394', '9396', '9396', '9394' ] }

Correct ids are showing up on the correct context.