WIP: tweak coop threads

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()
@@ -392,8 +393,10 @@ def pending(self):
     return not self.running() and self.ready_func is not None
 
   def ready(self):
-    assert(self.pending())
-    return self.ready_func()
+    return self.pending() and self.ready_func()
+
+  def ready_list(inst: ComponentInstance) -> list[Thread]:
+    return [t for t in inst.threads.array if t and t.ready()]
 
   def __init__(self, task, thread_func):
     self.task = task
@@ -410,17 +413,48 @@ def fiber_func():
       self.fiber_lock.acquire()
       assert(self.running() and self.suspend_result == SuspendResult.NOT_CANCELLED)
       self.suspend_result = None
-      thread_func(self)
-      assert(self.running())
-      self.task.thread_stop(self)
-      if self.index is not None:
-        self.task.inst.threads.remove(self.index)
-      self.parent_lock.release()
+      try:
+        thread_func(self)
+        self.exit()
+      except ThreadExit:
+        return
+      assert(False)
     self.fiber = threading.Thread(target = fiber_func)
     self.fiber.start()
     self.task.thread_start(self)
     assert(self.suspended())
 
+  def exit(self):
+    assert(self.running() and self.task.may_block())
+    if self.index is not None:
+      self.task.inst.threads.remove(self.index)
+    self.task.thread_stop(self)
+    if self.task.is_sync_before_return():
+      self.release_any_ready_sibling()
+    else:
+      self.parent_lock.release()
+    raise ThreadExit()
+
+  def release_any_ready_sibling(self):
+    other = random.choice(Thread.ready_list(self.task.inst))
+    assert(self is not other)
+    other.clear_pending()
+    self.release_sibling(other)
+
+  def clear_pending(self, suspend_result = SuspendResult.NOT_CANCELLED):
+    assert(suspend_result == SuspendResult.CANCELLED or self.ready_func())
+    self.ready_func = None
+    self.task.inst.store.pending.remove(self)
+
+  def release_sibling(self, other: Thread):
+    assert(other.suspend_result is None)
+    other.suspend_result = SuspendResult.NOT_CANCELLED
+    assert(self.parent_lock and not other.parent_lock)
+    other.parent_lock = self.parent_lock
+    self.parent_lock = None
+    assert(not self.running() and other.running())
+    other.fiber_lock.release()
+
   def resume_later(self):
     assert(self.suspended())
     self.ready_func = lambda: True
@@ -429,9 +463,7 @@ def resume_later(self):
   def resume(self, suspend_result = SuspendResult.NOT_CANCELLED):
     assert(not self.running())
     if self.pending():
-      assert(suspend_result == SuspendResult.CANCELLED or self.ready_func())
-      self.ready_func = None
-      self.task.inst.store.pending.remove(self)
+      self.clear_pending(suspend_result)
     assert(self.cancellable or suspend_result == SuspendResult.NOT_CANCELLED)
     assert(self.suspend_result is None)
     self.suspend_result = suspend_result
@@ -446,7 +478,10 @@ def suspend(self, cancellable) -> SuspendResult:
     assert(self.running() and self.task.may_block())
     assert(not self.cancellable and self.suspend_result is None)
     self.cancellable = cancellable
-    self.parent_lock.release()
+    if self.task.is_sync_before_return():
+      self.release_any_ready_sibling()
+    else:
+      self.parent_lock.release()
     self.fiber_lock.acquire()
     assert(self.running())
     self.cancellable = False
@@ -464,16 +499,11 @@ def suspend_until(self, ready_func, cancellable = False) -> SuspendResult:
     self.task.inst.store.pending.append(self)
     return self.suspend(cancellable)
 
-  def switch_to(self, cancellable, other: Thread) -> SuspendResult:
+  def suspend_to_suspended(self, cancellable, other: Thread) -> SuspendResult:
     assert(self.running() and other.suspended())
-    assert(not self.cancellable and self.suspend_result is None and other.suspend_result is None)
+    assert(not self.cancellable and self.suspend_result is None)
     self.cancellable = cancellable
-    other.suspend_result = SuspendResult.NOT_CANCELLED
-    assert(self.parent_lock and not other.parent_lock)
-    other.parent_lock = self.parent_lock
-    self.parent_lock = None
-    assert(not self.running() and other.running())
-    other.fiber_lock.release()
+    self.release_sibling(other)
     self.fiber_lock.acquire()
     assert(self.running())
     self.cancellable = False
@@ -483,11 +513,21 @@ def switch_to(self, cancellable, other: Thread) -> SuspendResult:
     assert(cancellable or suspend_result == SuspendResult.NOT_CANCELLED)
     return suspend_result
 
-  def yield_to(self, cancellable, other: Thread) -> SuspendResult:
-    assert(not self.ready_func)
+  def yield_to_suspended(self, cancellable, other: Thread) -> SuspendResult:
+    assert(self.running() and other.suspended())
     self.ready_func = lambda: True
     self.task.inst.store.pending.append(self)
-    return self.switch_to(cancellable, other)
+    return self.suspend_to_suspended(cancellable, other)
+
+  def suspend_to_ready(self, cancellable, other: Thread) -> SuspendResult:
+    assert(self.running() and other.ready())
+    other.clear_pending()
+    return self.suspend_to_suspended(cancellable, other)
+
+  def yield_to_ready(self, cancellable, other: Thread) -> SuspendResult:
+    assert(self.running() and other.ready())
+    other.clear_pending()
+    return self.yield_to_suspended(cancellable, other)
 
 #### Waitable State
 
