Refactor CABI: reimplement async in terms of Thread abstraction

Author: lukewagner

design/mvp/canonical-abi/definitions.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Optional, Callable, Awaitable, TypeVar, Generic, Literal
+from typing import Any, Optional, Callable, TypeVar, Generic, Literal
 from enum import Enum, IntEnum
 import math
 import struct
@@ -214,18 +214,18 @@ class CanonicalOptions(LiftLowerOptions):
 
 ### Runtime State
 
-scheduler = asyncio.Lock()
-
 #### Component Instance State
 
 class ComponentInstance:
+  store: Store
   table: Table
   may_leave: bool
   no_backpressure: asyncio.Event
   num_backpressure_waiters: int
   lock: asyncio.Lock
 
-  def __init__(self):
+  def __init__(self, store):
+    self.store = store
     self.table = Table()
     self.may_leave = True
     self.no_backpressure = asyncio.Event()
@@ -457,10 +457,6 @@ class Cancelled(IntEnum):
   FALSE = 0
   TRUE = 1
 
-OnStart = Callable[[], list[any]]
-OnResolve = Callable[[Optional[list[any]]], None]
-OnBlock = Callable[[Awaitable], Awaitable[Cancelled]]
-
 class Task:
   class State(Enum):
     INITIAL = 1
@@ -473,24 +469,23 @@ class State(Enum):
   inst: ComponentInstance
   ft: FuncType
   supertask: Optional[Task]
-  on_resolve: OnResolve
-  on_block: OnBlock
+  thread: Thread
+  on_resolve: Callable[[Optional[list[any]]], None]
   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 self.opts.sync or self.opts.callback:
       if self.inst.lock.locked():
@@ -530,15 +525,15 @@ async def wait_on(self, awaitable, cancellable = False, for_callback = False) ->
     if for_callback:
       self.inst.lock.release()
 
-    cancelled = await self.on_block(f)
+    cancelled = await self.thread.suspend(f)
     if cancelled and not cancellable:
-      assert(await self.on_block(f) == Cancelled.FALSE)
+      assert(await self.thread.suspend(f) == Cancelled.FALSE)
 
     if for_callback:
       acquired = asyncio.create_task(self.inst.lock.acquire())
-      cancelled |= await self.on_block(acquired)
+      cancelled |= await self.thread.suspend(acquired)
       if cancelled:
-        assert(self.on_block(acquired) == Cancelled.FALSE)
+        assert(self.thread.suspend(acquired) == Cancelled.FALSE)
 
     if cancelled:
       assert(self.state == Task.State.INITIAL)
@@ -551,16 +546,6 @@ async def wait_on(self, awaitable, cancellable = False, for_callback = False) ->
     else:
       return Cancelled.FALSE
 
-  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 Cancelled.FALSE
-
-    await callee(self, on_start, on_return, sync_on_block)
-
   async def wait_for_event(self, waitable_set, cancellable, for_callback) -> EventTuple:
     if self.state == Task.State.PENDING_CANCEL and cancellable:
       self.state = Task.State.CANCEL_DELIVERED
@@ -626,18 +611,16 @@ class State(IntEnum):
     CANCELLED_BEFORE_RETURNED = 4
 
   state: State
-  task: Task
+  thread: Optional[Thread]
   lenders: Optional[list[ResourceHandle]]
-  request_cancel_begin: asyncio.Future
-  request_cancel_end: asyncio.Future
+  cancellation_requested: bool
 
-  def __init__(self, task):
+  def __init__(self):
     Waitable.__init__(self)
     self.state = Subtask.State.STARTING
-    self.task = task
+    self.thread = None
     self.lenders = []
-    self.request_cancel_begin = asyncio.Future()
-    self.request_cancel_end = asyncio.Future()
+    self.cancellation_requested = False
 
   def resolved(self):
     match self.state:
@@ -649,44 +632,6 @@ def resolved(self):
             Subtask.State.CANCELLED_BEFORE_RETURNED):
         return True
 
-  async def request_cancel(self):
-    assert(not self.cancellation_requested() and not self.resolved())
-    self.request_cancel_begin.set_result(None)
-    await self.request_cancel_end
-
-  def cancellation_requested(self):
-    return self.request_cancel_begin.done()
-
-  async def call_async(self, callee, on_start, on_resolve):
-    async def do_call():
-      await callee(self.task, on_start, on_resolve, async_on_block)
-      relinquish_control()
-
-    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 Cancelled.TRUE
-      else:
-        await awaitable
-      assert(awaitable.done())
-      await scheduler.acquire()
-      return Cancelled.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
-
   def add_lender(self, lending_handle):
     assert(not self.resolve_delivered() and not self.resolved())
     lending_handle.num_lends += 1
@@ -927,6 +872,84 @@ def drop(self):
     trap_if(not self.done)
     FutureEnd.drop(self)
 
