Add thread.{get,set}-task and task.drop built-ins

Author: lukewagner

design/mvp/Binary.md

@@ -340,6 +340,9 @@ canon    ::= 0x00 0x00 f:<core:funcidx> opts:<opts> ft:<typeidx> => (canon lift
            | 0x2b 0x00                                           => (canon thread.yield-then-resume (core func)) 🧵
            | 0x2c 0x00                                           => (canon thread.suspend-then-promote (core func)) 🧵
            | 0x2d 0x00                                           => (canon thread.yield-then-promote (core func)) 🧵
+           | 0x30                                                => (canon thread.get-task (core func)) 🧵
+           | 0x31                                                => (canon thread.set-task (core func)) 🧵
+           | 0x32                                                => (canon task.drop (core func)) 🧵
            | 0x40 shared?:<sh?> ft:<core:typeidx>                => (canon thread.spawn-ref shared? ft (core func)) 🧵②
            | 0x41 shared?:<sh?> ft:<core:typeidx> tbl:<core:tableidx> => (canon thread.spawn-indirect shared? ft tbl (core func)) 🧵②
            | 0x42 shared?:<sh?>                                  => (canon thread.available-parallelism shared? (core func)) 🧵②

design/mvp/CanonicalABI.md

@@ -67,6 +67,9 @@ specified here.
   * [`canon thread.yield-then-resume`](#-canon-threadyield-then-resume) 🧵
   * [`canon thread.suspend-then-promote`](#-canon-threadsuspend-then-promote) 🧵
   * [`canon thread.yield-then-promote`](#-canon-threadyield-then-promote) 🧵
+  * [`canon thread.get-task`](#-canon-threadget-task) 🧵
+  * [`canon thread.set-task`](#-canon-threadset-task) 🧵
+  * [`canon task.drop`](#-canon-taskdrop) 🧵
   * [`canon error-context.new`](#-canon-error-contextnew) 📝
   * [`canon error-context.debug-message`](#-canon-error-contextdebug-message) 📝
   * [`canon error-context.drop`](#-canon-error-contextdrop) 📝
@@ -121,7 +124,7 @@ into the `handles` or `threads` fields of `ComponentInstance`.
 ```python
 class ComponentInstance:
   store: Store
-  handles: Table[ResourceHandle | Waitable | WaitableSet | ErrorContext]
+  handles: Table[ResourceHandle | Waitable | WaitableSet | ErrorContext | Task]
   threads: Table[Thread]
   may_leave: bool
   backpressure: int
@@ -3425,10 +3428,12 @@ are lifted according to `lift_flat_values` above, the optional `post-return`
 function (specified as a `canonopt` immediate of `canon lift`) is called,
 passing the same core wasm results as parameters so that the `post-return`
 function can free any associated allocations.
+TODO: copy from set-task-old (and also alll the other stuff in Concurrency.md)
 ```python
     if not opts.async_:
       flat_results = call_and_trap_on_throw(callee, flat_args)
       assert(types_match_values(flat_ft.results, flat_results))
+      trap_if(thread.task is not task)
       result = lift_flat_values(cx, MAX_FLAT_RESULTS, CoreValueIter(flat_results), ft.result_type())
       task.return_(result)
       if opts.post_return is not None:
@@ -4673,17 +4678,17 @@ class CoreFuncRef:
   callee: Callable[[list[CoreValType]], list[CoreValType]]
 
 def canon_thread_new_indirect(ft, ftbl: Table[CoreFuncRef], fi, c):
-  task = current_task()
-  trap_if(not task.inst.may_leave)
+  inst = current_instance()
+  trap_if(not inst.may_leave)
   f = ftbl.get(fi)
   assert(ft == CoreFuncType(['i32'], []) or ft == CoreFuncType(['i64'], []))
   trap_if(f.t != ft)
   def thread_func():
     [] = call_and_trap_on_throw(f.callee, [c])
-    task.inst.threads.remove(new_thread.index)
-  new_thread = Thread(task, thread_func)
+    inst.threads.remove(new_thread.index)
+  new_thread = Thread(current_task(), thread_func)
   assert(new_thread.suspended())
-  new_thread.index = task.inst.threads.add(new_thread)
+  new_thread.index = inst.threads.add(new_thread)
   return [new_thread.index]
 ```
 The newly-created thread starts out in a "suspended" state and so, to
@@ -4862,6 +4867,66 @@ def canon_thread_yield_then_promote(i):
 ```
 
 
+###### 🧵 `canon thread.get-task`
+
+For a canonical definition:
+```wat
+(canon thread.get-task (core func $thread.get-task))
+```
+validation specifies:
+* `$thread.get-task` is given type `(func (result i32))`
+
+Calling `$thread.get-task` invokes the following function which TODO
+```python
+def canon_thread_get_task():
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  taski = thread.task.inst.handles.add(thread.task)
+  return [taski]
+```
+
+
+###### 🧵 `canon thread.set-task`
+
+For a canonical definition:
+```wat
+(canon thread.set-task (core func $thread.set-task))
+```
+validation specifies:
+* `$thread.set-task` is given type `(func (param i32))`
+
+Calling `$thread.set-task` invokes the following function which TODO
+```python
+def canon_thread_set_task(taski):
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  new_task = thread.task.inst.handles.get(taski)
+  trap_if(not isinstance(new_task, Task))
+  thread.task = new_task
+  return []
+```
+
+
+###### 🧵 `canon task.drop`
+
+For a canonical definition:
+```wat
+(canon thread.drop (core func $task.drop))
+```
+validation specifies:
+* `$thread.drop` is given type `(func (param i32))`
+
+Calling `$task.drop` invokes the following function which TODO
+```python
+def canon_task_drop(taski):
+  inst = current_instance()
+  trap_if(not inst.may_leave)
+  task = inst.handles.remove(taski)
+  trap_if(not isinstance(task, Task))
+  return []
+```
+
+
 ### 📝 `canon error-context.new`
 
 For a canonical definition:

design/mvp/Concurrency.md

@@ -268,11 +268,14 @@ Thread
 where a **component store** is the top-level "thing" and analogous to a Core
 WebAssembly [store].
 
-The reason for the thread/task split is that, when one thread creates a new
-thread, the new thread is contained by the task of the original thread which
-creates an N:1 relationship between threads and tasks that ties N threads to
-the original export call (= "task") that transitively spawned those N threads.
-This relationship serves several purposes described in the following sections.
+The reason for the thread/task split is that a single component export call,
+which spawns a task, can spawn N more threads as part of the execution of that
+task, thereby creating an N:1 relationship between threads and tasks.
+
+While the store:instance and instance:task relationships are immutably set on
+creation, the task:thread relationship is *mutable*: guest code running inside a
+component instance can change which task a thread is currently executing on
+behalf of by calling the [`thread.set-task`] built-in, as described below.
 
 In the Canonical ABI explainer, threads, tasks, component instances and
 component stores are represented by the [`Thread`], [`Task`],
@@ -311,9 +314,15 @@ supertask, they can be thought of as a single node in the async call stack.
 
 A subtask/supertask relationship is immutably established when an import is
 called, setting the [current task](#current-thread-and-task) as the supertask
-of the new subtask created for the import call. Thus, one reason for
-associating every thread with a "containing task" is to ensure that there is
-always a well-defined async call stack.
+of the new subtask created for the import call. Thus, one reason for associating
+every thread with a "containing task" is to ensure that there is always a
+well-defined async call stack. Note that guest code can call [`thread.set-task`]
+to change the containing task of a thread and thus the async call stack can
+change completely between two program points while executing on a single thread.
+For example, when a JS runtime flushes its microtask queue and encounters a
+JS callback associated with a task other than the current thread's task, the JS
+runtime would call `thread.set-task` so that the async call stack matches the
+JS developer's expectation.
 
 The async call stack is not currently observable to running components, except
 that it may nondeterministically appear as part of the callstack stored in
@@ -336,7 +345,7 @@ not enforcing a stricter form of Structured Concurrency at the Component Model
 level is that there are important use cases where forcing a supertask's thread
 to stay resident just to wait for subtasks to finish would waste resources
 without tangible benefit. Instead, we can say that once a supertask's last
-thread finishes execution, the supertask semantically "tail calls" any still-
+thread exits or leaves, the supertask semantically "tail calls" any still-
 executing subtasks, staying technically-alive and on the async call stack until
 they complete, but not consuming real resources.
 
@@ -374,13 +383,23 @@ New threads are created with the [`thread.new-indirect`] built-in. As mentioned
 [above](#threads-and-tasks), a spawned thread inherits the task of the spawning
 thread which is why threads and tasks are N:1. `thread.new-indirect` adds a new
 thread to the component instance's threads table and returns the `i32` index of
-this table entry to the Core WebAssembly caller. Like [`pthread_create`],
-`thread.new-indirect` takes a Core WebAssembly function (via index into a
-`funcref` table) and a "closure" parameter to pass to the function when called
-on the new thread. However, unlike `pthread_create`, the new thread is
-initially in a "suspended" state and must be explicitly "resumed" using one of
-the following 3 thread built-ins. Once the thread is resumed, the thread can
-learn its own index by calling the [`thread.index`] built-in.
+this table entry to the Core WebAssembly caller.
+
+After creation, the implicitly-set containing task of a thread can be explicitly
+overridden using the [`thread.set-task`] built-in. `thread.set-task` sets the
+containing task of the current thread to the task of some other thread that was
+retrieved via [`thread.get-task`]. `thread.get-task` always allocates a fresh
+handle storing a reference to the current thread's containing task, returning
+the `i32` of this new handle which must be explicitly dropped via [`task.drop`]
+to avoid leaking the task. Tasks are thus (acyclicly) kept alive by any or all
+of: contained threads, subtask handles and task handles.
+
+Like [`pthread_create`], `thread.new-indirect` takes a Core WebAssembly function
+(via index into a `funcref` table) and a "closure" parameter to pass to the
+function when called on the new thread. However, unlike `pthread_create`, the
+new thread is initially in a "suspended" state and must be explicitly "resumed"
+using one of the following 3 thread built-ins. Once the thread is resumed, the
+thread can learn its own index by calling the [`thread.index`] built-in.
 
 A suspended thread (identified by thread-table index) can be resumed at some
 nondeterministic point in future via the [`thread.resume-later`] built-in. In
@@ -725,10 +744,7 @@ the "started" state.
 The way an `async` export returns its value using the async ABI is by calling
 [`task.return`], passing the core values that are to be lifted as *parameters*.
 When using the async ABI, *any* of the threads contained by a task can call
-`task.return`; there is no "main thread" of a task. When the last thread of a
-task returns, there is a trap if `task.return` has not been called. Thus, *some*
-thread (either the thread created implicitly for the initial export call or some
-thread transitively created by that thread) must call `task.return`.
+`task.return`; there is no "main thread" of a task.
 
 Returning values by calling `task.return` allows a task to continue executing
 even after it has passed its initial results to the caller. This is also
@@ -1560,6 +1576,9 @@ the concurrency story:
 [`thread.yield-then-resume`]: Explainer.md#-threadyield-then-resume
 [`thread.suspend-then-promote`]: Explainer.md#-threadsuspend-then-promote
 [`thread.yield-then-promote`]: Explainer.md#-threadyield-then-promote
+[`thread.set-task`]: Explainer.md#-threadset-task
+[`thread.get-task`]: Explainer.md#-threadget-task
+[`task.drop`]: Explainer.md#-taskdrop
 [`{stream,future}.new`]: Explainer.md#-streamnew-and-futurenew
 [`{stream,future}.{read,write}`]: Explainer.md#-streamread-and-streamwrite
 [`stream.cancel-write`]: Explainer.md#-streamcancel-read-streamcancel-write-futurecancel-read-and-futurecancel-write

design/mvp/Explainer.md

@@ -1590,6 +1590,9 @@ canon ::= ...
         | (canon thread.yield-then-resume (core func <id>?)) 🧵
         | (canon thread.suspend-then-promote (core func <id>?)) 🧵
         | (canon thread.yield-then-promote (core func <id>?)) 🧵
+        | (canon thread.get-task (core func <id>?)) 🧵
+        | (canon thread.set-task (core func <id>?)) 🧵
+        | (canon task.drop (core func <id>?)) 🧵
         | (canon error-context.new <canonopt>* (core func <id>?)) 📝
         | (canon error-context.debug-message <canonopt>* (core func <id>?)) 📝
         | (canon error-context.drop (core func <id>?)) 📝
@@ -2298,6 +2301,43 @@ returned `i32` is always `0` and may be removed in a future ABI revision.
 For details, see [Thread Built-ins] in the concurrency explainer and
 [`canon_thread_yield_then_promote`] in the Canonical ABI explainer.
 
+###### 🧵 `thread.get-task`
+
+| Synopsis                   |                  |
+| -------------------------- | ---------------- |
+| Approximate WIT signature  | `func() -> task` |
+| Canonical ABI signature    | `[] -> [i32]`    |
+
+TODO
+
+For details, see [Thread Built-ins] in the concurrency explainer and
+[`canon_thread_get_task`] in the Canonical ABI explainer.
+
+###### 🧵 `thread.set-task`
+
+| Synopsis                   |                 |
+| -------------------------- | --------------- |
+| Approximate WIT signature  | `func(t: task)` |
+| Canonical ABI signature    | `[t:i32] -> []` |
+
+TODO
+
+For details, see [Thread Built-ins] in the concurrency explainer and
+[`canon_thread_set_task`] in the Canonical ABI explainer.
+
+###### 🧵 `task.drop`
+
+| Synopsis                   |                 |
+| -------------------------- | --------------- |
+| Approximate WIT signature  | `func(t: task)` |
+| Canonical ABI signature    | `[t:i32] -> []` |
+
+TODO
+
+For details, see [Thread Built-ins] in the concurrency explainer and
+[`canon_task_drop`] in the Canonical ABI explainer.
+
+
 ###### 🧵② `thread.spawn-ref`
 
 | Synopsis                   |                                                                                            |
@@ -3440,6 +3480,9 @@ For some use-case-focused, worked examples, see:
 [`canon_thread_yield_then_resume`]: CanonicalABI.md#-canon-threadyield-then-resume
 [`canon_thread_suspend_then_promote`]: CanonicalABI.md#-canon-threadsuspend-then-promote
 [`canon_thread_yield_then_promote`]: CanonicalABI.md#-canon-threadyield-then-promote
+[`canon_thread_get_task`]: CanonicalABI.md#-canon-threadget-task
+[`canon_thread_set_task`]: CanonicalABI.md#-canon-threadset-task
+[`canon_task_drop`]: CanonicalABI.md#-canon-taskdrop
 [`canon_thread_spawn_ref`]: CanonicalABI.md#-canon-threadspawn-ref
 [`canon_thread_spawn_indirect`]: CanonicalABI.md#-canon-threadspawn-indirect
 [`canon_thread_available_parallelism`]: CanonicalABI.md#-canon-threadavailable_parallelism

design/mvp/canonical-abi/definitions.py

@@ -190,7 +190,7 @@ class FutureType(ValType):
 
 class ComponentInstance:
   store: Store
-  handles: Table[ResourceHandle | Waitable | WaitableSet | ErrorContext]
+  handles: Table[ResourceHandle | Waitable | WaitableSet | ErrorContext | Task]
   threads: Table[Thread]
   may_leave: bool
   backpressure: int
@@ -2090,6 +2090,7 @@ def thread_func():
     if not opts.async_:
       flat_results = call_and_trap_on_throw(callee, flat_args)
       assert(types_match_values(flat_ft.results, flat_results))
+      trap_if(thread.task is not task)
       result = lift_flat_values(cx, MAX_FLAT_RESULTS, CoreValueIter(flat_results), ft.result_type())
       task.return_(result)
       if opts.post_return is not None:
@@ -2622,17 +2623,17 @@ class CoreFuncRef:
   callee: Callable[[list[CoreValType]], list[CoreValType]]
 
 def canon_thread_new_indirect(ft, ftbl: Table[CoreFuncRef], fi, c):
-  task = current_task()
-  trap_if(not task.inst.may_leave)
+  inst = current_instance()
+  trap_if(not inst.may_leave)
   f = ftbl.get(fi)
   assert(ft == CoreFuncType(['i32'], []) or ft == CoreFuncType(['i64'], []))
   trap_if(f.t != ft)
   def thread_func():
     [] = call_and_trap_on_throw(f.callee, [c])
-    task.inst.threads.remove(new_thread.index)
-  new_thread = Thread(task, thread_func)
+    inst.threads.remove(new_thread.index)
+  new_thread = Thread(current_task(), thread_func)
   assert(new_thread.suspended())
-  new_thread.index = task.inst.threads.add(new_thread)
+  new_thread.index = inst.threads.add(new_thread)
   return [new_thread.index]
 
 ### 🧵 `canon thread.resume-later`
@@ -2704,6 +2705,33 @@ def canon_thread_yield_then_promote(i):
   thread.yield_then_promote(other_thread)
   return [0]
 
+### 🧵 `canon thread.get-task`
+
+def canon_thread_get_task():
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  taski = thread.task.inst.handles.add(thread.task)
+  return [taski]
+
+### 🧵 `canon thread.set-task`
+
+def canon_thread_set_task(taski):
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  new_task = thread.task.inst.handles.get(taski)
+  trap_if(not isinstance(new_task, Task))
+  thread.task = new_task
+  return []
+
+### 🧵 `canon task.drop`
+
+def canon_task_drop(taski):
+  inst = current_instance()
+  trap_if(not inst.may_leave)
+  task = inst.handles.remove(taski)
+  trap_if(not isinstance(task, Task))
+  return []
+
 ### 📝 `canon error-context.new`
 
 @dataclass