有趣的 Async hooks 模块

在 Node.js 中，Async hooks 是一个非常有意思且强大的模块（虽然性能上存在一些问题），在 APM 中，我们可以借助这个模块做很多事情。本文介绍两个有趣的用法。

AsyncLocalStorage

在 Node.js 中，上下文传递一直是一个非常困难的问题，Node.js 通过 AsyncLocalStorage 提供了一种解决方案，今天看到一个库中实现了类似 AsyncLocalStorage 的能力，还挺有意思的。代码如下。

class ALS {
    constructor() {
        this._contexts = new Map();
        this._stack = [];
        this.hook = createHook({
            init: this._init.bind(this),
            before: this._before.bind(this),
            after: this._after.bind(this),
            destroy: this._destroy.bind(this),
        });
    }

    context() {
        return this._stack[this._stack.length - 1];
    }

    run(context, fn, thisArg, ...args) {
        this._enterContext(context);
        try {
            return fn.call(thisArg, ...args);
        }
        finally {
            this._exitContext();
        }
    }

    enable() {
        this.hook.enable();
    }

    disable() {
        this.hook.disable();
        this._contexts.clear();
        this._stack = [];
    }

    _init(asyncId) {
        const context = this._stack[this._stack.length - 1];
        if (context !== undefined) {
            this._contexts.set(asyncId, context);
        }
    }
    
    _destroy(asyncId) {
        this._contexts.delete(asyncId);
    }
    
    _before(asyncId) {
        const context = this._contexts.get(asyncId);
        this._enterContext(context);
    }
    
    _after() {
        this._exitContext();
    }
    
    _enterContext(context) {
        this._stack.push(context);
    }
    
    _exitContext() {
        this._stack.pop();
    }
}

这个方式是基于 Async hooks 实现的，原理是在 init 钩子中获取当前的上下文，然后把当前的上下文传递到当前创建的异步资源的，接着在执行异步资源回调前，Node.js 会执行 before 钩子，before 钩子中会把当前异步资源（正在执行回调的这个资源）的上下文压入栈中，然后在回调里就可以通过 context 函数获取到当前的上下文，实际上获取的就是刚才压入栈中的内容，执行完回调后再出栈。前面介绍了其工作原理，主要是实现异步资源的上下文传递且在执行回调时通过栈的方式实现了上下文的管理，那么第一个上下文是如何来的呢？答案是通过 run 函数（Node.js 中还可以通过 enterWith），run 会把用户设置的上下文压入栈中，然后执行了一个用户传入的函数，如果这个函数中创建了异步资源，那么用户传入的上下文就会传递到这个新创建的异步资源中，后面执行这个异步资源的回调时，就可以拿到对应的上下文了。接着看一下使用效果。

const als = new ALS();
als.enable();

http.createServer((req, res) => {
    als.run(req, () => {
        setImmediate(() => {
            console.log(als.context().url);
            res.end();
        });
    })
}).listen(9999, () => {
    http.get({ port: 9999, host: '127.0.0.1' });
});

执行上面代码会输出 /。可以看到在 setImmediate 的回调中（setImmediate 会创建一个异步资源）成功拿到了 run 时设置的上下文。

监控异步回调的耗时

在 Node.js 中，代码执行耗时是一个非常值得关注的地方，Node.js 也提供了很多手段采集代码执行的耗时信息，下面介绍的是基于 Async hooks 实现的回调函数耗时监控。

const { createHook } = require('async_hooks');
const fs = require('fs');

const map = {};

createHook({
    init: (asyncId) => {
        map[asyncId] = { stack: new Error().stack };
    },
    before: (asyncId) => {
        if (map[asyncId]) {
            map[asyncId].start = Date.now();
        }
    },
    after: (asyncId) => {
        if (map[asyncId]) {
            fs.writeFileSync(1, `callback cost: ${Date.now() - map[asyncId].start}, stack: ${map[asyncId].stack}`);
        }
    },
    destroy: (asyncId) => {
        delete map[asyncId];
    }
}).enable();

setTimeout(() => {
    for (let i = 0; i < 1000000000; i++) {

    }
});

实现原理非常简单，主要是利用 before 和 after 钩子实现了回调的耗时统计，就不多介绍，社区中也有同学实现了这个能力，具体可以参考 https://github.com/naugtur/blocked-at/tree/master。