Add cooperative threads

Author: lukewagner

design/mvp/Async.md

@@ -240,6 +240,8 @@ as a parameter.
 
 ### Context-Local Storage
 
+TODO: update (also there are 2 now)
+
 Each task contains a distinct mutable **context-local storage** array. The
 current task's context-local storage can be read and written from core wasm
 code by calling the [`context.get`] and [`context.set`] built-ins.

design/mvp/canonical-abi/definitions.py

@@ -213,20 +213,31 @@ def tick(self):
 
 class Thread:
   task: Task
+  index: int
+  context: list[int]
   ready_func: Optional[Callable[[], bool]]
   run_lock: threading.Lock
   resume_lock: Optional[threading.Lock]
   stack: threading.Thread
+  cancellable: bool
+  waiting_for_callback: bool
+
+  CONTEXT_LENGTH = 2
 
   def __init__(self, task, thread_func):
     self.task = task
+    self.index = task.inst.table.add(self)
+    self.context = [0] * Thread.CONTEXT_LENGTH
     self.ready_func = None
     self.run_lock = threading.Lock()
     self.resume_lock = None
+    self.cancellable = False
+    self.waiting_for_callback = False
     def thread_stack_base():
       self.run_lock.acquire()
       thread_func(self)
       self.task.thread_stop(self)
+      self.task.inst.table.remove(self.index)
       self.resume_lock.release()
     self.stack = threading.Thread(target = thread_stack_base)
     self.run_lock.acquire()
@@ -247,14 +258,45 @@ def resume(self):
     self.resume_lock.acquire()
     self.resume_lock = None
 
-  def suspend_until(self, ready_func):
+  def block(self, cancellable):
+    assert(not self.cancellable)
+    self.cancellable = cancellable
+    self.resume_lock.release()
+    self.run_lock.acquire()
+    self.cancellable = False
+
+  def suspend_until(self, ready_func, cancellable = False):
     assert(not self.ready_func)
     if not DETERMINISTIC_PROFILE and ready_func():
       return
     self.ready_func = ready_func
     self.task.inst.store.waiting.append(self)
-    self.resume_lock.release()
+    self.block(cancellable)
+
+  async def switch_to(self, cancellable, other: Thread):
+    assert(self.task.inst is other.task.inst)
+    if other.ready_func:
+      other.ready_func = None
+      other.task.inst.store.waiting.remove(other)
+    assert(not self.cancellable)
+    self.cancellable = cancellable
+    assert(self.resume_lock and not other.resume_lock)
+    other.resume_lock = self.resume_lock
+    self.resume_lock = None
+    other.run_lock.release()
     self.run_lock.acquire()
+    self.cancellable = False
+
+  def yield_to(self, cancellable, other: Thread):
+    assert(not self.ready_func)
+    self.ready_func = lambda: True
+    self.task.inst.store.waiting.append(self)
+    self.switch_to(cancellable, other)
+
+  def unblock(self, other: Thread):
+    if not other.ready_func:
+      other.task.inst.store.waiting.append(other)
+    other.ready_func = lambda: True
 
 
 ### Lifting and Lowering Context
@@ -432,22 +474,6 @@ def write(self, vs):
       assert(all(v == () for v in vs))
     self.progress += len(vs)
 
-#### Context-Local Storage
-
-class ContextLocalStorage:
-  LENGTH = 1
-  array: list[int]
-
-  def __init__(self):
-    self.array = [0] * ContextLocalStorage.LENGTH
-
-  def set(self, i, v):
-    assert(types_match_values(['i32'], [v]))
-    self.array[i] = v
-
-  def get(self, i):
-    return self.array[i]
-
 #### Waitable State
 
 class EventCode(IntEnum):
@@ -458,6 +484,7 @@ class EventCode(IntEnum):
   FUTURE_READ = 4
   FUTURE_WRITE = 5
   TASK_CANCELLED = 6
+  THREAD_RESUMED = 7
 
 EventTuple = tuple[EventCode, int, int]
 
@@ -530,11 +557,8 @@ class State(Enum):
   ft: FuncType
   supertask: Optional[Task]
   on_resolve: OnResolve
-  thread: Optional[Thread]
-  cancellable: bool
-  waiting_for_callback: bool
+  threads: list[Thread]
   num_borrows: int
