✨ feat(aio): implement non-blocking socket stream and async connection lifecycle

Author: attid

firebirdsql/aio/fbcore.py

@@ -66,16 +66,21 @@ def __init__(self, trans):
         DEBUG_OUTPUT("AsyncStatement::__init__()")
         self.trans = trans
 
+        self._is_open = False
+        self.stmt_type = None
+        self.handle = -1
+
+    @classmethod
+    async def create(cls, trans):
+        self = cls(trans)
         self.trans.connection._op_allocate_statement()
         if (self.trans.connection.accept_type & ptype_MASK) == ptype_lazy_send:
             self.trans.connection.lazy_response_count += 1
             self.handle = -1
         else:
-            (h, oid, buf) = self.trans.connection._op_response()
+            (h, oid, buf) = await self.trans.connection._async_op_response()
             self.handle = h
-
-        self._is_open = False
-        self.stmt_type = None
+        return self
 
     async def fetch_generator(self, rows, more_data):
         DEBUG_OUTPUT("AsyncStatement::_fetch_generator()", self.handle, self.trans._trans_handle, self.trans.connection.db_handle)
@@ -172,7 +177,7 @@ class AsyncPreparedStatement(PreparedStatement):
     async def __init__(self, cur, sql, explain_plan=False):
         DEBUG_OUTPUT("AsyncPreparedStatement::__init__()")
         await cur.transaction.check_trans_handle()
-        self.stmt = await AsyncStatement(cur.transaction)
+        self.stmt = await AsyncStatement.create(cur.transaction)
         await self.stmt.prepare(sql, explain_plan)
         self.sql = sql
 
@@ -217,7 +222,7 @@ async def _get_stmt(self, query):
                 await self.stmt.drop()
                 self.stmt = None
             if self.stmt is None:
-                self.stmt = AsyncStatement(self.transaction)
+                self.stmt = await AsyncStatement.create(self.transaction)
             stmt = self.stmt
             await stmt.prepare(query)
         return stmt
@@ -508,6 +513,17 @@ async def check_trans_handle(self):
         if self._trans_handle is None:
             await self._begin()
 
+    async def close(self):
+        if self._trans_handle is None:
+            return
+        if not self.is_dirty:
+            return
+        DEBUG_OUTPUT("AsyncTransaction::close()", self._trans_handle, self.connection.db_handle)
+        self.connection._op_rollback(self._trans_handle)
+        (h, oid, buf) = await self.connection._async_op_response()
+        self._trans_handle = None
+        self.is_dirty = False
+
 
 class AsyncConnectionResponseMixin(ConnectionResponseMixin):
     async def _async_recv_channel(self, nbytes, word_alignment=False):
@@ -516,9 +532,10 @@ async def _async_recv_channel(self, nbytes, word_alignment=False):
             n += 4 - nbytes % 4  # 4 bytes word alignment
         r = bytes([])
         while n:
-            if (self.timeout is not None and select.select([self.sock._sock], [], [], self.timeout)[0] == []):
-                break
-            b = await self.sock.async_recv(n)
+            if self.timeout is not None:
+                b = await asyncio.wait_for(self.sock.async_recv(n), timeout=self.timeout)
+            else:
+                b = await self.sock.async_recv(n)
             if not b:
                 break
             r += b
@@ -713,7 +730,7 @@ async def _async_parse_connect_response(self):
                     raise OperationalError(
                         'Unknown wirecrypt plugin %s' % (enc_plugin.encode("utf-8"))
                     )
-                (h, oid, buf) = self._op_response()
+                (h, oid, buf) = await self._async_op_response()
             else:
                 # no matched wire encription plugin
                 # self.auth_data use _op_attach() and _op_create()
@@ -983,6 +1000,25 @@ async def drop_database(self):
         self.sock = None
         self.db_handle = None
 
+    async def close(self):
+        DEBUG_OUTPUT("AsyncConnection::close()", id(self), self.db_handle)
+        if self.sock is None:
+            return
+        if self.db_handle is not None:
+            # cleanup transaction
+            for trans in list(self._cursors.keys()):
+                await trans.close()
+            if self.is_services:
+                self._op_service_detach()
+            else:
+                self._op_detach()
+            (h, oid, buf) = await self._async_op_response()
+        self.sock.close()
+        self.sock = None
+        self.db_handle = None
+
     def __del__(self):
         if self.sock:
-            self.close()
+            # Async close cannot be called from __del__
+            # self.close() 
+            pass

firebirdsql/aio/stream.py

@@ -25,16 +25,53 @@
 #
 # Python DB-API 2.0 module for Firebird.
 ##############################################################################
+import asyncio
+
 from firebirdsql.stream import SocketStream
+from firebirdsql.utils import bytes_to_bint
 
 
 class AsyncSocketStream(SocketStream):
     def __init__(self, host, port, loop, timeout, cloexec):
         super().__init__(host, port, timeout, cloexec)
         self.loop = loop
+        self._send_lock = asyncio.Lock()
+        self._last_send_task = None
+        self._sock.setblocking(False)
+        self._buf = b''
+
+    async def _await_pending_send(self):
+        task = self._last_send_task
+        if task is not None:
+            await task
 
     async def async_recv(self, nbytes):
-        b = await self.loop.sock_recv(self._sock, nbytes)
-        if self.read_translator:
-            b = self.read_translator.decrypt(b)
-        return b
+        await self._await_pending_send()
+        
+        if len(self._buf) < nbytes:
+            read_size = max(8192, nbytes - len(self._buf))
+            chunk = await self.loop.sock_recv(self._sock, read_size)
+            if self.read_translator:
+                chunk = self.read_translator.decrypt(chunk)
+            self._buf += chunk
+
+        ret = self._buf[:nbytes]
+        self._buf = self._buf[nbytes:]
+        return ret
+
+    def send(self, b):
+        if not self.loop.is_running():
+            return super().send(b)
+        if self.write_translator:
+            b = self.write_translator.encrypt(b)
+
+        previous_task = self._last_send_task
+
+        async def _send_all(payload):
+            if previous_task is not None:
+                await previous_task
+            async with self._send_lock:
+                await self.loop.sock_sendall(self._sock, payload)
+
+        self._last_send_task = self.loop.create_task(_send_all(b))
+        return None