@@ -597,8 +637,11 @@ def thread_stop(self, thread):
   def needs_exclusive(self):
     return not self.opts.async_ or self.opts.callback
 
+  def is_sync_before_return(self):
+    return not self.ft.async_ and self.state != Task.State.RESOLVED
+
   def may_block(self):
-    return self.ft.async_ or self.state == Task.State.RESOLVED
+    return not self.is_sync_before_return() or Thread.ready_list(self.inst)
 
   def enter(self, thread):
     assert(thread in self.threads and thread.task is self)
@@ -1998,10 +2041,7 @@ def thread_func(thread):
       inst.exclusive = False
       match code:
         case CallbackCode.YIELD:
-          if task.may_block():
-            event = task.yield_until(lambda: not inst.exclusive, thread, cancellable = True)
-          else:
-            event = (EventCode.NONE, 0, 0)
+          event = task.yield_until(lambda: not inst.exclusive, thread, cancellable = True)
         case CallbackCode.WAIT:
           trap_if(not task.may_block())
           wset = inst.handles.get(si)
@@ -2514,6 +2554,13 @@ def canon_thread_resume_later(thread, i):
   other_thread.resume_later()
   return []
 
+### 🧵 `canon thread.exit`
+
+def canon_thread_exit(thread):
+  trap_if(not thread.task.inst.may_leave)
+  thread.exit()
+  assert(False)
+
 ### 🧵 `canon thread.suspend`
 
 def canon_thread_suspend(cancellable, thread):
@@ -2529,60 +2576,62 @@ def canon_thread_suspend(cancellable, thread):
 
 def canon_thread_yield(cancellable, thread):
   trap_if(not thread.task.inst.may_leave)
-  if not thread.task.may_block():
-    return [SuspendResult.NOT_CANCELLED]
   if thread.task.deliver_pending_cancel(cancellable):
     suspend_result = SuspendResult.CANCELLED
   else:
     suspend_result = thread.suspend_until(lambda:True, cancellable)
   return [suspend_result]
 
-### 🧵 `canon thread.switch-to`
+### 🧵 `canon thread.suspend-to-suspended`
 
-def canon_thread_switch_to(cancellable, thread, i):
+def canon_thread_suspend_to_suspended(cancellable, thread, i):
   trap_if(not thread.task.inst.may_leave)
   other_thread = thread.task.inst.threads.get(i)
   trap_if(not other_thread.suspended())
   if thread.task.deliver_pending_cancel(cancellable):
     suspend_result = SuspendResult.CANCELLED
   else:
-    suspend_result = thread.switch_to(cancellable, other_thread)
+    suspend_result = thread.suspend_to_suspended(cancellable, other_thread)
   return [suspend_result]
 
-### 🧵 `canon thread.yield-to`
+### 🧵 `canon thread.yield-to-suspended`
 
-def canon_thread_yield_to(cancellable, thread, i):
+def canon_thread_yield_to_suspended(cancellable, thread, i):
   trap_if(not thread.task.inst.may_leave)
   other_thread = thread.task.inst.threads.get(i)
   trap_if(not other_thread.suspended())
   if thread.task.deliver_pending_cancel(cancellable):
     suspend_result = SuspendResult.CANCELLED
   else:
-    suspend_result = thread.yield_to(cancellable, other_thread)
+    suspend_result = thread.yield_to_suspended(cancellable, other_thread)
   return [suspend_result]
 
-### 🧵 `canon thread.switch-to`
+### 🧵 `canon thread.suspend-then-promote`
 
-def canon_thread_switch_to(cancellable, thread, i):
+def canon_thread_suspend_then_promote(cancellable, thread, i):
   trap_if(not thread.task.inst.may_leave)
+  trap_if(not thread.task.may_block())
   other_thread = thread.task.inst.threads.get(i)
-  trap_if(not other_thread.suspended())
   if thread.task.deliver_pending_cancel(cancellable):
     suspend_result = SuspendResult.CANCELLED
+  elif other_thread.ready():
+    suspend_result = thread.suspend_to_ready(cancellable, other_thread)
   else:
-    suspend_result = thread.switch_to(cancellable, other_thread)
+    suspend_result = thread.suspend(cancellable)
   return [suspend_result]
 
-### 🧵 `canon thread.yield-to`
+### 🧵 `canon thread.yield-then-promote`
 
-def canon_thread_yield_to(cancellable, thread, i):
+def canon_thread_yield_then_promote(cancellable, thread, i):
   trap_if(not thread.task.inst.may_leave)
+  trap_if(not thread.task.may_block())
   other_thread = thread.task.inst.threads.get(i)
-  trap_if(not other_thread.suspended())
   if thread.task.deliver_pending_cancel(cancellable):
     suspend_result = SuspendResult.CANCELLED
+  elif other_thread.ready():
+    suspend_result = thread.yield_to_ready(cancellable, other_thread)
   else:
-    suspend_result = thread.yield_to(cancellable, other_thread)
+    suspend_result = thread.suspend_until(lambda:True, cancellable)
   return [suspend_result]
 
 ### 📝 `canon error-context.new`