CABI: improve and add cooperative thread built-ins

Author: lukewagner

design/mvp/canonical-abi/definitions.py

@@ -16,6 +16,7 @@
 
 class Trap(BaseException): pass
 class CoreWebAssemblyException(BaseException): pass
+class ThreadExit(BaseException): pass
 
 def trap():
   raise Trap()
@@ -321,13 +322,15 @@ def waiting(self):
     return not self.running() and self.ready_func is not None
 
   def ready(self):
-    assert(self.waiting())
-    return self.ready_func()
+    return self.waiting() and self.ready_func()
 
   def __init__(self, task, thread_func):
     def cont_func(cancelled):
       assert(self.running() and not cancelled)
-      thread_func()
+      try:
+        thread_func()
+      except ThreadExit:
+        pass
       return None
     self.cont = cont_new(cont_func)
     self.ready_func = None
@@ -337,6 +340,10 @@ def cont_func(cancelled):
     self.storage = [0,0]
     assert(self.suspended())
 
+  def exit(self):
+    assert(self.running())
+    raise ThreadExit()
+
   def resume_later(self):
     assert(self.suspended())
     self.ready_func = lambda: True
@@ -345,7 +352,8 @@ def resume_later(self):
 
   def resume(self, cancelled = Cancelled.FALSE):
     assert(not self.running() and (self.cancellable or not cancelled))
-    self.stop_waiting()
+    if self.waiting():
+      self.stop_waiting(cancelled)
     thread = self
     while thread is not None:
       cont = thread.cont
@@ -354,23 +362,23 @@ def resume(self, cancelled = Cancelled.FALSE):
       thread = switch_to
       cancelled = Cancelled.FALSE
 
-  def stop_waiting(self):
-    if self.waiting():
-      self.ready_func = None
-      self.task.inst.store.waiting.remove(self)
+  def stop_waiting(self, cancelled = Cancelled.FALSE):
+    assert(self.waiting())
+    assert(cancelled or self.ready_func())
+    self.ready_func = None
+    self.task.inst.store.waiting.remove(self)
 
   def suspend(self, cancellable) -> Cancelled:
-    assert(self.running() and self.task.may_block())
+    assert(self.running())
     if self.task.deliver_pending_cancel(cancellable):
-      self.stop_waiting()
       return Cancelled.TRUE
     self.cancellable = cancellable
     cancelled = block(switch_to = None)
     assert(self.running() and (cancellable or not cancelled))
     return cancelled
 
   def wait_until(self, ready_func, cancellable = False) -> Cancelled:
-    assert(self.running() and self.task.may_block())
+    assert(self.running())
     if self.task.deliver_pending_cancel(cancellable):
       return Cancelled.TRUE
     if ready_func() and not DETERMINISTIC_PROFILE and random.randint(0,1):
@@ -379,32 +387,41 @@ def wait_until(self, ready_func, cancellable = False) -> Cancelled:
     self.task.inst.store.waiting.append(self)
     return self.suspend(cancellable)
 
-  def yield_until(self, ready_func, cancellable) -> Cancelled:
-    assert(self.running())
-    if self.task.may_block():
-      return self.wait_until(ready_func, cancellable)
-    else:
-      assert(ready_func())
-      return Cancelled.FALSE
-
   def yield_(self, cancellable) -> Cancelled:
-    return self.yield_until(lambda: True, cancellable)
+    return self.wait_until(lambda: True, cancellable)
 
-  def switch_to(self, cancellable, other: Thread) -> Cancelled:
+  def suspend_then_resume(self, cancellable, other: Thread) -> Cancelled:
     assert(self.running() and other.suspended())
     if self.task.deliver_pending_cancel(cancellable):
-      self.stop_waiting()
       return Cancelled.TRUE
     self.cancellable = cancellable
     cancelled = block(switch_to = other)
     assert(self.running() and (cancellable or not cancelled))
     return cancelled
 
