Merge pull request #5 from dapper91/dev

- socket_timeout contex manager works with ssl sockets.
- graceful_timeout default set to 0.
- connection release bug fixed.

Author: dapper91

CHANGELOG.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+0.2.0 (2023-04-19)
+------------------
+
+- socket_timeout contex manager works with ssl sockets.
+- graceful_timeout default set to 0.
+- connection release bug fixed.
+
 
 0.1.1 (2023-03-17)
 ------------------

generic_connection_pool/asyncio.py

@@ -155,17 +155,14 @@ async def release(self, conn: ConnectionT, endpoint: EndpointT) -> None:
             dispose_batch_size = self._dispose_batch_size or int(math.log2(self._pool_size + 1)) + 1
             await self._collect_disposable_connections(dispose_batch_size)
 
-    async def close(self, graceful_timeout: Optional[float] = None, timeout: Optional[float] = None) -> None:
+    async def close(self, graceful_timeout: float = 0.0, timeout: Optional[float] = None) -> None:
         """
         Closes the connection pool.
 
-        :param graceful_timeout: timeout within which the pool waits all acquired connection to be released
+        :param graceful_timeout: timeout within which the pool waits for all acquired connection to be released
         :param timeout: timeout after which the pool closes all connection despite they are released or not
         """
 
-        if graceful_timeout is None:
-            graceful_timeout = timeout
-
         if graceful_timeout is not None and timeout is not None:
             assert timeout >= graceful_timeout, "timeout can't be less than graceful_timeout"
 
