Async Thin: bounded cancellation while BREAK/RESET waits on an open silent socket

# /// script
# requires-python = ">=3.14"
# dependencies = ["oracledb==26.0.1"]
# ///
import asyncio
import getpass
import os
import platform
import sys
from time import monotonic

import oracledb

WAIT = 5


class Relay:
    def __init__(self, host, port):
        self.upstream = host, port
        self.forwarding = asyncio.Event()
        self.forwarding.set()
        self.intercepted = asyncio.Event()
        self.writers = set()
        self.handlers = set()

    async def start(self):
        self.server = await asyncio.start_server(self.accept, "127.0.0.1", 0)
        self.port = self.server.sockets[0].getsockname()[1]

    def accept(self, reader, writer):
        self.writers.add(writer)
        self.handlers.add(asyncio.create_task(self.serve(reader, writer)))

    async def pump(self, reader, writer):
        while data := await reader.read(65536):
            if not self.forwarding.is_set():
                self.intercepted.set()
                await self.forwarding.wait()
            writer.write(data)
            await writer.drain()

    async def serve(self, reader, writer):
        streams, pumps = [writer], []
        try:
            remote, remote_writer = await asyncio.wait_for(
                asyncio.open_connection(*self.upstream), WAIT
            )
            streams.append(remote_writer)
            self.writers.add(remote_writer)
            pumps = [asyncio.create_task(self.pump(reader, remote_writer)),
                     asyncio.create_task(self.pump(remote, writer))]
            await asyncio.wait(pumps, return_when=asyncio.FIRST_COMPLETED)
        finally:
            for stream in streams:
                stream.transport.abort()
                self.writers.discard(stream)
            for task in pumps:
                task.cancel()
            results = await asyncio.gather(*pumps, return_exceptions=True)
            for result in results:
                if isinstance(result, Exception) and not isinstance(result, ConnectionError):
                    raise result

    async def cut(self):
        self.forwarding.set()
        for writer in tuple(self.writers):
            writer.transport.abort()

    async def close(self):
        self.server.close()
        await asyncio.wait_for(self.server.wait_closed(), WAIT)
        await self.cut()
        if self.handlers:
            done, pending = await asyncio.wait(self.handlers, timeout=WAIT)
            assert not pending, "Relay handlers did not stop"
            for task in done:
                task.result()
            self.handlers.clear()
        assert not self.writers


async def settle(task, relay):
    done, _ = await asyncio.wait((task,), timeout=WAIT)
    if not done:
        # Closing only this relay wakes the driver's pending receive on failure.
        await relay.cut()
        done, _ = await asyncio.wait((task,), timeout=WAIT)
    assert done, "Task did not stop after relay sockets were cut"
    try:
        task.result()
    except asyncio.CancelledError:
        return "CancelledError"
    except Exception as exc:
        return getattr(exc.args[0], "full_code", type(exc).__name__) if exc.args else type(exc).__name__
    return "success"


async def case(config, timeout_ms, cancel_after):
    relay = Relay(config.pop("host"), config.pop("port"))
    await relay.start()
    connection = None
    task = timer = None
    try:
        connection = await asyncio.wait_for(oracledb.connect_async(
            **config, host="127.0.0.1", port=relay.port), WAIT)
        with connection.cursor() as cursor:
            await cursor.execute("select sys_context('USERENV','SESSION_USER'), "
                                 "sys_context('USERENV','CON_NAME') from dual")
            assert await cursor.fetchone() == (config["user"].upper(), config["service_name"].upper())
        print("connection.version:", connection.version)
        connection.call_timeout = timeout_ms
        relay.forwarding.clear()
        task = asyncio.create_task(connection.ping())
        await asyncio.wait_for(relay.intercepted.wait(), WAIT)
        started = monotonic()
        timer = asyncio.get_running_loop().call_later(cancel_after, task.cancel)
        done, _ = await asyncio.wait((task,), timeout=cancel_after + 0.25)
        print(f"call_timeout={timeout_ms}ms cancel_after={cancel_after}s "
              f"observed={monotonic() - started:.3f}s done_before_restore={bool(done)} "
              f"cancellation_requests={task.cancelling()}")
        relay.forwarding.set()
        timer.cancel()
        print("outcome_after_restore:", await settle(task, relay))
    finally:
        if timer:
            timer.cancel()
        relay.forwarding.set()
        try:
            if task:
                await settle(task, relay)
            if connection:
                print("close:", await settle(asyncio.create_task(connection.close()), relay))
        finally:
            await relay.close()
    recovery = await asyncio.wait_for(oracledb.connect_async(**config,
        host=relay.upstream[0], port=relay.upstream[1]), WAIT)
    recovery.call_timeout = 250
    try:
        await asyncio.wait_for(recovery.ping(), WAIT)
        print("fresh_connection_recovery: success; relay handlers=0; streams=0")
    finally:
        await asyncio.wait_for(recovery.close(), WAIT)