-  def yield_to(self, cancellable, other: Thread) -> Cancelled:
+  def yield_then_resume(self, cancellable, other: Thread) -> Cancelled:
     assert(self.running() and other.suspended())
+    if self.task.deliver_pending_cancel(cancellable):
+      return Cancelled.TRUE
     self.ready_func = lambda: True
     self.task.inst.store.waiting.append(self)
-    return self.switch_to(cancellable, other)
+    return self.suspend_then_resume(cancellable, other)
+
+  def suspend_then_promote(self, cancellable, other: Thread) -> ResumeArg:
+    assert(self.running())
+    if other.ready():
+      other.stop_waiting()
+      return self.suspend_then_resume(cancellable, other)
+    else:
+      return self.suspend(cancellable)
+
+  def yield_then_promote(self, cancellable, other: Thread) -> ResumeArg:
+    assert(self.running())
+    if other.ready():
+      other.stop_waiting()
+      return self.yield_then_resume(cancellable, other)
+    else:
+      return self.yield_(cancellable)
 
 ### Tasks
 
@@ -446,9 +463,6 @@ def needs_exclusive(self):
     assert(self.ft.async_)
     return not self.opts.async_ or self.opts.callback
 
-  def may_block(self):
-    return self.ft.async_ or self.state == Task.State.RESOLVED
-
   def enter_implicit_thread(self):
     assert(self.state == Task.State.INITIAL)
     self.implicit_thread = current_thread()
@@ -672,6 +686,11 @@ def remove(self, i):
     self.free.append(i)
     return e
 
+  def __iter__(self):
+    for e in self.array:
+      if e is not None:
+        yield e
+
 ### Resource State
 
 class ResourceHandle:
@@ -2114,13 +2133,12 @@ def thread_func():
       inst.exclusive_thread = None
       match code:
         case CallbackCode.YIELD:
-          cancelled = thread.yield_until(lambda: not inst.exclusive_thread, cancellable = True)
+          cancelled = thread.wait_until(lambda: not inst.exclusive_thread, cancellable = True)
           if cancelled:
             event = (EventCode.TASK_CANCELLED, 0, 0)
           else:
             event = (EventCode.NONE, 0, 0)
         case CallbackCode.WAIT:
-          trap_if(not task.may_block())
           wset = inst.handles.get(si)
           trap_if(not isinstance(wset, WaitableSet))
           event = wset.wait_for_event_and(lambda: not inst.exclusive_thread, cancellable = True)
@@ -2137,6 +2155,11 @@ def thread_func():
   task = Task(ft, opts, inst, on_start, on_resolve)
   thread = Thread(task, thread_func)
   thread.resume()
+  if not ft.async_:
+    while task.state != Task.State.RESOLVED:
+      candidates = { t for t in inst.threads if t.ready() and t is not inst.exclusive_thread }
+      trap_if(not candidates)
+      random.choice(list(candidates)).resume()
   return task.request_cancellation
 
 class CallbackCode(IntEnum):
@@ -2164,7 +2187,6 @@ def call_and_trap_on_throw(callee, args):
 def canon_lower(callee, ft, opts, flat_args: list[CoreValType]) -> list[CoreValType]:
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block() and ft.async_ and not opts.async_)
 
   subtask = Subtask()
   cx = LiftLowerContext(opts, thread.task.inst, subtask)
@@ -2345,13 +2367,12 @@ def canon_waitable_set_new():
 ### 🔀 `canon waitable-set.wait`
 
 def canon_waitable_set_wait(cancellable, mem, si, ptr):
-  task = current_task()
-  trap_if(not task.inst.may_leave)
-  trap_if(not task.may_block())
-  wset = task.inst.handles.get(si)
+  inst = current_instance()
+  trap_if(not inst.may_leave)
+  wset = inst.handles.get(si)
   trap_if(not isinstance(wset, WaitableSet))
   event = wset.wait_for_event(cancellable)
