穿针引线之 AsyncLocalStorage

在 Node.js 中，如何更优雅地获取请求上下文一直是一个问题，看一下下面的例子。

背景

const http = require('http');
function handler1(req, res) {
    console.log(req.url);
}

function handler2(req, res) {
     console.log(req.url);
}

http.createServer((req, res) => {
    handler1(req, res);
    handler2(req, res);
    res.end();
}).listen();

上面的例子中，每次收到一个请求时都会执行 handler1 和 handler2，为了在不同的地方里都能拿到请求上下文，我们只能逐级进行传递，如果业务逻辑很复杂，这个维护性是非常差的，下接下来看看如何使用 AsyncLocalStorage 解决这个问题。

AsyncLocalStorage

AsyncLocalStorage 是基于 Async Hooks 实现的，它通过上下文传递实现了异步代码的上下文共享和隔离。下面看一个例子。

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

const asyncLocalStorage = new AsyncLocalStorage();
function logWithId(msg) {
  const id = asyncLocalStorage.getStore();
  console.log(`${id !== undefined ? id : '-'}:`, msg);
}

asyncLocalStorage.run(1, () => {
    logWithId('start');
    setImmediate(() => {
      logWithId('finish');
    });
 });

上面的代码会输出

1: start
1: finish

从中可以看到两个 logWithId 共享了同一个上下文，这个上下文是由 run 函数设置的 1，那这种技术如何解决我们刚开始提出的问题呢？看一下下面的例子。

const http = require('http');
const { AsyncLocalStorage } = require('async_hooks');

const asyncLocalStorage = new AsyncLocalStorage();

function handler1() {
    const { req } = asyncLocalStorage.getStore();
    console.log(req.url);
}

function handler2() {
    setImmediate(() => {
        const { req } = asyncLocalStorage.getStore();
        console.log(req.url);
    });
}

http.createServer((req, res) => {
    asyncLocalStorage.run({ req, res }, () => {   
        handler1();
        handler2();
    });
    res.end();
}).listen(9999, () => {
    http.get({ port: 9999, path: '/test' })
});

执行上面代码输出如下。

/test
/test

可以看到，我们不需要逐级地传递请求上下文并且可以在任意异步代码中获取请求上下文。这让代码的编写和维护带来了非常大的好处，不过缺点就是，因为 AsyncLocalStorage 是基于 Async hooks 的，所以会带来一些性能损耗，不同的版本可能不一样，但是 Node.js 也在不断地优化其性能，我印象中，社区有人提过使用其他技术实现 AsyncLocalStorage。

AsyncLocalStorage 原理

知其然知其所以然，只知道怎么使用是不够的，理解其原理可以帮助我们更好地使用它。下面来分析一下 AsyncLocalStorage 的原理。先看一下创建 AsyncLocalStorage 的逻辑

class AsyncLocalStorage {
  constructor() {
    this.kResourceStore = Symbol('kResourceStore');
    this.enabled = false;
  }
}

创建AsyncLocalStorage的时候很简单，主要是置状态为false，并且设置kResourceStore的值为Symbol(‘kResourceStore’)。设置为Symbol(‘kResourceStore’)而不是 ‘kResourceStore’ 很重要，我们后面会看到。继续看一下执行AsyncLocalStorage.run的逻辑。

const storageList = [];
const storageHook = createHook({
  init(asyncId, type, triggerAsyncId, resource) {
    const currentResource = executionAsyncResource();
    // 传递上下文
    for (let i = 0; i < storageList.length; ++i) {
      storageList[i]._propagate(resource, currentResource, type);
    }
  },
});

run(store, callback, ...args) {
    // 把当前 AsyncLocalStorage 加入队列
    ArrayPrototypePush(storageList, this);
    // 启动 AsyncHooks
    storageHook.enable();
   // 获取当前的异步资源，比如收到的请求
   const resource = executionAsyncResource();
   // 记录旧的上下文
   const oldStore = resource[this.kResourceStore];
   // 修改当前异步资源的上下文
   resource[this.kResourceStore] = store;
   // 在新的上下文中执行传入的回调函数
   try {
     return ReflectApply(callback, null, args);
   } finally {
     resource[this.kResourceStore] = oldStore;
   }
 }

回调函数里可以通过 asyncLocalStorage.getStore() 获得设置的公共上下文。

getStore() {
  const resource = executionAsyncResource();
  return resource[this.kResourceStore];
}

getStore的原理很简单，首先拿到当前的异步资源，然后根据AsyncLocalStorage的kResourceStore的值从resource中拿到公共上下文，如果是同步执行getStore（比如 handler1 中），那么executionAsyncResource返回的就是我们请求所对应的异步资源，上下文就是在run时设置的上下文({req, res})，但是如果是异步getStore那么怎么办呢？因为这时候executionAsyncResource返回的不再是请求所对应的异步资源，也就拿不到他挂载的公共上下文。为了解决这个问题，Node.js对公共上下文进行了传递。

const storageList = []; // AsyncLocalStorage对象数组
const storageHook = createHook({
  init(asyncId, type, triggerAsyncId, resource) {
    const currentResource = executionAsyncResource();
    for (let i = 0; i < storageList.length; ++i) {
      storageList[i]._propagate(resource, currentResource);
    }
  }
});

 _propagate(resource, triggerResource) {
    const store = triggerResource[this.kResourceStore];
    resource[this.kResourceStore] = store;
  }

我们看到在每次资源创建的时候，Node.js会把当前异步资源的上下文挂载到新创建的异步资源中。所以在asyncLocalStorage.getStore() 时即使不是我们在执行run时创建的资源对象，也可以获得具体asyncLocalStorage对象所设置的资源（ handler2 中）。关系图如下。

这样就实现了异步资源上下文的共享和隔离。

总结

AsyncLocalStorage 有很多用法和用处，我们在 Node.js APM 中也大量用到该技术，通过 AsyncLocalStorage，我们可以无侵入地实现对 Node.js 应用的内部进行观测，时间关系，本文简单地介绍了 AsyncLocalStorage 的使用和原理，有兴趣的同学可以自行探索。