CABI: rebase async onto Thread abstraction (no change)

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
@@ -227,7 +226,8 @@ class ComponentInstance:
   pending_tasks: list[tuple[Task, asyncio.Future]]
   starting_pending_task: bool
 
-  def __init__(self):
+  def __init__(self, store):
+    self.store = store
     self.table = Table()
     self.may_leave = True
     self.backpressure = False
@@ -461,10 +461,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
@@ -477,29 +473,28 @@ 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 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)
@@ -537,37 +532,24 @@ async def sync_wait(self, awaitable) -> None:
       return
     assert(self.inst.unblocked.is_set())
     self.inst.unblocked.clear()
-    if await self.on_block(awaitable) == Cancelled.TRUE:
+    if await self.thread.suspend(awaitable) == Cancelled.TRUE:
       assert(self.state == Task.State.INITIAL)
       self.state = Task.State.PENDING_CANCEL
-      assert(await self.on_block(awaitable) == Cancelled.FALSE)
+      assert(await self.thread.suspend(awaitable) == Cancelled.FALSE)
     self.inst.unblocked.set()
 
   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.FALSE
-    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())
     if cancelled:
       awaitable.cancel()
     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 Cancelled.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, wset, sync) -> EventTuple:
     if self.state == Task.State.PENDING_CANCEL:
       self.state = Task.State.CANCEL_DELIVERED
@@ -626,7 +608,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:
@@ -646,17 +627,17 @@ class State(IntEnum):
 
   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):
     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:
@@ -668,44 +649,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
@@ -946,6 +889,82 @@ def drop(self):
     trap_if(not self.done)
     FutureEnd.drop(self)
 
+#### Thread State
+
+class Thread:
+  store: Store
+  awaitable: Optional[Awaitable]
+  on_resume: Optional[asyncio.Future]
+  on_suspend_or_exit: Optional[asyncio.Future]
+
+  def __init__(self, store, lifted_func, caller, on_start, on_resolve):
+    self.store = store
+    self.awaitable = asyncio.Future()
+    self.on_resume = asyncio.Future()
+    self.on_suspend_or_exit = None
+    async def coro():
+      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)
+    asyncio.create_task(coro())
+    store.waiting.append(self)
+    self.awaitable.set_result(None)
+
+  async def resume(self, cancelled = Cancelled.FALSE):
+    assert(cancelled or self.awaitable.done())
+    self.awaitable = 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.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()
+    result = await self.on_resume
+    self.on_resume = None
+    return result
+
+#### 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 coro():
+      caller = None
+      thread = Thread(self, lifted_func, caller, on_start, on_resolve)
+      await thread.resume()
+      return thread
+    return self.loop.run_until_complete(coro())
+
+  def tick(self):
+    if not DETERMINISTIC_PROFILE:
+      random.shuffle(self.waiting)
+    for thread in self.waiting:
+      if thread.awaitable.done():
+        self.loop.run_until_complete(thread.resume())
+        return
+
+  def export_call_finished(self, export_call: ExportCall):
+    return export_call.awaitable is None
+
 ### Despecialization
 
 def despecialize(t):
@@ -1901,8 +1920,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
 
@@ -2001,7 +2020,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:
@@ -2013,13 +2032,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.sync_wait(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()
@@ -2196,11 +2221,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():