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()
@@ -395,8 +396,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
@@ -413,28 +416,57 @@ def fiber_func():
       self.fiber_lock.acquire()
       assert(self.running() and not cancelled(self.suspend_result))
       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 resume_later(self):
+  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_some_ready_sibling()
+    else:
+      self.parent_lock.release()
+    raise ThreadExit()
+
+  def release_some_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(cancelled(suspend_result) 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 unsuspend(self):
     assert(self.suspended())
     self.ready_func = lambda: True
     self.task.inst.store.pending.append(self)
 
   def resume(self, suspend_result = SuspendResult.NOT_CANCELLED):
     assert(not self.running())
     if self.pending():
-      assert(cancelled(suspend_result) or self.ready_func())
-      self.ready_func = None
-      self.task.inst.store.pending.remove(self)
+      self.clear_pending(suspend_result)
     assert(self.cancellable or not cancelled(suspend_result))
     assert(self.suspend_result is None)
     self.suspend_result = suspend_result
@@ -449,7 +481,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_some_ready_sibling()
+    else:
+      self.parent_lock.release()
     self.fiber_lock.acquire()
     assert(self.running())
     self.cancellable = False
@@ -477,16 +512,11 @@ def yield_until(self, ready_func, cancellable) -> SuspendResult:
   def yield_(self, cancellable) -> SuspendResult:
     return self.yield_until(lambda: True, 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
@@ -496,11 +526,21 @@ def switch_to(self, cancellable, other: Thread) -> SuspendResult:
     assert(cancellable or not cancelled(suspend_result))
     return suspend_result
 
-  def yield_to(self, cancellable, other: Thread) -> SuspendResult:
+  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
 
@@ -624,8 +664,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 bool(Thread.ready_list(self.inst))
 
   def enter(self, thread):
     assert(thread in self.threads and thread.task is self)
@@ -2513,15 +2556,22 @@ def thread_func(thread):
   new_thread.index = thread.task.inst.threads.add(new_thread)
   return [new_thread.index]
 
-### 🧵 `canon thread.resume-later`
+### 🧵 `canon thread.unsuspend`
 
-def canon_thread_resume_later(thread, i):
+def canon_thread_unsuspend(thread, i):
   trap_if(not thread.task.inst.may_leave)
   other_thread = thread.task.inst.threads.get(i)
   trap_if(not other_thread.suspended())
-  other_thread.resume_later()
+  other_thread.unsuspend()
   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):
@@ -2543,28 +2593,55 @@ def canon_thread_yield(cancellable, thread):
     suspend_result = thread.yield_(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.suspend-then-promote`
+
+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)
+  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.suspend(cancellable)
+  return [suspend_result]
+
+### 🧵 `canon thread.yield-then-promote`
+
+def canon_thread_yield_then_promote(cancellable, thread, i):
+  trap_if(not thread.task.inst.may_leave)
+  other_thread = thread.task.inst.threads.get(i)
+  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_(cancellable)
   return [suspend_result]
 
 ### 📝 `canon error-context.new`