如何让 AsyncLocalStorage 像 ContextVar 一样工作？

Question

我正在尝试使用 Node 的

AsyncLocalStorage

API 和 Typescript 5.2 中新引入的

using

关键字来模拟 python 的

ContextVar

日志行为。

以下 python 代码按照我的预期完美运行：

# Output
# ------
# push {'fn': 'main', 'main': True}
# info {'msg': 'main: before', 'fn': 'main', 'main': True}
# push {'fn': 'a', 'a': True}
# push {'fn': 'b', 'b': True}
# info {'msg': 'a', 'fn': 'a', 'main': True, 'a': True}
# pop {'fn': 'a', 'a': True}
# info {'msg': 'b', 'fn': 'b', 'main': True, 'b': True}
# pop {'fn': 'b', 'b': True}
# info {'msg': 'main: after', 'fn': 'main', 'main': True}
# pop {'fn': 'main', 'main': True}

from contextvars import ContextVar
from contextlib import contextmanager
import asyncio


class Logger:
    def __init__(self) -> None:
        self.store: ContextVar[list] = ContextVar("store", default=[])

    @contextmanager
    def with_meta(self, meta):
        try:
            stack = self.store.get()
            print("push", meta)
            self.store.set([*stack, meta])
            yield
        finally:
            stack = self.store.get()
            print("pop", stack.pop())

    @property
    def meta(self):
        stack = self.store.get()
        return {k: v for d in stack for k, v in d.items()}

    def info(self, msg):
        print("info", {**{"msg": msg}, **self.meta})


logger = Logger()


async def a():
    with logger.with_meta({"fn": "a", "a": True}):
        await asyncio.sleep(0.2)
        logger.info("a")


async def b():
    with logger.with_meta({"fn": "b", "b": True}):
        await asyncio.sleep(0.2)
        logger.info("b")


async def main():
    with logger.with_meta({"fn": "main", "main": True}):
        logger.info("main: before")
        await asyncio.gather(a(), b())
        logger.info("main: after")


asyncio.run(main())

但是，当我将其转换为打字稿时，道具变得一团糟：

// Output:
// ------
// push { fn: 'main', main: true }
// info { msg: 'main: before', fn: 'main', main: true }
// push { fn: 'a', a: true }
// push { fn: 'b', b: true }
// info { msg: 'a', fn: 'b', main: true, a: true, b: true }
// pop { fn: 'b', b: true }
// info { msg: 'b', fn: 'a', main: true, a: true }
// pop { fn: 'a', a: true }
// info { msg: 'main: after', fn: 'main', main: true }
// pop { fn: 'main', main: true }

import { AsyncLocalStorage } from 'node:async_hooks'

function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms))
}


class Logger {
  private als: AsyncLocalStorage<any[]>
  constructor() {
    this.als = new AsyncLocalStorage()
    this.als.enterWith([])
  }

  withMeta(meta: any,) {
    const stack = this.als.getStore()!
    stack.push(meta)
    console.log('push', meta)
    return this
  }

  [Symbol.dispose]() {
    const stack = this.als.getStore()!
    console.log('pop', stack.pop())
  }

  get meta() {
    const stack = this.als.getStore()!
    return stack.reduce((prev, curr) => ({ ...prev, ...curr }), {})
  }

  info(msg: string) {
    console.log('info', { msg, ...this.meta })
  }
}

const logger = new Logger()

async function a() {
  using _ = logger.withMeta({ fn: 'a', a: true })
  await sleep(200)
  logger.info('a')
}

async function b() {
  using _ = logger.withMeta({ fn: 'b', b: true })
  await sleep(200)
  logger.info('b')
}

async function main() {
  using _ = logger.withMeta({ fn: 'main', main: true })
  logger.info('main: before')
  await Promise.all([a(), b()])
  logger.info('main: after')
}

main().catch(console.error)

例如，

a

的上下文与

b

的上下文混合在一起

ContextVar

与

AsyncLocalStorage

有何不同？看起来

AsyncLocalStorage

无论如何都在使用同一个存储，而 ContextVar 似乎在创建新任务/承诺时“分叉/复制”。

有没有办法在node中实现python行为？

Answer 1

JavaScript 版本代码中的第一个问题是，您只有一个所有异步任务共享的堆栈数组。因此，特别是，当一个异步任务将一个值推送到数组时，所有其他异步任务都会看到该值。⁰ 如果您将

self.store.set([*stack, meta])

替换为

stack.append(meta)

，您将在 Python 版本中得到相同的结果。在 JavaScript 版本中，这意味着您需要使用新构造的数组而不是

enterWith

来调用

stack.push

。

除非这也行不通。正如 enterWith

 的 

文档警告：

此转换将持续到整个同步执行。这意味着，例如，如果在事件处理程序中输入上下文，则后续事件处理程序也将在该上下文中运行，除非使用

AsyncResource

专门绑定到另一个上下文。这就是为什么

run()

应优先于

enterWith()

，除非有充分理由使用后一种方法。

这几乎就是您在这里遇到的情况。尽管您不编写同步事件处理程序，但请记住，在 JavaScript 中，异步函数会同步运行，直到第一个

await

点或直到它们返回（以较早者为准）；此时，它们会同步向调用者返回一个承诺。 （这与 Python 不同，在 Python 中，在调用

a

之前，

b

和

asyncio.gather

任务甚至不会开始运行。）只有

await

之后的延续实际上是异步执行的。

console.log(0);
(async () => {
  console.log(1);
  await 0;
  console.log(3);
})();
console.log(2);

最简单的解决方法是在

await 0

和

a

的开头注入

b

。更一般地，它应该出现在调用

withMeta

之前；它不能出现在调用者之外的任何其他地方。这是为了确保在新的、单独的异步任务中调用

enterWith

，并且不会传播到

withMeta

的调用者之外。

但是，它不仅在代码中看起来很尴尬，而且每次都必须记住这样做，这使得这种方法几乎站不住脚。您可能还放弃了 disposal-based API¹，而是将

run

方法应用于延续回调。这正是文档告诉您要做的事情。

（截至 2023 年 9 月）创建此 API 的标准版本（

AsyncContext

提案）的努力可能有一个很好的理由，其中包含

run

，但没有与

enterWith

等效的内容。

---

⁰ 另外，请注意在 Python 版本中，每次调用

with_meta

都有一个单独的上下文管理器（因此也有一个单独的处理程序），但在 JavaScript 版本中，您只需返回

this

。就其本身而言，它可能对正确性影响不大，但它确实意味着处置器不以任何方式与它应该弹出的堆栈元素相关联，并且无法检查它弹出与推送的相同元素的不变量。 Python版本也没有这样做，但是添加会容易得多。

¹ 首先，使用处置 API 似乎并不是一个好主意。考虑一下如果您使用

DisposableStack

会发生什么：

{
    using stack = new DisposableStack();
    {
        using _ = logger.withMeta({ p: 123 });
        stack.use(logger.withMeta({ q: 456 }));
        logger.info('foo');
    }
    logger.info('bar');
}

对于几乎任何其他类型的资源来说，这都是完全合法的处置 API 使用。如果它不适合您的 API，那么也许它一开始就不应该公开处置 API。