-  context: ContextLocalStorage
 
   def __init__(self, opts, inst, ft, supertask, on_resolve):
     self.state = Task.State.INITIAL
@@ -543,11 +567,8 @@ def __init__(self, opts, inst, ft, supertask, on_resolve):
     self.ft = ft
     self.supertask = supertask
     self.on_resolve = on_resolve
-    self.thread = None
-    self.cancellable = False
-    self.waiting_for_callback = False
+    self.threads = []
     self.num_borrows = 0
-    self.context = ContextLocalStorage()
 
   def trap_if_on_the_stack(self, inst):
     c = self.supertask
@@ -559,15 +580,13 @@ def needs_exclusive(self):
     return self.opts.sync or self.opts.callback
 
   def enter(self, thread):
-    assert(thread is self.thread and thread.task is self)
+    assert(thread in self.threads and thread.task is self)
     def has_backpressure():
       return self.inst.backpressure or (self.needs_exclusive() and self.inst.exclusive)
     if has_backpressure() or self.inst.pending_tasks > 0:
       self.inst.pending_tasks += 1
-      self.cancellable = True
-      thread.suspend_until(lambda: not has_backpressure())
+      thread.suspend_until(lambda: not has_backpressure(), cancellable = True)
       self.inst.pending_tasks -= 1
-      self.cancellable = False
       if self.deliver_cancel():
         self.cancel()
         return False
@@ -586,27 +605,28 @@ def deliver_cancel(self) -> bool:
   def request_cancellation(self):
     assert(self.state == Task.State.INITIAL)
     self.state = Task.State.PENDING_CANCEL
-    if self.cancellable and not (self.waiting_for_callback and self.inst.exclusive):
-      self.thread.resume()
+    if not DETERMINISTIC_PROFILE:
+      random.shuffle(self.threads)
+    for thread in self.threads:
+      if thread.cancellable and not (thread.waiting_for_callback and self.inst.exclusive):
+        thread.resume()
+        break
 
   def wait_until(self, ready_func, thread, cancellable, for_callback):
-    assert(thread is self.thread and thread.task is self)
+    assert(thread in self.threads and thread.task is self)
     if cancellable and self.deliver_cancel():
       return True
-    assert(not self.cancellable)
-    self.cancellable = cancellable
     if for_callback:
       assert(self.inst.exclusive)
       self.inst.exclusive = False
-      self.waiting_for_callback = True
+      thread.waiting_for_callback = True
     def ready_and_allowed():
       return ready_func() and not (for_callback and self.inst.exclusive)
-    thread.suspend_until(ready_and_allowed)
+    thread.suspend_until(ready_and_allowed, cancellable)
     if for_callback:
       assert(not self.inst.exclusive)
       self.inst.exclusive = True
-      self.waiting_for_callback = False
-    self.cancellable = False
+      thread.waiting_for_callback = False
     if cancellable and self.deliver_cancel():
       return True
     return False
@@ -615,13 +635,15 @@ def wait_for_event(self, thread, wset, cancellable, for_callback) -> EventTuple:
     wset.num_waiting += 1
     cancelled = self.wait_until(wset.has_pending_event, thread, cancellable, for_callback)
     wset.num_waiting -= 1
+    # TODO: somehow get a THREAD_RESUME event...
     if cancelled:
       return (EventCode.TASK_CANCELLED, 0, 0)
     else:
       return wset.get_pending_event()
 
   def yield_(self, thread, cancellable, for_callback) -> EventTuple:
     cancelled = self.wait_until(lambda: True, thread, cancellable, for_callback)
+    # TODO: somehow get a THREAD_RESUME event...
     if cancelled:
       return (EventCode.TASK_CANCELLED, 0, 0)
     else:
@@ -631,6 +653,7 @@ def poll_for_event(self, thread, wset, cancellable, for_callback) -> Optional[Ev
     wset.num_waiting += 1
     cancelled = self.wait_until(lambda: True, thread, cancellable, for_callback)
     wset.num_waiting -= 1
+    # TODO: somehow get a THREAD_RESUME event...
     if cancelled:
       return (EventCode.TASK_CANCELLED, 0, 0)
     elif wset.has_pending_event():
@@ -652,20 +675,21 @@ def cancel(self):
     self.state = Task.State.RESOLVED
 
   def exit(self):
-    assert(self.thread is not None)
+    assert(len(self.threads) > 0)
     if self.needs_exclusive():
       assert(self.inst.exclusive)
       self.inst.exclusive = False
 
   def thread_start(self, thread):
-    assert(self.thread is None and thread.task is self)
-    self.thread = thread
+    assert(thread not in self.threads and thread.task is self)
+    self.threads.append(thread)
 
   def thread_stop(self, thread):
-    assert(thread is self.thread and thread.task is self)
-    self.thread = None
-    trap_if(self.state != Task.State.RESOLVED)
-    assert(self.num_borrows == 0)
+    assert(thread in self.threads and thread.task is self)
+    self.threads.remove(thread)
+    if len(self.threads) == 0:
+      trap_if(self.state != Task.State.RESOLVED)
+      assert(self.num_borrows == 0)
 
 #### Subtask State
 
@@ -2081,25 +2105,80 @@ def canon_resource_rep(rt, thread, i):
   trap_if(h.rt is not rt)
   return [h.rep]
 
+### 🧵 `canon thread.index`
+
+def canon_thread_index(shared, thread):
+  assert(not shared)
+  return [thread.index]
+
+### 🧵 `canon thread.new`
+
+def canon_thread_new(ft, ftbl, thread, i, c):
+  task = thread.task
+  trap_if(not task.inst.may_leave)
+  f = task.inst.ftbl.get(i)
+  trap_if(f.type != ft)
+  new_thread = Thread(task, f(c))
+  return [new_thread.index]
+
+### 🧵 `canon thread.switch-to`
+
+def canon_thread_switch_to(thread, cancellable, i):
+  trap_if(not thread.task.inst.may_leave)
+  other = thread.task.inst.table.get(i)
+  trap_if(not isinstance(other, Thread))
+  trap_if(not other.cancellable) # TODO: what about waiting_for_callback and exclusive
+  cancelled = thread.switch_to(cancellable, other)
+  return [ 1 if cancelled else 0 ]
+
+### 🧵 `canon thread.yield-to`
+
+def canon_thread_yield_to(thread, cancellable, i):
+  trap_if(not thread.task.inst.may_leave)
+  other = thread.task.inst.table.get(i)
+  trap_if(not isinstance(other, Thread))
+  trap_if(not other.cancellable) # TODO: what about waiting_for_callback and exclusive
+  other.yield_to(cancellable, other)
+  return []
+
+### 🧵 `canon thread.unblock`
+
+def canon_thread_unblock(thread, i):
+  trap_if(not thread.task.inst.may_leave)
+  other = thread.task.inst.table.get(i)
+  trap_if(not isinstance(other, Thread))
+  trap_if(not other.cancellable) # TODO: what about waiting_for_callback and exclusive
+  thread.unblock()
+  return []
+
+### 🧵 `canon thread.block`
+
+def canon_thread_block(thread, cancellable, i):
+  trap_if(not thread.task.inst.may_leave)
+  other = thread.task.inst.table.get(i)
+  trap_if(not isinstance(other, Thread))
+  cancelled = thread.block(cancellable)
+  return [ 1 if cancelled else 0 ]
+
 ### 🔀 `canon context.get`
 
 def canon_context_get(t, i, thread):
   assert(t == 'i32')
-  assert(i < ContextLocalStorage.LENGTH)
-  return [thread.task.context.get(i)]
+  assert(i < Thread.CONTEXT_LENGTH)
+  return [thread.context[i]]
 
 ### 🔀 `canon context.set`
 
 def canon_context_set(t, i, thread, v):
   assert(t == 'i32')
-  assert(i < ContextLocalStorage.LENGTH)
-  thread.task.context.set(i, v)
+  assert(i < Thread.CONTEXT_LENGTH)
+  thread.context[i] = v
   return []
 
 ### 🔀 `canon backpressure.set`
 
 def canon_backpressure_set(thread, flat_args):
-  trap_if(thread.task.opts.sync)
+  # TODO: remove trap_if(thread.task.opts.sync)
   assert(len(flat_args) == 1)
   thread.task.inst.backpressure = bool(flat_args[0])
   return []