async def main():
    # ONLY use a dedicated disposable test account/database; never production.
    assert os.environ.get("ORACLE_TEST_ONLY") == "1", "Set ORACLE_TEST_ONLY=1 for a dedicated test environment"
    config = {"host": os.environ["ORACLE_HOST"], "port": int(os.environ.get("ORACLE_PORT", "1521")),
              "service_name": os.environ["ORACLE_SERVICE"], "user": os.environ["ORACLE_USER"],
              "password": os.environ.get("ORACLE_PASSWORD") or getpass.getpass("Test database password: ")}
    print("platform.platform:", platform.platform())
    print("sys.maxsize > 2**32:", sys.maxsize > 2**32)
    print("platform.python_version:", platform.python_version())
    print("oracledb.__version__:", oracledb.__version__, "thin:", oracledb.is_thin_mode())
    await case(config.copy(), 30_000, 1.0)
    await case(config.copy(), 250, 2.0)


if __name__ == "__main__":
    asyncio.run(main())

This sketch makes every cancellation destructive. Preserving the current default would require an explicit opt-in design; the diff alone does not provide that API.

--- a/src/oracledb/impl/thin/protocol.pyx
+++ b/src/oracledb/impl/thin/protocol.pyx
@@ -77,7 +77,7 @@
                 or self._read_buf._transport._transport is None:
             self._transport = None
 
-    cdef int _disconnect(self) except -1:
+    cdef int _disconnect(self, bint abort=False) except -1:
         """
         Disconnects from the transport.
         """
@@ -86,7 +86,7 @@
             self._transport = None
             self._read_buf._transport = None
             self._write_buf._transport = None
-            transport.disconnect()
+            transport.disconnect(abort)
 
     cdef bint _get_is_healthy(self):
         """
@@ -384,6 +384,13 @@
                                                              description)
 
     async def _process_message(self, Message message):
+        try:
+            await self._process_message_owned(message)
+        except asyncio.CancelledError:
+            self._disconnect(True)
+            raise
+
+    async def _process_message_owned(self, Message message):
         """
         Sends a message to the server and processes its response.
         """
@@ -393,6 +400,8 @@
         try:
             coroutine = self._process_message_helper(message)
             await asyncio.wait_for(coroutine, timeout_obj)
+        except asyncio.CancelledError:
+            raise
         except asyncio.TimeoutError:
             try:
                 coroutine = self._process_timeout_helper(message, timeout)
--- a/src/oracledb/impl/thin/transport.pyx
+++ b/src/oracledb/impl/thin/transport.pyx
@@ -113,7 +113,7 @@
             offset += 8
         self._print_output("\n".join(output_lines))
 
-    cdef int disconnect(self) except -1:
+    cdef int disconnect(self, bint abort=False) except -1:
         """
         Disconnects the transport.
         """
@@ -122,7 +122,10 @@
                 self._print_output(
                     self._get_debugging_header("Disconnecting transport")
                 )
-            self._transport.close()
+            if abort and self._is_async:
+                self._transport.abort()
+            else:
+                self._transport.close()
             self._transport = None
 
     cdef int create_ssl_context(self, ConnectParamsImpl params,