@@ -403,7 +400,7 @@ async def _close_connections(
                 try:
                     while released:
                         conn_info = released[-1]
-                        await self._dispose_connection(conn_info, timeout=graceful_timer.remains)
+                        await self._dispose_connection(conn_info, timeout=global_timer.remains)
                         released.pop()
                 except asyncio.TimeoutError:
                     await asyncio.shield(self._return_released_conns(released))
@@ -422,5 +419,5 @@ async def _return_released_conns(self, released: List[ConnectionInfo[EndpointT,
             for conn_info in released:
                 pool = self._pools[conn_info.endpoint]
                 pool.queue[conn_info.conn] = conn_info
-                pool.access_queue.remove((conn_info.accessed_at, conn_info.conn))
+                pool.access_queue.push((conn_info.accessed_at, conn_info.conn))
                 self._pool_size += 1

generic_connection_pool/common.py

@@ -42,7 +42,7 @@ def remains(self) -> Optional[float]:
         if self._timeout is None:
             return None
 
-        return self._timeout - self.elapsed
+        return max(0.0, self._timeout - self.elapsed)
 
     @property
     def timedout(self) -> bool:

generic_connection_pool/contrib/socket.py

@@ -1,7 +1,8 @@
+import contextlib
 import socket
 from ipaddress import IPv4Address, IPv6Address
 from ssl import SSLContext, SSLSocket
-from typing import Optional, Tuple, Union
+from typing import Generator, Optional, Tuple, Union
 
 from generic_connection_pool.threding import BaseConnectionManager
 
@@ -11,6 +12,25 @@
 TcpEndpoint = Tuple[IpAddress, Port]
 
 
+@contextlib.contextmanager
+def socket_timeout(sock: socket.socket, timeout: Optional[float]) -> Generator[None, None, None]:
+    if timeout is None:
+        yield
+        return
+
+    orig_timeout = sock.gettimeout()
+    sock.settimeout(max(timeout, 1e-6))  # if timeout is 0 set it a small value to prevent non-blocking socket mode
+    try:
+        yield
+    except OSError as e:
+        if 'timed out' in str(e):
+            raise TimeoutError
+        else:
+            raise
+    finally:
+        sock.settimeout(orig_timeout)
+
+
 class TcpSocketConnectionManager(BaseConnectionManager[TcpEndpoint, socket.socket]):
     """
     TCP socket connection manager.
@@ -28,17 +48,12 @@ def create(self, endpoint: TcpEndpoint, timeout: Optional[float] = None) -> sock
 
         sock = socket.socket(family=family, type=socket.SOCK_STREAM)
 
-        orig_timeout = sock.gettimeout()
-        sock.settimeout(timeout)
-        try:
+        with socket_timeout(sock, timeout):
             sock.connect((str(addr), port))
-        finally:
-            sock.settimeout(orig_timeout)
 
         return sock
 
     def dispose(self, endpoint: TcpEndpoint, conn: socket.socket, timeout: Optional[float] = None) -> None:
-        conn.settimeout(timeout)
         try:
             conn.shutdown(socket.SHUT_RDWR)
         except OSError:
@@ -62,17 +77,13 @@ def create(self, endpoint: SslEndpoint, timeout: Optional[float] = None) -> SSLS
         hostname, port = endpoint
 
         sock = self._ssl.wrap_socket(socket.socket(type=socket.SOCK_STREAM), server_hostname=hostname)
-        orig_timeout = sock.gettimeout()
-        sock.settimeout(timeout)
-        try:
+
+        with socket_timeout(sock, timeout):
             sock.connect((hostname, port))
-        finally:
-            sock.settimeout(orig_timeout)
 
         return sock
 
     def dispose(self, endpoint: SslEndpoint, conn: SSLSocket, timeout: Optional[float] = None) -> None:
-        conn.settimeout(timeout)
         try:
             conn.shutdown(socket.SHUT_RDWR)
         except OSError:

generic_connection_pool/contrib/socket_async.py

@@ -49,7 +49,7 @@ class TcpStreamConnectionManager(BaseConnectionManager[TcpStreamEndpoint, TcpStr
     TCP stream connection manager.
     """
 
-    def __init__(self, ssl: Union[None, bool, SSLContext]):
+    def __init__(self, ssl: Union[None, bool, SSLContext] = None):
         self._ssl = ssl
 
     async def create(self, endpoint: TcpStreamEndpoint) -> TcpStream:

generic_connection_pool/threding.py

@@ -159,17 +159,14 @@ def release(self, conn: ConnectionT, endpoint: EndpointT) -> None:
             dispose_batch_size = self._dispose_batch_size or int(math.log2(self._pool_size + 1)) + 1
             self._collect_disposable_connections(dispose_batch_size)
 
-    def close(self, graceful_timeout: Optional[float] = None, timeout: Optional[float] = None) -> None:
+    def close(self, graceful_timeout: float = 0.0, timeout: Optional[float] = None) -> None:
         """
         Closes the connection pool.
 
-        :param graceful_timeout: timeout within which the pool waits all acquired connection to be released
+        :param graceful_timeout: timeout within which the pool waits for all acquired connection to be released
         :param timeout: timeout after which the pool closes all connection despite they are released or not
         """
 
-        if graceful_timeout is None:
-            graceful_timeout = timeout
-
         if graceful_timeout is not None and timeout is not None:
             assert timeout >= graceful_timeout, "timeout can't be less than graceful_timeout"
 
@@ -401,7 +398,7 @@ def _close_connections(self, graceful_timeout: Optional[float], timeout: Optiona
                 try:
                     while released:
                         conn_info = released[-1]
-                        self._dispose_connection(conn_info, timeout=graceful_timer.remains)
+                        self._dispose_connection(conn_info, timeout=global_timer.remains)
                         released.pop()
                 except TimeoutError:
                     self._return_released_conns(released)
@@ -420,5 +417,5 @@ def _return_released_conns(self, released: List[ConnectionInfo[EndpointT, Connec
             for conn_info in released:
                 pool = self._pools[conn_info.endpoint]
                 pool.queue[conn_info.conn] = conn_info
-                pool.access_queue.remove((conn_info.accessed_at, conn_info.conn))
+                pool.access_queue.push((conn_info.accessed_at, conn_info.conn))
                 self._pool_size += 1

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "generic-connection-pool"
-version = "0.1.1"
+version = "0.2.0"
 description = "generic connection pool"
 authors = ["Dmitry Pershin <[email protected]>"]
 license = "Unlicense"

tests/conftest.py

@@ -1,24 +1,34 @@
 import random
 from pathlib import Path
+from typing import Generator
 
 import pytest
 
 
-@pytest.fixture
+@pytest.fixture(scope='session')
 def test_dir() -> Path:
     return Path(__file__).parent
 
 
-@pytest.fixture
+@pytest.fixture(scope='session')
 def resource_dir(test_dir) -> Path:
     return test_dir / 'resources'
 
 
 @pytest.fixture(autouse=True)
-def init_random():
+def init_random() -> None:
     random.seed(0)
 
 
 @pytest.fixture(scope='session')
-def sleep_delay(pytestconfig) -> float:
-    return pytestconfig.getoption('--sleep-delay', 0.05)
+def delay(pytestconfig) -> float:
+    return pytestconfig.getoption('--delay', 0.05)
+
+
+@pytest.fixture(scope='session')
+def port_gen() -> Generator[int, None, None]:
+    def gen():
+        for port in range(10000, 65535):
+            yield port
+
+    return gen()

tests/contrib/test_async_socket_manager.py

@@ -2,48 +2,72 @@
 import ssl
 from ipaddress import IPv4Address
 from pathlib import Path
-from typing import Tuple
+from typing import AsyncGenerator, Generator, Optional, Tuple
 
 import pytest
-import pytest_asyncio.plugin
 
 from generic_connection_pool.asyncio import ConnectionPool
 from generic_connection_pool.contrib.socket_async import TcpSocketConnectionManager, TcpStreamConnectionManager
 
 
-@pytest.fixture
-async def tcp_server(request: pytest_asyncio.plugin.SubRequest, resource_dir: Path):
-    params = getattr(request, 'param', {})
-    user_ssl = params.get('use_ssl', False)
-
-    if user_ssl:
-        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
-        context.load_cert_chain(resource_dir / 'ssl.cert', resource_dir / 'ssl.key')
-    else:
-        context = None
-
-    hostname, addr, port = 'localhost', IPv4Address('127.0.0.1'),  10000
-
-    async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
-        data = await reader.read(1024)
+class TCPServer:
+    @staticmethod
+    async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
+        data = await reader.read(1500)
         writer.write(data)
         await writer.drain()
         writer.close()
         await writer.wait_closed()
 
-    async with (server := await asyncio.start_server(echo, host=hostname, port=port, ssl=context, reuse_port=True)):
-        server_task = asyncio.create_task(server.serve_forever())
-        yield hostname, addr, port
-        server_task.cancel()
-        try:
-            await server_task
-        except asyncio.CancelledError:
-            pass
+    def __init__(self, hostname: str, port: int, ssl_ctx: Optional[ssl.SSLContext] = None):
+        self._hostname = hostname
+        self._port = port
+        self._ssl_ctx = ssl_ctx
+        self._server_task: Optional[asyncio.Task[None]] = None
+
+    async def start(self) -> None:
+        server = await asyncio.start_server(
+            self.echo_handler,
+            host=self._hostname,
+            port=self._port,
+            ssl=self._ssl_ctx,
+            reuse_port=True,
+        )
+        self._server_task = asyncio.create_task(server.serve_forever())
+
+    async def stop(self) -> None:
+        if (server_task := self._server_task) is not None:
+            server_task.cancel()
+            try:
+                await server_task
+            except asyncio.CancelledError:
+                pass
+
+
+@pytest.fixture
+async def tcp_server(port_gen: Generator[int, None, None]) -> AsyncGenerator[Tuple[IPv4Address, int], None]:
+    addr, port = IPv4Address('127.0.0.1'), next(port_gen)
+    server = TCPServer(str(addr), port)
+    await server.start()
+    yield addr, port
+    await server.stop()
+
+
+@pytest.fixture
+async def ssl_server(resource_dir: Path, port_gen: Generator[int, None, None]) -> AsyncGenerator[Tuple[str, int], None]:
+    hostname, port = 'localhost', next(port_gen)
+    ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+    ssl_ctx.load_cert_chain(resource_dir / 'ssl.cert', resource_dir / 'ssl.key')
+
+    server = TCPServer(hostname, port, ssl_ctx=ssl_ctx)
+    await server.start()
+    yield hostname, port
+    await server.stop()
 
 
-async def test_tcp_socket_manager(tcp_server: Tuple[str, IPv4Address, int]):
+async def test_tcp_socket_manager(tcp_server: Tuple[IPv4Address, int]):
     loop = asyncio.get_running_loop()
-    hostname, addr, port = tcp_server
+    addr, port = tcp_server
 
     pool = ConnectionPool(TcpSocketConnectionManager())
     async with pool.connection((addr, port)) as sock:
@@ -55,20 +79,23 @@ async def test_tcp_socket_manager(tcp_server: Tuple[str, IPv4Address, int]):
     await pool.close()
 
 
-@pytest.mark.parametrize(
-    'use_ssl, tcp_server', [
-        (True, {'use_ssl': True}),
-        (False, {'use_ssl': False}),
-    ],
-    indirect=['tcp_server'],
-)
-async def test_ssl_stream_manager(resource_dir: Path, use_ssl: bool, tcp_server: Tuple[str, IPv4Address, int]):
-    hostname, addr, port = tcp_server
-
-    if use_ssl:
-        ssl_context = ssl.create_default_context(cafile=resource_dir / 'ssl.cert')
-    else:
-        ssl_context = None
+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
+
+    await pool.close()
+
+
+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):