lift: settle an idle-exited async export on task.return, not thread drain

runtime/src/exec/boundary.ts

@@ -219,7 +219,7 @@
     stringEncoding: opts.stringEncoding,
     memory: opts.memory,
     realloc: opts.realloc === null ? null : (o, os, a, n) => {
       const realloc = require(opts.realloc, "realloc")!;
       const p = callCore(realloc, [o, os, a, n]);
       trapIf(p.length !== 1 || typeof p[0] !== "number", "realloc result");
       return (p[0] as number) >>> 0;
@@ -1690,6 +1690,12 @@
 
     let resolved: ComponentValue[] | null = null;
     let resolvedSeen = false;
+    /**
+     * One-shot: set only by `backgroundCompletion` below, fired here when the
+     * task resolves. See there for why the resolve callback — not thread
+     * drain — is the host's answer event (polyengine#313).
+     */
+    let onResolvedHook: (() => void) | null = null;
     const task = new Task(
       ft,
       taskOpts,
@@ -1699,6 +1705,11 @@
         resolved = result;
         resolvedSeen = true;
         stats.tasksResolved++;
+        if (onResolvedHook !== null) {
+          const f = onResolvedHook;
+          onResolvedHook = null;
+          f();
+        }
       },
     );
 
@@ -1884,14 +1895,24 @@
      * may have moved on since (polyengine#310).
      *
      * Resolution rides `finishHostEntry` unchanged — it already holds
-     * `completed`/`resultsToHost` — fired from `Task.onFinished`, i.e. the
-     * moment this task's last thread unregisters. That point is safe for the
-     * `done` predicate's task-scoped clauses: `resolvedSeen` is guaranteed
-     * (`unregisterThread`'s own `trapIf(state !== "resolved")` fires first
-     * otherwise), and this task's threads are gone from `store.awaiting`, so
-     * `midWasmCall()` is false by construction. `hopParked()` plays no part:
-     * a hop is the obligation of the driver that put it in flight (#280),
-     * never of a Promise waiting on another task.
+     * `completed`/`resultsToHost` — fired from THE TASK'S RESOLVE CALLBACK,
+     * i.e. `task.return` (polyengine#313). That is the reference's own answer
+     * event: definitions.py delivers a task's result to its caller through
+     * `on_resolve`, called from `Task.return_`, and run_tests.py's
+     * `lift_and_run` keeps ticking the store afterwards for OTHER work, not to
+     * produce the result. wasmtime's `call_concurrent` is the same shape.
+     *
+     * It used to fire on the task's LAST thread unregistering, which is not an
+     * event a callback-ABI task need ever reach: a guest that keeps spawned
+     * futures alive for the instance's life (wit-bindgen `spawn_local` — an
+     * event loop, a driver, an accept loop) leaves `task.threads` non-empty
+     * forever, so a lift that went idle before `task.return` and was later
+     * woken by another driver had its results captured and its host Promise
+     * left hanging. Nor does the answer need deferring to "no wasm call in
+     * flight": `driveDone`'s `midWasmCall`/`hopParked` clauses keep a DRIVER
+     * driving under a suspended activation, and on this path there is no lift
+     * driver left to stop — whichever driver is running when `task.return`
+     * happens keeps driving the store.
      *
      * Rejection has two sources: `finishHostEntry` itself throwing, and the
      * instance being poisoned by a later driver that ran this task into a
@@ -1904,10 +1925,11 @@
      */
     const backgroundCompletion = (): Promise<unknown> =>
       new Promise((resolve, reject) => {
-        // Degenerate case: the task is already over (its threads unregistered
-        // during the drive) and only a foreign hop kept `done` false. Nothing
-        // will fire `onFinished`, so settle now.
-        if (task.threads.length === 0) {
+        // Already resolved, and only a driver-liveness clause of `driveDone`
+        // (a foreign hop, this task's own suspended thread, `driveAsync`'s
+        // idle probe) kept the verdict off `"done"`. The answer is in hand:
+        // settle now, nothing further will fire.
+        if (resolvedSeen) {
           try {
             resolve(finishHostEntry());
           } catch (e) {
@@ -1916,13 +1938,17 @@
           return;
         }
         const onPoison = (cause: unknown): void => {
-          task.onFinished = null;
+          onResolvedHook = null;
           reject(cause);
         };
         registerPendingLift(inst, onPoison);
-        task.onFinished = () => {
+        onResolvedHook = () => {
           unregisterPendingLift(inst, onPoison);
           try {
+            // Safe synchronously inside the resolve callback: pure over
+            // already-lifted `ComponentValue`s (`canon_task_return` lifted
+            // them into `resolved`; `resultsToHost` reshapes). Host `.then`
+            // handlers run on a microtask regardless.
             resolve(finishHostEntry());
           } catch (e) {
             reject(e);
@@ -2028,10 +2054,12 @@
      * then the settlement pump (servicing a settled host call belonging to
      * another task) resumed a background activation that transiently
      * hop-parked, so `hopParked()` read true in the continuation and the lift
-     * took `backgroundCompletion()` — which waits for the task's LAST thread
-     * to unregister, i.e. never, for a task holding long-lived spawned
-     * futures. The verdict cannot rot that way: `"done"` means the driver
-     * saw the task finished, which stays true.
+     * took `backgroundCompletion()` — a wait on an event that had already
+     * happened for this task (today a harmless detour; when the misroute was
+     * traced, `backgroundCompletion` waited for the task's LAST thread to
+     * unregister, i.e. never, for a task holding long-lived spawned futures).
+     * The verdict cannot rot that way: `"done"` means the driver saw the task
+     * finished, which stays true.
      */
     const finish = (verdict: DriveExit): unknown =>
       verdict === "idle" ? backgroundCompletion() : finishHostEntry();
@@ -2404,7 +2432,7 @@
     task.return_(results);
     // Post-return runs after the results were read out of guest memory,
     // with may_leave cleared (reference canon_lift).
     const postReturn = require(opts.postReturn, `${name} post-return`);
     if (postReturn !== null) {
       assert_(inst.mayLeave, "post-return with may_leave already false");
       inst.mayLeave = false;
@@ -2450,7 +2478,7 @@
   // *mixed* activation, which pin (c) punishes: the first Suspending import
   // it reached would trap.
   const callback = enterWasm(
     require(opts.callback, `${name} callback`)!,
     input.mode,
   );
   const [packed] = normalizeCoreValues(

runtime/src/task/mod.ts

@@ -147,26 +147,6 @@ export class Task {
   /** TaskBorrowScope (cabi/context.ts): live borrows lowered into this task. */
   numBorrows = 0;
 
-  /**
-   * Fired once, when this task's LAST thread unregisters (#292).
-   *
-   * The reference has no such hook: `canon_lift` for an async-typed export
-   * returns to the embedder right after the first `thread.resume()`
-   * (definitions.py line 2189 — the driving loop with the empty-candidate-set
-   * `trap_if` is guarded by `if not ft.async_`), and the embedder's own
-   * `Store.tick` never traps on idle. Polyengine surfaces such an export as a
-   * Promise, so it needs the one thing a Python generator embedding gets for
-   * free: a signal for "the task is over" once its driver has gone home. Set
-   * only by exec/boundary.ts, only on the background path — a lifted async
-   * export whose driver exited idle with the task unresolved — and only ever
-   * once per task.
-   *
-   * Fires AFTER `unregisterThread`'s `trapIf(state !== "resolved")`, so a
-   * task that finished without resolving traps there as before and this never
-   * runs. Must not throw: it is called from inside guest-driven thread exit,
-   * under some OTHER call's driver.
-   */
-  onFinished: (() => void) | null = null;
   implicitThread: Thread | null = null;
   readonly threads: Thread[] = [];
   /**
@@ -329,14 +309,6 @@ export class Task {
     assert_(thread.index !== null, "unregister of an unindexed thread");
     this.inst.threads.remove(thread.index);
     thread.index = null;
-    if (this.threads.length === 0 && this.onFinished !== null) {
-      // Last, and after the table removal: the callback resolves a host-facing
-      // Promise off this task's results, so everything this task owns must
-      // already be released when it runs. One-shot (see `onFinished`).
-      const f = this.onFinished;
-      this.onFinished = null;
-      f();
-    }
   }
 
   /**

runtime/tests/lift_background_return_test.ts

@@ -0,0 +1,156 @@
+// An async-typed lift whose driver exited `idle` BEFORE `task.return` settles
+// when the task returns — polyengine#313.
+//
+// THE SHAPE. A callback-ABI export spawns futures that outlive the call (the
+// wit-bindgen `spawn_local` pattern: an event loop, a driver, an accept loop)
+// and goes idle without an answer, so its own driver exits `"idle"` and the
+// host gets a Promise from `backgroundCompletion()`. A later export call on
+// the same store drives the store, resumes one of those futures, and it calls
+// `task.return`. The task's result exists from that moment on; the task's
+// LAST thread, meanwhile, never unregisters — the immortal future keeps
+// `task.threads` non-empty for the instance's life. Settling on thread drain
+// therefore hangs the host forever; settling on the task's resolve callback
+// (definitions.py `Task.return_` -> `on_resolve`) is both correct and enough.
+//
+// Store-level, in the style of `lift_done_verdict_test.ts`: no checked-in
+// example guest has the participants.
+
+import { assertEq } from "./support/asserts.ts";
+import {
+  createLiftedFunction,
+  newStats,
+  type ResolvedOptions,
+} from "../src/exec/mod.ts";
+import {
+  ComponentInstanceState,
+  currentThread,
+  Store,
+  Thread,
+} from "../src/task/mod.ts";
+import type { FuncType } from "../src/cabi/types.ts";
+
+function assert(cond: boolean, msg: string): asserts cond {
+  if (!cond) throw new Error(`assertion failed: ${msg}`);
+}
+
+/** `async func()` — no params, no results, so no memory is needed. */
+const FT: FuncType = { params: [], results: [], async: true };
+
+/** The task seam the cores below reach through `currentThread()`. */
+type TaskSeam = {
+  return_(r: never[]): void;
+  registerThread(t: Thread): void;
+};
+
+function taskOf(): TaskSeam {
+  const thread = currentThread() as unknown as { task: unknown };
+  return thread.task as TaskSeam;
+}
+
+Deno.test({
+  name:
+    "an async lift that went idle before task.return settles when the task returns (#313)",
+  fn: async () => {
+    const store = new Store();
+    const inst = new ComponentInstanceState(0, store);
+    const opts: ResolvedOptions = {
+      stringEncoding: "utf8",
+      memory: null,
+      realloc: null,
+      postReturn: null,
+      // Callback ABI: the packed code the core returns is CallbackCode.EXIT
+      // (0), so the loop exits at once and the implicit thread unregisters.
+      callback: () => (() => 0) as never,
+      async: true,
+      cancellable: false,
+      coreType: { params: [], results: ["i32"] },
+      instance: inst,
+    };
+
+    // The test-owned condition the returner future parks on: the "later
+    // event" that gives the export its answer.
+    let wake = false;
+    let immortal!: Thread;
+
+    const boot = createLiftedFunction({
+      name: "boot",
+      ft: FT,
+      opts,
+      core: () => {
+        const task = taskOf();
+        // A spawned future that never becomes ready: the model of the
+        // consumer's long-lived background work. It alone keeps this task's
+        // thread list non-empty forever.
+        immortal = new Thread(
+          task as never,
+          (function* () {
+            yield { readyFunc: () => false, cancellable: false };
+          })(),
+        );
+        task.registerThread(immortal);
+        immortal.resume();
+        // The future that will eventually produce the export's result.
+        const returner = new Thread(
+          task as never,
+          (function* () {
+            yield { readyFunc: () => wake, cancellable: false };
+            task.return_([]);
+          })(),
+        );
+        task.registerThread(returner);
+        returner.resume();
+        // EXIT synchronously: the implicit thread leaves at once (releasing
+        // the exclusive slot), the task is unresolved with two live threads,
+        // and the driver has no ready candidate and no host call outstanding
+        // — the `"idle"` verdict, taken synchronously.
+        return 0;
+      },
+      stats: newStats(),
+    });
+
+    const out = boot() as Promise<unknown>;
+    assert(
+      out instanceof Promise,
+      "an async-typed lift that exits idle must return a Promise",
+    );
+
+    // Pinning that the test really went through the idle path: nothing can
+    // give this task an answer until `wake` flips.
+    const early = await Promise.race([
+      out.then(() => "resolved"),
+      new Promise((r) => setTimeout(() => r("pending"), 0)),
+    ]);
+    assertEq(early, "pending");
+
+    // The later event, plus a driver to notice it: a second export call on
+    // the same instance. `boot`'s implicit thread released the exclusive slot
+    // when it exited, so this one enters; its driver's tick drain resumes the
+    // returner, which calls `task.return` on the FIRST task.
+    wake = true;
+    const ping = createLiftedFunction({
+      name: "ping",
+      ft: FT,
+      opts,
+      core: () => {
+        taskOf().return_([]);
+        return 0;
+      },
+      stats: newStats(),
+    });
+    await ping();
+
+    const settled = await Promise.race([
+      out.then((v) => ({ v })),
+      new Promise((r) => setTimeout(() => r("pending"), 0))
+        .then(() => new Promise((r) => setTimeout(() => r("pending"), 0))),
+    ]);
+    assertEq(settled, { v: undefined });
+    assert(
+      immortal.waiting(),
+      "the immortal thread must still be parked: otherwise the task drained " +
+        "and the test proves nothing about resolving on task.return",
+    );
+    assertEq(store.hostFailure, undefined);
+    assertEq(store.pendingHostCalls.size, 0);
+  },
+});