feat(js-runtime): async context (#684)

* feat(js-runtime): async context

* feat(docs): add docs

* fix(runtime): tests

* chore: add changeset

.changeset/grumpy-meals-repair.md

@@ -0,0 +1,7 @@
+---
+'@lagon/cli': patch
+'@lagon/docs': patch
+'@lagon/js-runtime': patch
+---
+
+Add `AsyncLocalStorage` & `AsyncContext` APIs

crates/runtime/tests/async_context.rs

@@ -0,0 +1,215 @@
+use lagon_runtime_http::{Request, Response, RunResult};
+use lagon_runtime_isolate::options::IsolateOptions;
+
+mod utils;
+
+// Tests ported from https://github.com/tc39/proposal-async-context/blob/master/tests/async-context.test.ts
+#[tokio::test]
+async fn inital_undefined() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "const ctx = new AsyncContext();
+const actual = ctx.get();
+
+if (actual !== undefined) {
+    throw new Error('Expected undefined');
+}
+
+export function handler() {
+    return new Response(actual === undefined);
+}"
+        .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from("true"))
+    );
+}
+
+#[tokio::test]
+async fn return_value() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "const ctx = new AsyncContext();
+const expected = { id: 1 };
+const actual = ctx.run({ id: 2 }, () => expected);
+
+if (actual !== expected) {
+    throw new Error('Expected expected');
+}
+
+export function handler() {
+    return new Response();
+}"
+        .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from(""))
+    );
+}
+
+#[tokio::test]
+async fn get_returns_current_context_value() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "const ctx = new AsyncContext();
+const expected = { id: 1 };
+
+ctx.run(expected, () => {
+    if (ctx.get() !== expected) {
+        throw new Error('Expected expected');
+    }
+});
+
+export function handler() {
+    return new Response();
+}"
+        .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from(""))
+    );
+}
+
+#[tokio::test]
+#[serial_test::serial]
+async fn get_within_nesting_contexts() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "const ctx = new AsyncContext();
+const first = { id: 1 };
+const second = { id: 2 };
+
+ctx.run(first, () => {
+    if (ctx.get() !== first) {
+        throw new Error('Expected first');
+    }
+    ctx.run(second, () => {
+        if (ctx.get() !== second) {
+            throw new Error('Expected second');
+        }
+    });
+    if (ctx.get() !== first) {
+        throw new Error('Expected first');
+    }
+});
+
+if (ctx.get() !== undefined) {
+    throw new Error('Expected undefined');
+}
+
+export function handler() {
+    return new Response();
+}"
+        .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from(""))
+    );
+}
+
+#[tokio::test]
+#[serial_test::serial]
+async fn get_within_nesting_different_contexts() {
+    utils::setup();
+    let (send, receiver) = utils::create_isolate(IsolateOptions::new(
+        "const a = new AsyncContext();
+const b = new AsyncContext();
+const first = { id: 1 };
+const second = { id: 2 };
+
+a.run(first, () => {
+    if (a.get() !== first) {
+        throw new Error('Expected first');
+    }
+    if (b.get() !== undefined) {
+        throw new Error('Expected undefined');
+    }
+    b.run(second, () => {
+        if (a.get() !== first) {
+            throw new Error('Expected first');
+        }
+        if (b.get() !== second) {
+            throw new Error('Expected second');
+        }
+    });
+    if (a.get() !== first) {
+        throw new Error('Expected first');
+    }
+    if (b.get() !== undefined) {
+        throw new Error('Expected undefined');
+    }
+});
+if (a.get() !== undefined) {
+    throw new Error('Expected undefined');
+}
+if (b.get() !== undefined) {
+    throw new Error('Expected undefined');
+}
+
+export function handler() {
+    return new Response();
+}"
+        .into(),
+    ));
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from(""))
+    );
+}
+
+#[tokio::test]
+#[serial_test::serial]
+async fn timers() {
+    utils::setup();
+    let log_rx = utils::setup_logger();
+    let (send, receiver) = utils::create_isolate(
+        IsolateOptions::new(
+            "const store = new AsyncLocalStorage();
+let id = 1;
+export async function handler() {
+    const result = store.run(id++, () => {
+        setTimeout(() => {
+            console.log(store.getStore() * 2);
+        }, 100);
+
+        return store.getStore() * 2;
+    });
+    // Make sure the console.log is executed before returning the response
+    await new Promise((resolve) => setTimeout(resolve, 150));
+
+    return new Response(result);
+}"
+            .into(),
+        )
+        .metadata(Some(("".to_owned(), "".to_owned()))),
+    );
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from("2"))
+    );
+
+    send(Request::default());
+
+    assert_eq!(
+        receiver.recv_async().await.unwrap(),
+        RunResult::Response(Response::from("4"))
+    );
+
+    assert_eq!(log_rx.recv_async().await.unwrap(), "2");
+}

crates/runtime/tests/timers.rs

@@ -179,29 +179,6 @@ async fn queue_microtask() {
     );
 }
 