-  return unpack_event(mem, task.inst, ptr, event)
+  return unpack_event(mem, inst, ptr, event)
 
 def unpack_event(mem, inst, ptr, e: EventTuple):
   event, p1, p2 = e
@@ -2402,7 +2423,6 @@ def canon_waitable_join(wi, si):
 def canon_subtask_cancel(async_, i):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block() and not async_)
   subtask = thread.task.inst.handles.get(i)
   trap_if(not isinstance(subtask, Subtask))
   trap_if(subtask.resolve_delivered())
@@ -2463,7 +2483,6 @@ def canon_stream_write(stream_t, opts, i, ptr, n):
 def stream_copy(EndT, BufferT, event_code, stream_t, opts, i, ptr, n):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block() and not opts.async_)
 
   e = thread.task.inst.handles.get(i)
   trap_if(not isinstance(e, EndT))
@@ -2518,7 +2537,6 @@ def canon_future_write(future_t, opts, i, ptr):
 def future_copy(EndT, BufferT, event_code, future_t, opts, i, ptr):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block() and not opts.async_)
 
   e = thread.task.inst.handles.get(i)
   trap_if(not isinstance(e, EndT))
@@ -2571,7 +2589,6 @@ def canon_future_cancel_write(future_t, async_, i):
 def cancel_copy(EndT, event_code, stream_or_future_t, async_, i):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block() and not async_)
   e = thread.task.inst.handles.get(i)
   trap_if(not isinstance(e, EndT))
   trap_if(e.shared.t != stream_or_future_t.t)
@@ -2618,6 +2635,14 @@ def canon_thread_index():
   assert(thread.index is not None)
   return [thread.index]
 
+### 🧵 `canon thread.exit`
+
+def canon_thread_exit():
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  thread.exit()
+  assert(False)
+
 ### 🧵 `canon thread.new-indirect`
 
 @dataclass
@@ -2654,7 +2679,6 @@ def canon_thread_resume_later(i):
 def canon_thread_suspend(cancellable):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
-  trap_if(not thread.task.may_block())
   cancelled = thread.suspend(cancellable)
   return [cancelled]
 
@@ -2666,24 +2690,42 @@ def canon_thread_yield(cancellable):
   cancelled = thread.yield_(cancellable)
   return [cancelled]
 
-### 🧵 `canon thread.switch-to`
+### 🧵 `canon thread.suspend-then-resume`
 
-def canon_thread_switch_to(cancellable, i):
+def canon_thread_suspend_then_resume(cancellable, i):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
   other_thread = thread.task.inst.threads.get(i)
   trap_if(not other_thread.suspended())
-  cancelled = thread.switch_to(cancellable, other_thread)
+  cancelled = thread.suspend_then_resume(cancellable, other_thread)
   return [cancelled]
 
-### 🧵 `canon thread.yield-to`
+### 🧵 `canon thread.yield-then-resume`
 
-def canon_thread_yield_to(cancellable, i):
+def canon_thread_yield_then_resume(cancellable, i):
   thread = current_thread()
   trap_if(not thread.task.inst.may_leave)
   other_thread = thread.task.inst.threads.get(i)
   trap_if(not other_thread.suspended())
-  cancelled = thread.yield_to(cancellable, other_thread)
+  cancelled = thread.yield_then_resume(cancellable, other_thread)
+  return [cancelled]
+
+### 🧵 `canon thread.suspend-then-promote`
+
+def canon_thread_suspend_then_promote(cancellable, i):
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  other_thread = thread.task.inst.threads.get(i)
+  cancelled = thread.suspend_then_promote(cancellable, other_thread)
+  return [cancelled]
+
+### 🧵 `canon thread.yield-then-promote`
+
+def canon_thread_yield_then_promote(cancellable, i):
+  thread = current_thread()
+  trap_if(not thread.task.inst.may_leave)
+  other_thread = thread.task.inst.threads.get(i)
+  cancelled = thread.yield_then_promote(cancellable, other_thread)
   return [cancelled]
 
 ### 📝 `canon error-context.new`