WIP: cooperative threads

Author: lukewagner

design/mvp/canonical-abi/definitions.py

@@ -214,11 +214,10 @@ class CanonicalOptions(LiftLowerOptions):
 
 ### Runtime State
 
-scheduler = asyncio.Lock()
-
 #### Component Instance State
 
 class ComponentInstance:
+  store: Store
   table: Table
   may_leave: bool
   backpressure: bool
@@ -455,6 +454,51 @@ def drop(self):
     trap_if(len(self.elems) > 0)
     trap_if(self.num_waiting > 0)
 
+#### Thread State
+
+class Cancelled(IntEnum):
+  FALSE = 0
+  TRUE = 1
+
+class Thread:
+  store: Store
+  awaitable: Optional[Awaitable]
+  on_resume: Optional[asyncio.Future]
+  on_suspend_or_exit: Optional[asyncio.Future]
+
+  def __init__(self, store, callee, on_start, on_resolve):
+    self.store = store
+    self.awaitable = None
+    self.on_resume = asyncio.Future()
+    self.on_suspend_or_exit = None
+    async def thread():
+      assert(await self.on_resume == Cancelled.FALSE)
+      callee(self, on_start, on_resolve)
+      self.on_suspend_or_exit.set_result(None)
+    asyncio.create_task(thread())
+    store.waiting.append(self)
+
+  async def resume(self, cancelled = Cancelled.FALSE):
+    assert(cancelled or (not self.awaitable or self.awaitable.done()))
+    self.awaitable = None
+    self.store.waiting.remove(self)
+    self.on_resume.set_result(cancelled)
+    self.on_resume = None
+    assert(not self.on_suspend_or_exit)
+    self.on_suspend_or_exit = asyncio.Future()
+    await self.on_suspend_or_exit
+    assert(self.awaitable)
+    self.store.waiting.append(self)
+
+  async def suspend(self, awaitable) -> Cancelled:
+    assert(not self.awaitable)
+    self.awaitable = awaitable
+    self.on_suspend_or_exit.set_result(None)
+    self.on_suspend_or_exit = None
+    assert(not self.on_resume)
+    self.on_resume = asyncio.Future()
+    return await thread.on_resume
+
 #### Task State
 
 class Cancelled(IntEnum):
@@ -477,29 +521,28 @@ class State(Enum):
   inst: ComponentInstance
   ft: FuncType
   supertask: Optional[Task]
+  thread: Thread
   on_resolve: OnResolve
-  on_block: OnBlock
   num_borrows: int
   context: ContextLocalStorage
 
-  def __init__(self, opts, inst, ft, supertask, on_resolve, on_block):
+  def __init__(self, opts, inst, ft, supertask, thread, on_resolve):
     self.state = Task.State.INITIAL
     self.opts = opts
     self.inst = inst
     self.ft = ft
     self.supertask = supertask
+    self.thread = thread
     self.on_resolve = on_resolve
-    self.on_block = on_block
     self.num_borrows = 0
     self.context = ContextLocalStorage()
 
   async def enter(self):
-    assert(scheduler.locked())
     self.trap_if_on_the_stack(self.inst)
     if not self.may_enter(self) or self.inst.pending_tasks:
       f = asyncio.Future()
       self.inst.pending_tasks.append((self, f))
-      if await self.on_block(f) == Cancelled.TRUE:
+      if await self.thread.suspend(f) == Cancelled.TRUE:
         [i] = [i for i,(t,_) in enumerate(self.inst.pending_tasks) if t == self]
         self.inst.pending_tasks.pop(i)
         self.on_resolve(None)
@@ -531,44 +574,28 @@ def maybe_start_pending_task(self):
         pending_future.set_result(None)
         return
 
-  async def sync_wait(self, awaitable, cancellable = False) -> Cancelled:
+  async def sync_wait(self, awaitable) -> None:
     awaitable = asyncio.ensure_future(awaitable)
     if awaitable.done() and not DETERMINISTIC_PROFILE and random.randint(0,1):
       return
     assert(self.inst.unblocked.is_set())
     self.inst.unblocked.clear()
