Merge pull request #18 from dapper91/dev

- unix socket managers implemented.
- hooks usage example added.

Author: dapper91

CHANGELOG.rst

@@ -1,6 +1,12 @@
 Changelog
 =========
 
+0.5.0 (2023-08-17)
+------------------
+
+- unix socket managers implemented.
+
+
 0.4.1 (2023-08-16)
 ------------------
 

docs/source/pages/api.rst

@@ -35,3 +35,6 @@ _______
 
 .. automodule:: generic_connection_pool.contrib.unix
     :members:
+
+.. automodule:: generic_connection_pool.contrib.unix_async
+    :members:

docs/source/pages/quickstart.rst

@@ -125,6 +125,14 @@ The following example illustrate how to implement a custom connection manager fo
   :language: python
 
 
+Connection manager allows to define methods to be called on connection acquire, release or when
+a connection determined to be dead. That helps to log pool actions or collect metrics.
+The following examples illustrate how to collect pool metrics and export them to prometheus.
+
+.. literalinclude:: ../../../examples/manager_hooks.py
+  :language: python
+
+
 Examples
 ________
 

examples/manager_hooks.py

@@ -0,0 +1,46 @@
+import ssl
+import time
+from ssl import SSLSocket
+from typing import Any, Dict, Tuple
+
+import prometheus_client as prom
+
+from generic_connection_pool.contrib.socket import SslSocketConnectionManager
+from generic_connection_pool.threading import ConnectionPool
+
+Hostname = str
+Port = int
+Endpoint = Tuple[Hostname, Port]
+
+acquire_latency_hist = prom.Histogram('acquire_latency', 'Connections acquire latency', labelnames=['hostname'])
+acquire_total = prom.Counter('acquire_total', 'Connections acquire count', labelnames=['hostname'])
+dead_conn_total = prom.Counter('dead_conn_total', 'Dead connections count', labelnames=['hostname'])
+
+
+class ObservableConnectionManager(SslSocketConnectionManager):
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._acquires: Dict[SSLSocket, float] = {}
+
+    def on_acquire(self, endpoint: Endpoint, conn: SSLSocket) -> None:
+        hostname, port = endpoint
+
+        acquire_total.labels(hostname).inc()
+        self._acquires[conn] = time.time()
+
+    def on_release(self, endpoint: Endpoint, conn: SSLSocket) -> None:
+        hostname, port = endpoint
+
+        acquired_at = self._acquires.pop(conn)
+        acquire_latency_hist.labels(hostname).observe(time.time() - acquired_at)
+
+    def on_connection_dead(self, endpoint: Endpoint, conn: SSLSocket) -> None:
+        hostname, port = endpoint
+
+        dead_conn_total.labels(hostname).inc()
+
+
+http_pool = ConnectionPool[Endpoint, SSLSocket](
+    ObservableConnectionManager(ssl.create_default_context()),
+)

examples/ssl_pool.py

@@ -1,4 +1,3 @@
-import socket
 import ssl
 import urllib.parse
 from http.client import HTTPResponse
@@ -10,7 +9,7 @@
 Hostname = str
 Port = int
 Endpoint = Tuple[Hostname, Port]
-Connection = socket.socket
+Connection = ssl.SSLSocket
 
 
 http_pool = ConnectionPool[Endpoint, Connection](

examples/tcp_pool.py

@@ -1,12 +1,13 @@
 import socket
-from ipaddress import IPv4Address
-from typing import Tuple
+from ipaddress import IPv4Address, IPv6Address
+from typing import Tuple, Union
 
 from generic_connection_pool.contrib.socket import TcpSocketConnectionManager
 from generic_connection_pool.threading import ConnectionPool
 
 Port = int
-Endpoint = Tuple[IPv4Address, Port]
+IpAddress = Union[IPv4Address, IPv6Address]
+Endpoint = Tuple[IpAddress, Port]
 Connection = socket.socket
 
 
@@ -21,7 +22,7 @@
 )
 
 
-def command(addr: IPv4Address, port: int, cmd: str) -> None:
+def command(addr: IpAddress, port: int, cmd: str) -> None:
     with redis_pool.connection(endpoint=(addr, port), timeout=5.0) as sock:
         sock.sendall(cmd.encode() + b'\n')
         response = sock.recv(1024)

generic_connection_pool/contrib/socket_async.py

@@ -61,7 +61,11 @@ async def create(self, endpoint: TcpEndpoint) -> socket.socket:
         return sock
 
     async def dispose(self, endpoint: TcpEndpoint, conn: socket.socket) -> None:
-        conn.shutdown(socket.SHUT_RDWR)
+        try:
+            conn.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+
         conn.close()
 
 

generic_connection_pool/contrib/unix.py

@@ -3,24 +3,29 @@
 """
 
 import errno
+import pathlib
 import socket
 import sys
-from typing import Optional
+from typing import Generic, Optional
 
-from generic_connection_pool.contrib.socket import TcpEndpoint
+from generic_connection_pool.contrib.socket import BaseConnectionManager
+from generic_connection_pool.threading import EndpointT
 
 from .socket import socket_timeout
 
-if sys.platform not in ('linux', 'darwin'):
-    raise AssertionError('this module is supported only on linux and darwin platform')
+if sys.platform not in ('linux', 'darwin', 'freebsd'):
+    raise AssertionError('this module is supported only on unix platforms')
 
 
-class CheckSocketAlivenessMixin:
+UnixSocketEndpoint = pathlib.Path
+
+
+class CheckSocketAlivenessMixin(Generic[EndpointT]):
     """
     Socket aliveness checking mixin.
     """
 
-    def check_aliveness(self, endpoint: TcpEndpoint, conn: socket.socket, timeout: Optional[float] = None) -> bool:
+    def check_aliveness(self, endpoint: EndpointT, conn: socket.socket, timeout: Optional[float] = None) -> bool:
         try:
             with socket_timeout(conn, timeout):
                 resp = conn.recv(1, socket.MSG_PEEK | socket.MSG_DONTWAIT)
@@ -33,3 +38,28 @@ def check_aliveness(self, endpoint: TcpEndpoint, conn: socket.socket, timeout: O
             return False
 
         return True
+
+
+class UnixSocketConnectionManager(
+    CheckSocketAlivenessMixin[UnixSocketEndpoint],
+    BaseConnectionManager[UnixSocketEndpoint, socket.socket],
+):
+    """
+    Unix socket connection manager.
+    """
+
+    def create(self, endpoint: UnixSocketEndpoint, timeout: Optional[float] = None) -> socket.socket:
+        sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM)
+
+        with socket_timeout(sock, timeout):
+            sock.connect(str(endpoint))
+
+        return sock
+
+    def dispose(self, endpoint: UnixSocketEndpoint, conn: socket.socket, timeout: Optional[float] = None) -> None:
+        try:
+            conn.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+
+        conn.close()

generic_connection_pool/contrib/unix_async.py

@@ -0,0 +1,66 @@
+"""
+Asynchronous unix specific functionality.
+"""
+
+import asyncio
+import pathlib
+import socket
+import sys
+
+from generic_connection_pool.contrib.socket_async import BaseConnectionManager, SocketAlivenessCheckingMixin, Stream
+from generic_connection_pool.contrib.socket_async import StreamAlivenessCheckingMixin
+
+if sys.platform not in ('linux', 'darwin', 'freebsd'):
+    raise AssertionError('this module is supported only on unix platforms')
+
+
+UnixSocketEndpoint = pathlib.Path
+
+
+class UnixSocketConnectionManager(
+    SocketAlivenessCheckingMixin[UnixSocketEndpoint],
+    BaseConnectionManager[UnixSocketEndpoint, socket.socket],
+):
+    """
+    Asynchronous unix socket connection manager.
+    """
+
+    async def create(self, endpoint: UnixSocketEndpoint) -> socket.socket:
+        loop = asyncio.get_running_loop()
+
+        sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM)
+        sock.setblocking(False)
+
+        await loop.sock_connect(sock, address=(str(endpoint)))
+
+        return sock
+
+    async def dispose(self, endpoint: UnixSocketEndpoint, conn: socket.socket) -> None:
+        try:
+            conn.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+
+        conn.close()
+
+
+class UnixSocketStreamConnectionManager(
+    StreamAlivenessCheckingMixin[UnixSocketEndpoint],
+    BaseConnectionManager[UnixSocketEndpoint, Stream],
+):
+    """
+    Asynchronous unix socket stream connection manager.
+    """
+
+    async def create(self, endpoint: UnixSocketEndpoint) -> Stream:
+        reader, writer = await asyncio.open_unix_connection(path=endpoint)
+
+        return reader, writer
+
+    async def dispose(self, endpoint: UnixSocketEndpoint, conn: Stream) -> None:
+        reader, writer = conn
+        if writer.can_write_eof():
+            writer.write_eof()
+
+        writer.close()
+        await writer.wait_closed()

pyproject.toml

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