Merge pull request #13 from dapper91/dev

- socket connection manager aliveness check added
- shared lock bug fixed

Author: dapper91

CHANGELOG.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+0.4.0 (2023-08-15)
+------------------
+
+- socket connection manager aliveness check added
+- shared lock bug fixed
+
+
 0.3.0 (2023-08-10)
 ------------------
 

generic_connection_pool/asyncio/__init__.py

@@ -1,7 +1,13 @@
+"""
+Asyncio connection pool implementation.
+"""
+
 from .locks import SharedLock
-from .pool import BaseConnectionManager, ConnectionPool
+from .pool import BaseConnectionManager, ConnectionPool, ConnectionT, EndpointT
 
 __all__ = [
     'BaseConnectionManager',
     'ConnectionPool',
+    'ConnectionT',
+    'EndpointT',
 ]

generic_connection_pool/asyncio/locks.py

@@ -151,6 +151,6 @@ async def _acquire_exclusive(self, timeout: Optional[float] = None) -> bool:
         return True
 
     def _release_exclusive(self) -> None:
-        self._lock.release()
         self._exclusive = False
         self._exclusive_owner = None
+        self._lock.release()

generic_connection_pool/contrib/socket.py

@@ -3,19 +3,30 @@
 """
 
 import contextlib
+import errno
 import socket
 from ipaddress import IPv4Address, IPv6Address
 from ssl import SSLContext, SSLSocket
 from typing import Generator, Optional, Tuple, Union
 
-from generic_connection_pool.threading import BaseConnectionManager
+from generic_connection_pool.threading import BaseConnectionManager, EndpointT
 
 IpAddress = Union[IPv4Address, IPv6Address]
 Hostname = str
 Port = int
 TcpEndpoint = Tuple[IpAddress, Port]
 
 
+@contextlib.contextmanager
+def socket_nonblocking(sock: socket.socket) -> Generator[None, None, None]:
+    orig_timeout = sock.gettimeout()
+    sock.settimeout(0)
+    try:
+        yield
+    finally:
+        sock.settimeout(orig_timeout)
+
+
 @contextlib.contextmanager
 def socket_timeout(sock: socket.socket, timeout: Optional[float]) -> Generator[None, None, None]:
     if timeout is None:
@@ -65,6 +76,19 @@ def dispose(self, endpoint: TcpEndpoint, conn: socket.socket, timeout: Optional[
 
         conn.close()
 
+    def check_aliveness(self, endpoint: EndpointT, conn: socket.socket, timeout: Optional[float] = None) -> bool:
+        try:
+            with socket_nonblocking(conn):
+                if conn.recv(1, socket.MSG_PEEK) == b'':
+                    return False
+        except BlockingIOError as exc:
+            if exc.errno != errno.EAGAIN:
+                raise
+        except OSError:
+            return False
+
+        return True
+
 
 SslEndpoint = Tuple[Hostname, Port]
 
@@ -94,3 +118,17 @@ def dispose(self, endpoint: SslEndpoint, conn: SSLSocket, timeout: Optional[floa
             pass
 
         conn.close()
+
+    def check_aliveness(self, endpoint: EndpointT, conn: socket.socket, timeout: Optional[float] = None) -> bool:
+        try:
+            with socket_nonblocking(conn):
+                # peek into the plain socket since ssl socket doesn't support flags
+                if socket.socket.recv(conn, 1, socket.MSG_PEEK) == b'':
+                    return False
+        except BlockingIOError as exc:
+            if exc.errno != errno.EAGAIN:
+                raise
+        except OSError:
+            return False
+
+        return True

generic_connection_pool/contrib/socket_async.py

@@ -3,6 +3,7 @@
 """
 
 import asyncio
+import errno
 import socket
 from ipaddress import IPv4Address, IPv6Address
 from ssl import SSLContext
@@ -42,6 +43,18 @@ async def dispose(self, endpoint: TcpEndpoint, conn: socket.socket) -> None:
         conn.shutdown(socket.SHUT_RDWR)
         conn.close()
 
+    async def check_aliveness(self, endpoint: TcpEndpoint, conn: socket.socket) -> bool:
+        try:
+            if conn.recv(1, socket.MSG_PEEK) == b'':
+                return False
+        except BlockingIOError as exc:
+            if exc.errno != errno.EAGAIN:
+                raise
+        except OSError:
+            return False
+
+        return True
+
 
 Hostname = str
 TcpStreamEndpoint = Tuple[Hostname, Port]

generic_connection_pool/threading/__init__.py

@@ -1,6 +1,12 @@
-from .pool import BaseConnectionManager, ConnectionPool
+"""
+Threading connection pool implementation.
+"""
+
+from .pool import BaseConnectionManager, ConnectionPool, ConnectionT, EndpointT
 
 __all__ = [
     'BaseConnectionManager',
     'ConnectionPool',
+    'ConnectionT',
+    'EndpointT',
 ]

generic_connection_pool/threading/locks.py