-    cancelled = await self.on_block(awaitable)
-    if cancelled and not cancellable:
+    if await self.thread.suspend(awaitable) == Cancelled.TRUE:
       assert(self.state == Task.State.INITIAL)
       self.state = Task.State.PENDING_CANCEL
-      cancelled = await self.on_block(awaitable)
-      assert(not cancelled)
+      assert(await self.thread.suspend(awaitable) == Cancelled.FALSE)
     self.inst.unblocked.set()
-    return cancelled
 
   async def async_wait(self, awaitable) -> Cancelled:
     self.maybe_start_pending_task()
     awaitable = asyncio.ensure_future(awaitable)
     if awaitable.done() and not DETERMINISTIC_PROFILE and random.randint(0,1):
       return
-    cancelled = await self.on_block(awaitable)
+    cancelled = await self.thread.suspend(awaitable)
     while not self.inst.unblocked.is_set():
-      cancelled |= await self.on_block(self.inst.unblocked.wait())
+      cancelled |= await self.thread.suspend(self.inst.unblocked.wait())
     return cancelled
 
-  async def call_sync(self, callee, on_start, on_return):
-    async def sync_on_block(awaitable):
-      if await self.on_block(awaitable) == Cancelled.TRUE:
-        assert(self.state == Task.State.INITIAL)
-        self.state = Task.State.PENDING_CANCEL
-        assert(await self.on_block(awaitable) == Cancelled.FALSE)
-      return False
-
-    assert(self.inst.unblocked.is_set())
-    self.inst.unblocked.clear()
-    await callee(self, on_start, on_return, sync_on_block)
-    self.inst.unblocked.set()
-
   async def wait_for_event(self, waitable_set, sync) -> EventTuple:
     if self.state == Task.State.PENDING_CANCEL:
       self.state = Task.State.CANCEL_DELIVERED
@@ -579,14 +606,13 @@ async def wait_for_event(self, waitable_set, sync) -> EventTuple:
       while not e:
         maybe_event = waitable_set.maybe_has_pending_event.wait()
         if sync:
-          cancelled = await self.sync_wait(maybe_event, cancellable = True)
+          await self.sync_wait(maybe_event)
         else:
-          cancelled = await self.async_wait(maybe_event)
-        if cancelled:
-          assert(self.state == Task.State.INITIAL)
-          self.state = Task.State.CANCEL_DELIVERED
-          e = (EventCode.TASK_CANCELLED, 0, 0)
-          break
+          if await self.async_wait(maybe_event) == Cancelled.TRUE:
+            assert(self.state == Task.State.INITIAL)
+            self.state = Task.State.CANCEL_DELIVERED
+            e = (EventCode.TASK_CANCELLED, 0, 0)
+            break
         e = waitable_set.poll()
       waitable_set.num_waiting -= 1
     return e
@@ -595,12 +621,11 @@ async def yield_(self, sync) -> EventTuple:
     if self.state == Task.State.PENDING_CANCEL:
       self.state = Task.State.CANCEL_DELIVERED
       return (EventCode.TASK_CANCELLED, 0, 0)
+    if sync:
+      await self.sync_wait(asyncio.sleep(0))
+      return (EventCode.NONE, 0, 0)
     else:
-      if sync:
-        cancelled = await self.sync_wait(asyncio.sleep(0), cancellable = True)
-      else:
-        cancelled = await self.async_wait(asyncio.sleep(0))
-      if cancelled:
+      if await self.async_wait(asyncio.sleep(0)) == Cancelled.TRUE:
         assert(self.state == Task.State.INITIAL)
         self.state = Task.State.CANCEL_DELIVERED
         return (EventCode.TASK_CANCELLED, 0, 0)
