otp: fix tracked values race

otp/js/e2e/helpers.ts

@@ -46,13 +46,85 @@ export type Result<T> =
 
 export { assert, expect };
 
+type PopcornHooksGlobal = typeof globalThis & {
+  popcorn: {
+    cleanups: number;
+    cleanup: () => void;
+    runJs: {
+      isPaused: () => boolean;
+      pause: () => Promise<void>;
+      finish: () => void;
+    };
+  };
+};
+
+async function addPopcornHooks(page: Page): Promise<void> {
+  await page.evaluate(() => {
+    let paused = false;
+    let finishRunJs!: () => void;
+    const resume = new Promise<void>((resolve) => {
+      finishRunJs = resolve;
+    });
+    const scope = globalThis as PopcornHooksGlobal;
+    scope.popcorn = {
+      cleanups: 0,
+      cleanup() {
+        this.cleanups += 1;
+      },
+      runJs: {
+        isPaused: () => paused,
+        pause: async () => {
+          paused = true;
+          await resume;
+          paused = false;
+        },
+        finish: finishRunJs,
+      },
+    };
+  });
+}
+
+export async function getCleanups(page: Page): Promise<number> {
+  return await page.evaluate(() => {
+    return (globalThis as PopcornHooksGlobal).popcorn.cleanups;
+  });
+}
+
+export async function waitForRunJsSuspension(page: Page): Promise<void> {
+  await page.waitForFunction(() => {
+    return (globalThis as PopcornHooksGlobal).popcorn.runJs.isPaused();
+  });
+}
+
+export async function finishRunJs(page: Page): Promise<void> {
+  await page.evaluate(() => {
+    (globalThis as PopcornHooksGlobal).popcorn.runJs.finish();
+  });
+}
+
+export function valuesWithKey<K extends PropertyKey>(
+  values: Iterable<unknown>,
+  key: K,
+): unknown[] {
+  return Array.from(values)
+    .filter((value): value is Record<K, unknown> => {
+      return (
+        typeof value === "object" &&
+        value !== null &&
+        Object.hasOwn(value, key)
+      );
+    })
+    .map((value) => value[key]);
+}
+
 export const test = base.extend<Fixtures>({
   page: async ({ page }, use) => {
     await page.goto("/");
     await page.waitForFunction(() => window.Popcorn !== undefined);
     await use(page);
   },
   createOtp: async ({ page }, use) => {
+    await addPopcornHooks(page);
     const handles = new Set<OtpHandle>();
     const createOtp = async (id = randomOtpId()) => {
       const otp = await OtpHandle.create(page, id);

otp/js/e2e/otp.spec.ts

@@ -1,4 +1,13 @@
-import { assert, expect, test, trimLeft } from "./helpers";
+import {
+  assert,
+  expect,
+  finishRunJs,
+  getCleanups,
+  test,
+  trimLeft,
+  valuesWithKey,
+  waitForRunJsSuspension,
+} from "./helpers";
 
 function evalOpts(code: string) {
   return {
@@ -317,6 +326,109 @@ test("nested tracked values", async ({ otp }) => {
   expect(otp.events).toContainEqual({ nested: 10 });
 });
 
+test("final tracked value sends are delivered before cleanup", async ({
+  otp,
+  page,
+}) => {
+  const count = 8;
+  const RUN_JS_BOOT_EVAL = trimLeft(`
+    spawn(fun() ->
+      lists:foreach(fun(I) ->
+        H = wasm:run_js(
+          <<"({i}) => new TrackedValue(
+            {label: 'tracked-' + i},
+            () => globalThis.popcorn.cleanup()
+          )">>,
+          #{i => I}
+        ),
+        ok = wasm:send(#{final_ref => H}),
+        erlang:garbage_collect(self())
+      end, lists:seq(1, ${count})),
+      ok = wasm:send(#{final_ref_done => true})
+    end).
+  `);
+  const boot = await otp.boot(evalOpts(RUN_JS_BOOT_EVAL));
+  assert(boot.ok);
+
+  await otp.waitForEvent("final_ref_done");
+
+  const refs = valuesWithKey(otp.events, "final_ref");
+  const labels = valuesWithKey(refs, "label");
+  expect(labels).toEqual(
+    Array.from({ length: count }, (_, i) => `tracked-${i + 1}`),
+  );
+  await expect.poll(() => getCleanups(page)).toBe(count);
+});
+
+test("tracked argument cleanup waits for async run_js to finish", async ({
+  otp,
+  page,
+}) => {
+  const RUN_JS_BOOT_EVAL = trimLeft(`
+    register(controller, self()),
+    Runner = spawn(fun() ->
+      receive
+        {tracked, H} ->
+          % Runner owns the final BEAM reference while JavaScript is suspended.
+          V = wasm:run_js(
+            <<"async ({h}) => {
+              await globalThis.popcorn.runJs.pause();
+              return {
+                value: h.value,
+                cleanups: globalThis.popcorn.cleanups,
+              };
+            }">>,
+            #{h => H},
+            [{timeout, 10000}]
+          ),
+          ok = wasm:send(#{async_tracked => V})
+      end
+    end),
+    spawn(fun() ->
+      H = wasm:run_js(
+        <<"() => new TrackedValue(
+          {value: 'tracked argument'},
+          () => globalThis.popcorn.cleanup()
+        )">>,
+        #{}
+      ),
+      Runner ! {tracked, H},
+      ok = wasm:send(#{async_runner_ready => true})
+    end),
+    receive
+      {wasm, PayloadJson, _} ->
+        #{<<"collect">> := true} = json:decode(PayloadJson),
+        % Force collection while Runner is blocked inside wasm:run_js/3.
+        erlang:garbage_collect(Runner),
+        ok = wasm:send(#{async_runner_collected => true})
+    end.
+  `);
+  const boot = await otp.boot(evalOpts(RUN_JS_BOOT_EVAL));
+  assert(boot.ok);
+
+  await otp.waitForEvent("async_runner_ready");
+  await waitForRunJsSuspension(page);
+
+  const collect = await otp.send("controller", { collect: true });
+  assert(collect.ok);
+  await otp.waitForEvent("async_runner_collected");
+
+  await expect
+    .poll(() => getCleanups(page), {
+      timeout: 500,
+      intervals: [50],
+    })
+    .toBe(0);
+
+  await finishRunJs(page);
+
+  const event = await otp.waitForEvent("async_tracked");
+  expect(event).toEqual({
+    async_tracked: { value: "tracked argument", cleanups: 0 },
+  });
+  await expect.poll(() => getCleanups(page)).toBe(1);
+});
+
 test("isolated tracked values", async ({ createOtp }) => {
   const evalFor = (id: string) =>
     trimLeft(`

otp/patches/0001-emscripten-support.patch

@@ -1029,10 +1029,10 @@ index 3b672dc41a554cf7d5440ccfd94e85bd44b61d35..4e5efb04e7eb04556ff14dc3ba8e0d83
  	    ]},
 diff --git a/erts/preloaded/src/wasm.erl b/erts/preloaded/src/wasm.erl
 new file mode 100644
-index 0000000000000000000000000000000000000000..9b79d3cfd08e70efc765a74ef07de90836f6854d
+index 0000000000000000000000000000000000000000..0e99092004c681b76ab1fe42831cf74197c06ee3
 --- /dev/null
 +++ b/erts/preloaded/src/wasm.erl
-@@ -0,0 +1,89 @@
+@@ -0,0 +1,99 @@
 +%%
 +%% %CopyrightBegin%
 +%%
@@ -1059,7 +1059,7 @@ index 0000000000000000000000000000000000000000..9b79d3cfd08e70efc765a74ef07de908
 +
 +%% on_load/0 is called explicitly from erl_init:start/2; do NOT use the
 +%% -on_load attribute (a preloaded module with one aborts the VM at boot).
-+-export([on_load/0, send/1, run_js/2, run_js/3]).
++-export([on_load/0, send/1, run_js/2, run_js/3, keep_alive/1]).
 +
 +-nifs([send_raw/1, materialize_ref/1, ref_key/1]).
 +
@@ -1068,7 +1068,9 @@ index 0000000000000000000000000000000000000000..9b79d3cfd08e70efc765a74ef07de908
 +-spec send(json:encode_value()) -> ok.
 +send(Message) ->
 +    Envelope = #{type => vm_message, data => Message},
-+    ok = send_raw(iolist_to_binary(json:encode(Envelope, fun encode_arg/2))).
++    ok = send_raw(iolist_to_binary(json:encode(Envelope, fun encode_arg/2))),
++    ?MODULE:keep_alive(Envelope),
++    ok.
 +
 +-spec run_js(binary(), map()) -> json:decode_value().
 +run_js(Code, Args) when is_binary(Code), is_map(Args) ->
@@ -1081,13 +1083,21 @@ index 0000000000000000000000000000000000000000..9b79d3cfd08e70efc765a74ef07de908
 +    ReplyTo = base64:encode(term_to_binary({self(), Token})),
 +    Envelope = #{type => run_js, code => Code, args => Args, reply_to => ReplyTo},
 +    ok = send_raw(iolist_to_binary(json:encode(Envelope, fun encode_arg/2))),
-+    receive
-+        {wasm, {Token, PayloadJson}, _MetaJson} ->
-+            handle_run_js_reply(PayloadJson)
-+    after Timeout ->
-+        error(run_js_timeout)
++    Reply = receive
++                {wasm, {Token, PayloadJson}, _MetaJson} -> PayloadJson
++            after Timeout ->
++                timeout
++            end,
++    ?MODULE:keep_alive(Envelope),
++    case Reply of
++        timeout -> error(run_js_timeout);
++        ReplyJson -> handle_run_js_reply(ReplyJson)
 +    end.
 +
++-spec keep_alive(Term) -> Term when Term :: term().
++keep_alive(Term) ->
++    Term.
++
 +handle_run_js_reply(PayloadJson) ->
 +    case json:decode(PayloadJson) of
 +        #{<<"ok">> := true, <<"value">> := Value} -> revive_refs(Value);