Expose UpstreamConnectionPool to web & proxy plugins (#946)

* Expose conn pool to plugins

* Fix reusable state handling

* Separate `release` and `retain` methods

* Fix conn pool tests

* Fix tests

Author: abhinavsingh

proxy/core/connection/pool.py

@@ -77,57 +77,50 @@ def __init__(self) -> None:
         self.connections: Dict[int, TcpServerConnection] = {}
         self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
 
-    def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
-        """Creates and add a new connection to the pool."""
-        new_conn = TcpServerConnection(addr[0], addr[1])
-        new_conn.connect()
-        self._add(new_conn)
-        return new_conn
-
     def acquire(self, addr: Tuple[str, int]) -> Tuple[bool, TcpServerConnection]:
         """Returns a reusable connection from the pool.
 
         If none exists, will create and return a new connection."""
+        created, conn = False, None
         if addr in self.pools:
             for old_conn in self.pools[addr]:
                 if old_conn.is_reusable():
-                    old_conn.mark_inuse()
+                    conn = old_conn
                     logger.debug(
                         'Reusing connection#{2} for upstream {0}:{1}'.format(
                             addr[0], addr[1], id(old_conn),
                         ),
                     )
-                    return False, old_conn
-        new_conn = self.add(addr)
-        logger.debug(
-            'Created new connection#{2} for upstream {0}:{1}'.format(
-                addr[0], addr[1], id(new_conn),
-            ),
-        )
-        return True, new_conn
+                    break
+        if conn is None:
+            created, conn = True, self.add(addr)
+        conn.mark_inuse()
+        return created, conn
 
     def release(self, conn: TcpServerConnection) -> None:
         """Release a previously acquired connection.
 
-        If the connection has not been closed,
-        then it will be retained in the pool for reusability.
+        Releasing a connection will shutdown and close the socket
+        including internal pool cleanup.
         """
         assert not conn.is_reusable()
-        if conn.closed:
-            logger.debug(
-                'Removing connection#{2} from pool from upstream {0}:{1}'.format(
-                    conn.addr[0], conn.addr[1], id(conn),
-                ),
-            )
-            self._remove(conn.connection.fileno())
-        else:
-            logger.debug(
-                'Retaining connection#{2} to upstream {0}:{1}'.format(
-                    conn.addr[0], conn.addr[1], id(conn),
-                ),
-            )
-            # Reset for reusability
-            conn.reset()
+        logger.debug(
+            'Removing connection#{2} from pool from upstream {0}:{1}'.format(
+                conn.addr[0], conn.addr[1], id(conn),
+            ),
+        )
+        self._remove(conn.connection.fileno())
+
+    def retain(self, conn: TcpServerConnection) -> None:
+        """Retained previously acquired connection in the pool for reusability."""
+        assert not conn.closed
+        logger.debug(
+            'Retaining connection#{2} to upstream {0}:{1}'.format(
+                conn.addr[0], conn.addr[1], id(conn),
+            ),
+        )
+        # Reset for reusability
+        conn.reset()
 
     async def get_events(self) -> SelectableEvents:
         """Returns read event flag for all reusable connections in the pool."""
@@ -152,10 +145,28 @@ async def handle_events(self, readables: Readables, _writables: Writables) -> bo
             self._remove(fileno)
         return False
 
+    def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
+        """Creates, connects and adds a new connection to the pool.
+
+        Returns newly created connection.
+
+        NOTE: You must not use the returned connection, instead use `acquire`.
+        """
+        new_conn = TcpServerConnection(addr[0], addr[1])
+        new_conn.connect()
+        self._add(new_conn)
+        logger.debug(
+            'Created new connection#{2} for upstream {0}:{1}'.format(
+                addr[0], addr[1], id(new_conn),
+            ),
+        )
+        return new_conn
+
     def _add(self, conn: TcpServerConnection) -> None:
         """Adds a new connection to internal data structure."""
         if conn.addr not in self.pools:
             self.pools[conn.addr] = set()
+        conn._reusable = True
         self.pools[conn.addr].add(conn)
         self.connections[conn.connection.fileno()] = conn
 

proxy/http/plugin.py

@@ -52,7 +52,7 @@ def __init__(
             flags: argparse.Namespace,
             client: TcpClientConnection,
             request: HttpParser,
-            event_queue: Optional[EventQueue],
+            event_queue: Optional[EventQueue] = None,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ):
         self.uid: str = uid

proxy/http/proxy/plugin.py

@@ -11,14 +11,17 @@
 import argparse
 
 from abc import ABC
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 from ..parser import HttpParser
 from ..descriptors import DescriptorsHandlerMixin
 
 from ...core.event import EventQueue
 from ...core.connection import TcpClientConnection
 
+if TYPE_CHECKING:
+    from ...core.connection import UpstreamConnectionPool
+
 
 class HttpProxyBasePlugin(DescriptorsHandlerMixin, ABC):
     """Base HttpProxyPlugin Plugin class.
@@ -31,11 +34,13 @@ def __init__(
             flags: argparse.Namespace,
             client: TcpClientConnection,
             event_queue: EventQueue,
+            upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ) -> None:
         self.uid = uid                  # pragma: no cover
         self.flags = flags              # pragma: no cover
         self.client = client            # pragma: no cover
         self.event_queue = event_queue  # pragma: no cover
+        self.upstream_conn_pool = upstream_conn_pool
 
     def name(self) -> str:
         """A unique name for your plugin.

proxy/http/proxy/server.py

@@ -159,6 +159,7 @@ def __init__(
                     self.flags,
                     self.client,
                     self.event_queue,
+                    self.upstream_conn_pool,
                 )
                 self.plugins[instance.name()] = instance
 

proxy/http/server/plugin.py

@@ -11,7 +11,7 @@
 import argparse
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 
 from ..websocket import WebsocketFrame
 from ..parser import HttpParser
@@ -20,6 +20,9 @@
 from ...core.connection import TcpClientConnection
 from ...core.event import EventQueue
 
+if TYPE_CHECKING:
+    from ...core.connection import UpstreamConnectionPool
+
 
 class HttpWebServerBasePlugin(DescriptorsHandlerMixin, ABC):
     """Web Server Plugin for routing of requests."""
@@ -30,11 +33,13 @@ def __init__(
             flags: argparse.Namespace,
             client: TcpClientConnection,
             event_queue: EventQueue,
+            upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ):
         self.uid = uid
         self.flags = flags
         self.client = client
         self.event_queue = event_queue
+        self.upstream_conn_pool = upstream_conn_pool
 
     def name(self) -> str:
         """A unique name for your plugin.

proxy/http/server/web.py

@@ -106,6 +106,7 @@ def _initialize_web_plugins(self) -> None:
                 self.flags,
                 self.client,
                 self.event_queue,
+                self.upstream_conn_pool,
             )
             self.plugins[instance.name()] = instance
             for (protocol, route) in instance.routes():

tests/core/test_conn_pool.py

@@ -21,30 +21,34 @@
 class TestConnectionPool(unittest.TestCase):
 
     @mock.patch('proxy.core.connection.pool.TcpServerConnection')
-    def test_acquire_and_release_and_reacquire(self, mock_tcp_server_connection: mock.Mock) -> None:
+    def test_acquire_and_retain_and_reacquire(self, mock_tcp_server_connection: mock.Mock) -> None:
         pool = UpstreamConnectionPool()
         # Mock
         mock_conn = mock_tcp_server_connection.return_value
         addr = mock_conn.addr
         mock_conn.is_reusable.side_effect = [
-            False, True, True,
+            True,
         ]
         mock_conn.closed = False
         # Acquire
         created, conn = pool.acquire(addr)
-        self.assertTrue(created)
         mock_tcp_server_connection.assert_called_once_with(addr[0], addr[1])
+        mock_conn.mark_inuse.assert_called_once()
+        mock_conn.reset.assert_not_called()
+        self.assertTrue(created)
         self.assertEqual(conn, mock_conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])
-        # Release (connection must be retained because not closed)
-        pool.release(conn)
+        self.assertEqual(len(pool.connections), 1)
+        self.assertEqual(pool.connections[conn.connection.fileno()], mock_conn)
+        # Retail
+        pool.retain(conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])
+        mock_conn.reset.assert_called_once()
         # Reacquire
         created, conn = pool.acquire(addr)
         self.assertFalse(created)
-        mock_conn.reset.assert_called_once()
         self.assertEqual(conn, mock_conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])