@@ -630,7 +655,6 @@ def cancel(self):
     self.state = Task.State.RESOLVED
 
   def exit(self):
-    assert(scheduler.locked())
     trap_if(self.state != Task.State.RESOLVED)
     assert(self.num_borrows == 0)
     if self.opts.sync:
@@ -650,34 +674,46 @@ class State(IntEnum):
 
   state: State
   task: Task
+  callee_thread: Optional[Thread]
   lenders: Optional[list[ResourceHandle]]
-  request_cancel_begin: asyncio.Future
-  request_cancel_end: asyncio.Future
+  cancelled: bool
 
   def __init__(self, task):
     Waitable.__init__(self)
     self.state = Subtask.State.STARTING
     self.task = task
+    self.callee_thread = None
     self.lenders = []
-    self.request_cancel_begin = asyncio.Future()
-    self.request_cancel_end = asyncio.Future()
+    self.cancelled = False
 
-  async def call_sync(self, callee, on_start, on_resolve):
-    def sync_on_start():
+  async def call(self, callee, on_start, on_resolve):
+    def update_on_start():
       assert(self.state == Subtask.State.STARTING)
       self.state = Subtask.State.STARTED
       return on_start()
 
-    def sync_on_resolve(result):
-      assert(result is not None)
+    def update_on_resolve(result):
       assert(self.state == Subtask.State.STARTED)
-      self.state = Subtask.State.RETURNED
+      if result is None:
+        assert(self.cancelled)
+        if self.state == Subtask.State.STARTING:
+          self.state = Subtask.State.CANCELLED_BEFORE_STARTED
+        else:
+          assert(self.state == Subtask.State.STARTED)
+          self.state = Subtask.State.CANCELLED_BEFORE_RETURNED
+      else:
+        assert(self.state == Subtask.State.STARTED)
+        self.state = Subtask.State.RETURNED
       on_resolve(result)
 
-    await Task.call_sync(self.task, callee, sync_on_start, sync_on_resolve)
+    assert(not self.callee_thread)
+    self.callee_thread = Thread(self.supertask.inst.store, callee, update_on_start, update_on_resolve)
+    await self.callee_thread.resume()
 
-  def cancelled(self):
-    return self.request_cancel_begin.done()
+  async def wait_until_resolved(self):
+    while not self.resolved():
+      await self.supertask.sync_wait(self.callee_thread.awaitable)
+      await self.callee_thread.resume()
 
   def resolved(self):
     match self.state:
@@ -689,57 +725,13 @@ def resolved(self):
             Subtask.State.CANCELLED_BEFORE_RETURNED):
         return True
 
-  async def request_cancel(self):
-    assert(not self.cancelled() and not self.resolved())
-    self.request_cancel_begin.set_result(None)
-    await self.request_cancel_end
-
-  async def call_async(self, callee, on_start, on_resolve):
-    async def do_call():
-      await callee(self.task, async_on_start, async_on_resolve, async_on_block)
-      relinquish_control()
+  def cancellation_requested(self):
+    return self.cancelled
 
-    def async_on_start():
-      assert(self.state == Subtask.State.STARTING)
-      self.state = Subtask.State.STARTED
-      return on_start()
-
-    def async_on_resolve(result):
-      if result is None:
-        if self.state == Subtask.State.STARTING:
-          self.state = Subtask.State.CANCELLED_BEFORE_STARTED
-        else:
-          assert(self.state == Subtask.State.STARTED)
-          self.state = Subtask.State.CANCELLED_BEFORE_RETURNED
-      else:
-        assert(self.state == Subtask.State.STARTED)
-        self.state = Subtask.State.RETURNED
-      on_resolve(result)
-
-    async def async_on_block(awaitable):
-      relinquish_control()
-      if not self.request_cancel_end.done():
-        await asyncio.wait([awaitable, self.request_cancel_begin],
-                           return_when = asyncio.FIRST_COMPLETED)
-        if self.request_cancel_begin.done():
-          return True
-      else:
-        await awaitable
-      assert(awaitable.done())
-      await scheduler.acquire()
-      return False
-
-    def relinquish_control():
-      if not ret.done():
-        ret.set_result(None)
-      elif self.request_cancel_begin.done() and not self.request_cancel_end.done():
-        self.request_cancel_end.set_result(None)
-      else:
-        scheduler.release()
-
-    ret = asyncio.Future()
-    asyncio.create_task(do_call())
-    await ret
+  async def request_cancel(self):
+    assert(not self.cancellation_requested() and not self.resolved())
+    self.cancelled = True
+    await self.callee_thread.resume(Cancelled.TRUE)
 
   def add_lender(self, lending_handle):
     assert(not self.resolve_delivered() and not self.resolved())
