AsyncLocalStorage loses its store when using WASM JSPI after first asynchronous call

blexrob commented 7 months ago

Version

v21.6.2

Platform

Darwin Robs-Mac.local 23.3.0 Darwin Kernel Version 23.3.0: Wed Dec 20 21:28:58 PST 2023; root:xnu-10002.81.5~7/RELEASE_X86_64 x86_64

Subsystem

async_hooks

What steps will reproduce the bug?

Run the following code with node --experimental-wasm-stack-switching

const { AsyncLocalStorage } = require("node:async_hooks");

/* Tests that the AsyncLocalStore store is retained when WASM resumes
   after calling a JS function that returns a promise.
   Currently, the commandline switch `--experimental-wasm-stack-switching` is
   needed to run this code.
*/

async function testAsyncContextLoss() {
  /* The following code calls the import `m.import` twice, and returns the
     value returned by the second call.

     Test adapted from test TestStackSwitchGC2 from v8 test test/mjsunit/wasm/stack-switching.js

     Steps to generate the module bytecode
     - make a copy of the file /deps/v8/test/mjsunit/wasm/wasm-module-builder.js
     - append the following code

  let sig = makeSig([kWasmExternRef], [kWasmI32]);
  let builder = new WasmModuleBuilder();
  let import_index = builder.addImport('m', 'import', sig);
  builder.addFunction("test", sig)
      .addBody([
          kExprLocalGet, 0,
          kExprCallFunction, import_index,
          kExprDrop,
          kExprLocalGet, 0,
          kExprCallFunction, import_index,
      ]).exportFunc();
  console.log(builder.toBuffer());
  ```
 - run using node, copy the bytes in the following array

*/

const wasmModuleBinaryData = new Uint8Array([ 0, 97, 115, 109, 1, 0, 0, 0, 1, 11, 2, 96, 1, 111, 1, 127, 96, 1, 111, 1, 127, 2, 12, 1, 1, 109, 6, 105, 109, 112, 111, 114, 116, 0, 0, 3, 2, 1, 1, 7, 8, 1, 4, 116, 101, 115, 116, 0, 1, 10, 13, 1, 11, 0, 32, 0, 16, 0, 26, 32, 0, 16, 0, 11, 0, 14, 4, 110, 97, 109, 101, 1, 7, 1, 1, 4, 116, 101, 115, 116 ]);

const als = new AsyncLocalStorage;

// Imported asynchronous function, returns the .id value of the data stored in als const jsImport = new WebAssembly.Function({ parameters: ['externref'], results: ['i32'] }, () => { console.log(AsyncLocalStorage store called from wasm: + JSON.stringify(als.getStore())); return Promise.resolve(als.getStore()?.id ?? 0); }, { suspending: 'first' });

const wasmModule = new WebAssembly.Module(wasmModuleBinaryData); const wasmInstance = new WebAssembly.Instance(wasmModule, { 'm': { 'import': jsImport } });

const exportWrapper = new WebAssembly.Function({ parameters: [], results: ["externref"] }, wasmInstance.exports.test, { promising: 'first' });

const retval = await als.run({ id: 1 }, () => exportWrapper()); if (retval != 1) throw new Error(Lost the AsyncLocalStorage context, got ${retval}); }

testAsyncContextLoss();


### How often does it reproduce? Is there a required condition?

The error is 100% reproducible, no required conditions (other than enabling the `--experimental-wasm-stack-switching` CLI switch)

### What is the expected behavior? Why is that the expected behavior?

```console
Robs-Mac:_1 rob$ /usr/local/opt/node@20/bin/node -v
v20.11.1
Robs-Mac:_1 rob$ /usr/local/opt/node@20/bin/node --experimental-wasm-stack-switching jspi-asynccontext-loss.js 
AsyncLocalStorage store called from wasm: {"id":1}
AsyncLocalStorage store called from wasm: {"id":1}

The AsyncLocalStorage should retain the store after calling an asynchronous function from WASM.

What do you see instead?

Robs-Mac:_1 rob$ node -v
v21.6.2
Robs-Mac:_1 rob$ node --experimental-wasm-stack-switching jspi-asynccontext-loss.js 
AsyncLocalStorage store called from wasm: {"id":1}
AsyncLocalStorage store called from wasm: undefined
/private/tmp/_1/jspi-asynccontext-loss.js:68
    throw new Error(`Lost the AsyncLocalStorage context, got ${retval}`);
          ^

Error: Lost the AsyncLocalStorage context, got 0
    at testAsyncContextLoss (/private/tmp/_1/jspi-asynccontext-loss.js:68:11)
    at process.processTicksAndRejections (node:internal/process/task_queues:95:5)

Node.js v21.6.2

After the first call to an asynchronous function, the correct store isn't returned anymore.

Additional information

This used to work in node v20.

unilynx commented 5 months ago

retested with docker node:21 container (v21.7.2) - still an issue

unilynx commented 1 month ago

@blexrob I'm also still seeing it with v23.0.0-nightly202408194f94397650

nodejs / node