+#### Thread State
+
+class Thread:
+  store: Store
+  future: Optional[asyncio.Future]
+  on_resume: Optional[asyncio.Future]
+  on_suspend_or_exit: Optional[asyncio.Future]
+  returned: bool
+
+  def __init__(self, store, lifted_func, caller, on_start, on_resolve):
+    self.store = store
+    self.future = None
+    self.on_resume = asyncio.Future()
+    self.on_suspend_or_exit = None
+    self.returned = False
+    async def async_impl():
+      assert(await self.on_resume == Cancelled.FALSE)
+      self.on_resume = None
+      await lifted_func(caller, self, on_start, on_resolve)
+      self.on_suspend_or_exit.set_result(None)
+      self.returned = True
+    asyncio.create_task(async_impl())
+
+  async def resume(self, cancelled = Cancelled.FALSE):
+    if self.future:
+      assert(cancelled or self.future.done())
+      self.future = None
+      self.store.waiting.remove(self)
+    self.on_resume.set_result(cancelled)
+    assert(not self.on_suspend_or_exit)
+    self.on_suspend_or_exit = asyncio.Future()
+    await self.on_suspend_or_exit
+    self.on_suspend_or_exit = None
+    if self.future:
+      self.store.waiting.append(self)
+
+  async def suspend(self, future) -> Cancelled:
+    assert(not self.future)
+    self.future = future
+    self.on_suspend_or_exit.set_result(None)
+    self.on_suspend_or_exit = None
+    assert(not self.on_resume)
+    self.on_resume = asyncio.Future()
+    cancelled = await self.on_resume
+    self.on_resume = None
+    return cancelled
+
+#### Store State / Embedding API
+
+class Store:
+  loop: asyncio.AbstractEventLoop
+  waiting: list[Thread]
+
+  def __init__(self):
+    self.loop = asyncio.new_event_loop()
+    self.waiting = []
+
+  ExportCall = Thread
+
+  def start_export_call(self, lifted_func, on_start, on_resolve) -> ExportCall:
+    async def async_impl():
+      caller = None
+      thread = Thread(self, lifted_func, caller, on_start, on_resolve)
+      await thread.resume()
+      return thread
+    return self.loop.run_until_complete(async_impl())
+
+  def tick(self):
+    if not DETERMINISTIC_PROFILE:
+      random.shuffle(self.waiting)
+    for thread in self.waiting:
+      if thread.future.done():
+        self.loop.run_until_complete(thread.resume())
+        return
+
+  def export_call_finished(self, export_call: ExportCall):
+    return export_call.returned
+
 ### Despecialization
 
 def despecialize(t):
@@ -1882,8 +1905,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 await task.enter() == Cancelled.TRUE:
     task.cancel()
     task.exit()
@@ -1958,7 +1981,7 @@ async def call_and_trap_on_throw(callee, task, args):
 
 async def canon_lower(opts, ft, callee, task, flat_args):
   trap_if(not task.inst.may_leave)
-  subtask = Subtask(task)
+  subtask = Subtask()
 
   cx = LiftLowerContext(opts, task.inst, subtask)
   flat_ft = flatten_functype(opts, ft, 'lower')
@@ -1984,7 +2007,7 @@ def on_start():
   def on_resolve(result):
     on_progress()
     if result is None:
-      assert(subtask.cancellation_requested())
+      assert(subtask.cancellation_requested)
       if subtask.state == Subtask.State.STARTING:
         subtask.state = Subtask.State.CANCELLED_BEFORE_STARTED
       else:
@@ -1996,13 +2019,19 @@ def on_resolve(result):
       nonlocal flat_results
       flat_results = lower_flat_values(cx, max_flat_results, result, ft.result_type(), flat_args)
 
+  subtask.thread = Thread(task.inst.store, callee, task, on_start, on_resolve)
+  await subtask.thread.resume()
+
   if opts.sync:
-    await task.call_sync(callee, on_start, on_resolve)
+    if not subtask.resolved():
+      done = asyncio.Event()
+      def on_progress():
+        done.set()
+      await task.wait_on(done.wait())
     assert(types_match_values(flat_ft.results, flat_results))
     subtask.deliver_resolve()
     return flat_results
   else:
-    await subtask.call_async(callee, on_start, on_resolve)
     if subtask.resolved():
       assert(flat_results == [])
       subtask.deliver_resolve()
@@ -2182,11 +2211,12 @@ async def canon_subtask_cancel(sync, task, i):
   subtask = task.inst.table.get(i)
   trap_if(not isinstance(subtask, Subtask))
   trap_if(subtask.resolve_delivered())
-  trap_if(subtask.cancellation_requested())
+  trap_if(subtask.cancellation_requested)
   if subtask.resolved():
     assert(subtask.has_pending_event())
   else:
-    await subtask.request_cancel()
+    subtask.cancellation_requested = True
+    await subtask.thread.resume(Cancelled.TRUE)
     if sync:
       while not subtask.resolved():
         if subtask.has_pending_event():