cancellation w/ asyncio -> greenlet -> asyncio

test.py

import asyncio
import contextvars
import sys

import asyncpg
import greenlet


current_greenlet_context = contextvars.ContextVar("current greenlet context")


def await_(coroutine):
    current_greenlet = current_greenlet_context.get()
    if current_greenlet is None:
        raise Exception(
            "not running inside a greenlet right now, "
            "can't use await_() function"
        )
    return current_greenlet.switch(coroutine)


async def greenlet_spawn(fn, *args):

    result_future = asyncio.Future()

    def run_greenlet_target():
        result_future.set_result(fn(*args))
        return None

    async def run_greenlet():
        gl = greenlet.greenlet(run_greenlet_target)
        greenlet_coroutine = gl.switch()

        while greenlet_coroutine is not None:
            task = asyncio.create_task(greenlet_coroutine)
            try:
                await task
            except:
                # this allows an exception to be raised within
                # the moderated greenlet so that it can continue
                # its expected flow.
                greenlet_coroutine = gl.throw(*sys.exc_info())
            else:
                greenlet_coroutine = gl.switch(task.result())

    current_greenlet = greenlet.greenlet(run_greenlet)
    current_greenlet_context.set(current_greenlet)
    try:
        await current_greenlet.switch()
    finally:
        current_greenlet_context.set(None)

    return result_future.result()


if __name__ == "__main__":

    def sleep_on_db(conn):
        await_(conn.fetchrow("select pg_sleep(10)"))
        result = await_(conn.fetchrow("select 1"))

        return result[0]

    async def run_request():
        conn = await (
            asyncpg.connect(
                user="scott",
                password="tiger",
                host="localhost",
                database="test",
            )
        )

        print("starting, kill the connection")
        # ps -ef | grep -e "postgres.*SELECT"  | grep -v "grep" | awk '{print $2}' | xargs kill -9

        retval = await greenlet_spawn(sleep_on_db, conn)
        print("ok: %s" % retval)

        await (conn.close())

    async def main():
        await run_request()

    asyncio.run(main())

stack trace looks good:

Traceback (most recent call last):
  File "test3.py", line 83, in <module>
    asyncio.run(main())
  File "/opt/python-3.8.3/lib/python3.8/asyncio/runners.py", line 43, in run
    return loop.run_until_complete(main)
  File "/opt/python-3.8.3/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete
    return future.result()
  File "test3.py", line 81, in main
    await run_request()
  File "test3.py", line 75, in run_request
    retval = await greenlet_spawn(sleep_on_db, conn)
  File "test3.py", line 49, in greenlet_spawn
    await current_greenlet.switch()
  File "test3.py", line 42, in run_greenlet
    greenlet_coroutine = gl.throw(*sys.exc_info())
  File "test3.py", line 27, in run_greenlet_target
    result_future.set_result(fn(*args))
  File "test3.py", line 59, in sleep_on_db
    await_(conn.fetchrow("select pg_sleep(10)"))
  File "test3.py", line 19, in await_
    return current_greenlet.switch(coroutine)
  File "test3.py", line 37, in run_greenlet
    await task
  File "/home/classic/.venv3/lib/python3.8/site-packages/asyncpg/connection.py", line 454, in fetchrow
    data = await self._execute(query, args, 1, timeout)
  File "/home/classic/.venv3/lib/python3.8/site-packages/asyncpg/connection.py", line 1402, in _execute
    result, _ = await self.__execute(
  File "/home/classic/.venv3/lib/python3.8/site-packages/asyncpg/connection.py", line 1411, in __execute
    return await self._do_execute(query, executor, timeout)
  File "/home/classic/.venv3/lib/python3.8/site-packages/asyncpg/connection.py", line 1433, in _do_execute
    result = await executor(stmt, None)
  File "asyncpg/protocol/protocol.pyx", line 196, in bind_execute
asyncpg.exceptions.ConnectionDoesNotExistError: connection was closed in the middle of operation