Repository: abhinavsingh/proxy.py
Pull request #946 (2022-01-07 21:58:44)

Expose `UpstreamConnectionPool` to web & proxy plugins

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
> Merging [#946](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (ad77664) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/a2e1fc68a09416a6418694e31fd0ef9362114117?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (a2e1fc6) will **decrease** coverage by `0.32%`.
> The diff coverage is `50.00%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)

```diff
@@             Coverage Diff             @@
##           develop     #946      +/-   ##
===========================================
- Coverage    86.28%   85.95%   -0.33%     
===========================================
  Files          144      144              
  Lines         6371     6374       +3     
  Branches       642      642              
===========================================
- Hits          5497     5479      -18     
- Misses         755      767      +12     
- Partials       119      128       +9     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `85.95% <50.00%> (-0.19%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/http/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9wbHVnaW4ucHk=) | `93.54% <ø> (ø)` | |
| [proxy/http/server/web.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9zZXJ2ZXIvd2ViLnB5) | `59.47% <ø> (ø)` | |
| [proxy/http/server/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9zZXJ2ZXIvcGx1Z2luLnB5) | `86.20% <50.00%> (-6.11%)` | :arrow_down: |
| [tests/http/test\_proxy\_protocol.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC90ZXN0X3Byb3h5X3Byb3RvY29sLnB5) | `96.07% <0.00%> (-3.93%)` | :arrow_down: |
| [tests/test\_set\_open\_file\_limit.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdF9zZXRfb3Blbl9maWxlX2xpbWl0LnB5) | `96.42% <0.00%> (-3.58%)` | :arrow_down: |
| [tests/testing/test\_embed.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdGluZy90ZXN0X2VtYmVkLnB5) | `90.90% <0.00%> (-3.04%)` | :arrow_down: |
| [proxy/common/utils.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29tbW9uL3V0aWxzLnB5) | `91.59% <0.00%> (-2.53%)` | :arrow_down: |
| [proxy/http/handler.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9oYW5kbGVyLnB5) | `71.17% <0.00%> (-2.19%)` | :arrow_down: |
| [proxy/testing/test\_case.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvdGVzdGluZy90ZXN0X2Nhc2UucHk=) | `97.82% <0.00%> (-2.18%)` | :arrow_down: |
| [proxy/common/flag.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29tbW9uL2ZsYWcucHk=) | `91.17% <0.00%> (-1.48%)` | :arrow_down: |
| ... and [3 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [a2e1fc6...ad77664](https://codecov.io/gh/abhinavsingh/proxy.py/pull/946?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).


Patch:
diff --git a/proxy/core/connection/pool.py b/proxy/core/connection/pool.py
index 2f45e8a..3fed38d 100644
--- a/proxy/core/connection/pool.py
+++ b/proxy/core/connection/pool.py
@@ -77,57 +77,50 @@ class UpstreamConnectionPool(Work[TcpServerConnection]):
         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 @@ class UpstreamConnectionPool(Work[TcpServerConnection]):
             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
 
diff --git a/proxy/http/plugin.py b/proxy/http/plugin.py
index 931b6f4..eb9c070 100644
--- a/proxy/http/plugin.py
+++ b/proxy/http/plugin.py
@@ -52,7 +52,7 @@ class HttpProtocolHandlerPlugin(DescriptorsHandlerMixin, ABC):
             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
diff --git a/proxy/http/proxy/plugin.py b/proxy/http/proxy/plugin.py
index a5c62c3..8e769a5 100644
--- a/proxy/http/proxy/plugin.py
+++ b/proxy/http/proxy/plugin.py
@@ -11,7 +11,7 @@
 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
@@ -19,6 +19,9 @@ 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 @@ class HttpProxyBasePlugin(DescriptorsHandlerMixin, ABC):
             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.
diff --git a/proxy/http/proxy/server.py b/proxy/http/proxy/server.py
index 0d94ee4..ba87bb2 100644
--- a/proxy/http/proxy/server.py
+++ b/proxy/http/proxy/server.py
@@ -159,6 +159,7 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
                     self.flags,
                     self.client,
                     self.event_queue,
+                    self.upstream_conn_pool,
                 )
                 self.plugins[instance.name()] = instance
 
diff --git a/proxy/http/server/plugin.py b/proxy/http/server/plugin.py
index c9c47ce..5ea507e 100644
--- a/proxy/http/server/plugin.py
+++ b/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 ..descriptors import DescriptorsHandlerMixin
 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 @@ class HttpWebServerBasePlugin(DescriptorsHandlerMixin, ABC):
             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.
diff --git a/proxy/http/server/web.py b/proxy/http/server/web.py
index d31b3e6..fafc89e 100644
--- a/proxy/http/server/web.py
+++ b/proxy/http/server/web.py
@@ -106,6 +106,7 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
                 self.flags,
                 self.client,
                 self.event_queue,
+                self.upstream_conn_pool,
             )
             self.plugins[instance.name()] = instance
             for (protocol, route) in instance.routes():


Test patch:
diff --git a/tests/core/test_conn_pool.py b/tests/core/test_conn_pool.py
index ae2d2c7..7f38d99 100644
--- a/tests/core/test_conn_pool.py
+++ b/tests/core/test_conn_pool.py
@@ -21,30 +21,34 @@ from proxy.core.connection import UpstreamConnectionPool
 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])