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

[blexrob]

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?

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]

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

[unilynx]

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