@@ -981,6 +973,27 @@ def drop(self):
     trap_if(not self.done)
     FutureEnd.drop(self)
 
+### Store State
+
+class Store:
+  loop: asyncio.AbstractEventLoop
+  waiting: list[Thread]
+
+  def __init__(self):
+    self.loop = asyncio.new_event_loop()
+    self.waiting = []
+
+  def call_export(self, callee, on_start, on_resolve):
+    self.run_until_complete(Thread(self, callee, on_start, on_resolve).resume())
+
+  def tick(self, i):
+    if not DETERMINISTIC_PROFILE:
+      random.shuffle(self.waiting)
+    for thread in self.waiting:
+      if thread.awaitable.done():
+        self.run_until_complete(thread.resume())
+        return
+
 ### Despecialization
 
 def despecialize(t):
@@ -1936,8 +1949,8 @@ def lower_flat_values(cx, max_flat, vs, ts, out_param = None):
 
 ### `canon lift`
 
-async def canon_lift(opts, inst, ft, callee, caller, on_start, on_resolve, on_block):
-  task = Task(opts, inst, ft, caller, on_resolve, on_block)
+async def canon_lift(opts, inst, ft, callee, caller, thread, on_start, on_resolve):
+  task = Task(opts, inst, ft, caller, thread, on_resolve)
   if not await task.enter():
     return
 
@@ -2022,10 +2035,13 @@ def on_start():
 
     flat_results = None
     def on_resolve(result):
+      assert(result is not None)
       nonlocal flat_results
       flat_results = lower_flat_values(cx, MAX_FLAT_RESULTS, result, ft.result_type(), flat_args)
 
-    await subtask.call_sync(callee, on_start, on_resolve)
+    await subtask.call(callee, on_start, on_resolve)
+    await subtask.wait_until_resolved()
+
     assert(types_match_values(flat_ft.results, flat_results))
     subtask.deliver_resolve()
     return flat_results
@@ -2048,7 +2064,7 @@ def subtask_event():
         return (EventCode.SUBTASK, subtaski, subtask.state)
       subtask.set_event(subtask_event)
 
-  await subtask.call_async(callee, on_start, on_resolve)
+  await subtask.call(callee, on_start, on_resolve)
   if subtask.resolved():
     subtask.deliver_resolve()
     return [Subtask.State.RETURNED]
@@ -2217,16 +2233,13 @@ async def canon_subtask_cancel(sync, task, i):
   trap_if(not task.inst.may_leave)
   subtask = task.inst.table.get(i)
   trap_if(not isinstance(subtask, Subtask))
-  trap_if(subtask.resolve_delivered() or subtask.cancelled())
+  trap_if(subtask.resolve_delivered() or subtask.cancellation_requested())
   if subtask.resolved():
     assert(subtask.has_pending_event())
   else:
     await subtask.request_cancel()
     if sync:
-      while not subtask.resolved():
-        if subtask.has_pending_event():
-          _ = subtask.get_event()
-        await task.sync_wait(subtask.wait_for_pending_event())
+      await subtask.wait_until_resolved()
     else:
       if not subtask.resolved():
         return [BLOCKED]