-#[tokio::test]
-async fn queue_microtask_throw_not_function() {
-    utils::setup();
-    let (send, receiver) = utils::create_isolate(
-        IsolateOptions::new(
-            "export async function handler() {
-    queueMicrotask(true);
-    return new Response('Hello world');
-}"
-            .into(),
-        )
-        .metadata(Some(("".to_owned(), "".to_owned()))),
-    );
-    send(Request::default());
-
-    assert_eq!(
-        receiver.recv_async().await.unwrap(),
-        RunResult::Error(
-            "Uncaught TypeError: Parameter 1 is not of type 'Function'\n  at handler (2:5)".into()
-        )
-    );
-}
-
 #[tokio::test]
 #[serial]
 async fn timers_order() {

packages/docs/pages/runtime-apis.mdx

@@ -52,6 +52,17 @@ The standard `AbortController` object. [See the documentation on MDN](https://de
 
 The standard `AbortSignal` object. [See the documentation on MDN](https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal).
 
+### `AsyncContext`
+
+An early implementation of the [Async Context proposal](https://github.com/tc39/proposal-async-context). You shouldn't use this API yet, as it is still experimental and subject to change.
+
+### `AsyncLocalStorage`
+
+A minimal implementation of [Node.js's AsyncLocalStorage](https://nodejs.org/api/async_context.html#class-asynclocalstorage). The following methods are supported:
+
+- `getStore()`
+- `run(store, callback, ...args)`
+
 ### `Blob`
 
 The standard `Blob` object. [See the documentation on MDN](https://developer.mozilla.org/en-US/docs/Web/API/Blob).

packages/js-runtime/src/index.ts

@@ -5,6 +5,7 @@ import './runtime/encoding/base64';
 import './runtime/core';
 import './runtime/streams';
 import './runtime/abort';
+import './runtime/global/context';
 import './runtime/global/event';
 import './runtime/global/blob';
 import './runtime/global/file';
@@ -28,6 +29,29 @@ import './runtime/http/fetch';
 // NOTE: we use `var` to that we can refer to these variables
 // using `globalThis.VARIABLE`.
 declare global {
+  interface AsyncContextConstructor {
+    new (): AsyncContext;
+    wrap(callback: (...args: unknown[]) => void): (...args: unknown[]) => void;
+  }
+
+  interface AsyncContext<T = unknown> {
+    get(): T;
+    run<R>(store: T, callback: (...args: unknown[]) => R, ...args: unknown[]): R;
+  }
+
+  var AsyncContext: AsyncContextConstructor;
+
+  interface AsyncLocalStorageConstructor {
+    new (): AsyncLocalStorage;
+  }
+
+  interface AsyncLocalStorage<T = unknown> {
+    getStore(): T;
+    run<R>(store: T, callback: (...args: unknown[]) => R, ...args: unknown[]): R;
+  }
+
+  var AsyncLocalStorage: AsyncLocalStorageConstructor;
+
   var LagonSync: {
     log: (level: string, message: string) => void;
     pullStream: (id: number, done: boolean, chunk?: Uint8Array) => void;
@@ -73,6 +97,7 @@ declare global {
     TEXT_ENCODER: TextEncoder;
     TEXT_DECODER: TextDecoder;
   };
+  var __storage__: Map<AsyncContext, unknown>;
   var handler: (request: Request) => Promise<Response>;
   var masterHandler: (
     id: number,

packages/js-runtime/src/runtime/global/context.ts

@@ -0,0 +1,43 @@
+(globalThis => {
+  globalThis.__storage__ = new Map();
+
+  globalThis.AsyncContext = class {
+    get() {
+      return globalThis.__storage__.get(this);
+    }
+
+    static wrap(callback: (...args: unknown[]) => void): (...args: unknown[]) => void {
+      const snapshot = globalThis.__storage__;
+
+      return function (...args: unknown[]) {
+        const prev = globalThis.__storage__;
+        try {
+          globalThis.__storage__ = snapshot;
+          // @ts-expect-error we want to get this from the current function
+          return callback.apply(this, args);
+        } finally {
+          globalThis.__storage__ = prev;
+        }
+      };
+    }
+
+    run<R>(store: unknown, callback: (...args: unknown[]) => R, ...args: unknown[]): R {
+      const prev = globalThis.__storage__;
+
+      try {
+        const n = new Map(globalThis.__storage__);
+        n.set(this, store);
+        globalThis.__storage__ = n;
+        return callback(...args);
+      } finally {
+        globalThis.__storage__ = prev;
+      }
+    }
+  };
+
+  globalThis.AsyncLocalStorage = class extends AsyncContext {
+    getStore() {
+      return this.get();
+    }
+  };
+})(globalThis);

packages/js-runtime/src/runtime/global/timers.ts

@@ -11,7 +11,7 @@
     const id = counter++;
 
     timers.set(id, {
-      handler,
+      handler: AsyncContext.wrap(handler),
       repeat,
     });
 
@@ -46,6 +46,6 @@
   };
 
   globalThis.queueMicrotask = callback => {
-    LagonSync.queueMicrotask(callback);
+    LagonSync.queueMicrotask(AsyncContext.wrap(callback));
   };
 })(globalThis);