@@ -131,9 +131,9 @@ def _acquire_exclusive(self, blocking: bool = True, timeout: Optional[float] = N
         return True
 
     def _release_exclusive(self) -> None:
-        self._lock.release()
         self._exclusive = False
         self._exclusive_owner = None
+        self._lock.release()
 
     def _is_owned(self) -> bool:  # to be compatible with threading.Condition
         return self._exclusive_owner == threading.get_ident()

pyproject.toml

@@ -1,13 +1,13 @@
 [tool.poetry]
 name = "generic-connection-pool"
-version = "0.3.0"
+version = "0.4.0"
 description = "generic connection pool"
 authors = ["Dmitry Pershin <[email protected]>"]
 license = "Unlicense"
 readme = "README.rst"
 homepage = "https://github.com/dapper91/generic-connection-pool"
 repository = "https://github.com/dapper91/generic-connection-pool"
-documentation = "https://github.com/dapper91/generic-connection-pool"
+documentation = "https://generic-connection-pool.readthedocs.io"
 keywords = ['pool', 'connection-pool', 'asyncio', 'socket', 'tcp']
 classifiers = [
     "Development Status :: 5 - Production/Stable",
@@ -44,6 +44,7 @@ pytest-asyncio = "^0.21.1"
 pytest-mock = "^3.11.1"
 pytest-cov = "^4.1.0"
 types-psycopg2 = "^2.9.21"
+pytest-timeout = "^2.1.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/contrib/test_async_socket_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import socket
 import ssl
 from ipaddress import IPv4Address
 from pathlib import Path
@@ -13,9 +14,10 @@
 class TCPServer:
     @staticmethod
     async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
-        data = await reader.read(1500)
-        writer.write(data)
-        await writer.drain()
+        while data := await reader.read(1024):
+            writer.write(data)
+            await writer.drain()
+
         writer.close()
         await writer.wait_closed()
 
@@ -31,6 +33,7 @@ async def start(self) -> None:
             host=self._hostname,
             port=self._port,
             ssl=self._ssl_ctx,
+            family=socket.AF_INET,
             reuse_port=True,
         )
         self._server_task = asyncio.create_task(server.serve_forever())
@@ -68,44 +71,73 @@ async def ssl_server(
     await server.stop()
 
 
+@pytest.mark.timeout(5.0)
 async def test_tcp_socket_manager(tcp_server: Tuple[IPv4Address, int]):
     loop = asyncio.get_running_loop()
     addr, port = tcp_server
 
     pool = ConnectionPool(TcpSocketConnectionManager())
-    async with pool.connection((addr, port)) as sock:
-        request = b'test'
-        await loop.sock_sendall(sock, request)
-        response = await loop.sock_recv(sock, len(request))
-        assert response == request
+
+    attempts = 3
+    request = b'test'
+    for _ in range(attempts):
+        async with pool.connection((addr, port)) as sock1:
+            await loop.sock_sendall(sock1, request)
+            response = await loop.sock_recv(sock1, len(request))
+            assert response == request
+
+            async with pool.connection((addr, port)) as sock2:
+                await loop.sock_sendall(sock2, request)
+                response = await loop.sock_recv(sock2, len(request))
+                assert response == request
 
     await pool.close()
 
 
+@pytest.mark.timeout(5.0)
 async def test_tcp_stream_manager(resource_dir: Path, tcp_server: Tuple[IPv4Address, int]):
     addr, port = tcp_server
 
     pool = ConnectionPool(TcpStreamConnectionManager(ssl=None))
-    async with pool.connection((str(addr), port)) as (reader, writer):
-        request = b'test'
-        writer.write(request)
-        await writer.drain()
-        response = await reader.read()
-        assert response == request
+
+    attempts = 3
+    request = b'test'
+    for _ in range(attempts):
+        async with pool.connection((str(addr), port)) as (reader1, writer1):
+            writer1.write(request)
+            await writer1.drain()
+            response = await reader1.read(len(request))
+            assert response == request
+
+            async with pool.connection((str(addr), port)) as (reader2, writer2):
+                writer2.write(request)
+                await writer2.drain()
+                response = await reader2.read(len(request))
+                assert response == request
 
     await pool.close()
 
 
+@pytest.mark.timeout(5.0)
 async def test_ssl_stream_manager(resource_dir: Path, ssl_server: Tuple[str, int]):
     hostname, port = ssl_server
     ssl_context = ssl.create_default_context(cafile=resource_dir / 'ssl.cert')
 
     pool = ConnectionPool(TcpStreamConnectionManager(ssl_context))
-    async with pool.connection((hostname, port)) as (reader, writer):
-        request = b'test'
-        writer.write(request)
-        await writer.drain()
-        response = await reader.read()
-        assert response == request
+
+    attempts = 3
+    request = b'test'
+    for _ in range(attempts):
+        async with pool.connection((hostname, port)) as (reader1, writer1):
+            writer1.write(request)
+            await writer1.drain()
+            response = await reader1.read(len(request))
+            assert response == request
+
+            async with pool.connection((hostname, port)) as (reader2, writer2):
+                writer2.write(request)
+                await writer2.drain()
+                response = await reader2.read(len(request))
+                assert response == request
 
     await pool.close()