Merge pull request #51 from thread/connection-pooling

Connection Pooling

routemaster/app.py

@@ -1,11 +1,16 @@
 """Core App singleton that holds state for the application."""
 import threading
+from typing import Dict
 
 from sqlalchemy.engine import Engine
 
 from routemaster.db import initialise_db
-from routemaster.config import Config
+from routemaster.config import Config, StateMachine
 from routemaster.logging import BaseLogger, SplitLogger, register_loggers
+from routemaster.webhooks import (
+    WebhookRunner,
+    webhook_runner_for_state_machine,
+)
 
 
 class App(threading.local):
@@ -14,6 +19,7 @@ class App(threading.local):
     db: Engine
     config: Config
     logger: BaseLogger
+    _webhook_runners: Dict[str, WebhookRunner]
 
     def __init__(
         self,
@@ -24,3 +30,14 @@ def __init__(
         self.db = initialise_db(self.config.database)
 
         self.logger = SplitLogger(config, loggers=register_loggers(config))
+
+        # Webhook runners may choose to persist a session, so we instantiate
+        # up-front to ensure we re-use state.
+        self._webhook_runners = {
+            x: webhook_runner_for_state_machine(y)
+            for x, y in self.config.state_machines.items()
+        }
+
+    def get_webhook_runner(self, state_machine: StateMachine) -> WebhookRunner:
+        """Get the webhook runner for a state machine."""
+        return self._webhook_runners[state_machine.name]

routemaster/conftest.py

@@ -36,7 +36,10 @@
 from routemaster.server import server
 from routemaster.context import Context
 from routemaster.logging import BaseLogger
-from routemaster.webhooks import WebhookResult
+from routemaster.webhooks import (
+    WebhookResult,
+    webhook_runner_for_state_machine,
+)
 from routemaster.state_machine import LabelRef
 from routemaster.exit_conditions import ExitConditionProgram
 
@@ -228,6 +231,10 @@ def __init__(self, config):
         self.config = config
         self.database_used = False
         self.logger = mock.MagicMock()
+        self._webhook_runners = {
+            x: webhook_runner_for_state_machine(y)
+            for x, y in self.config.state_machines.items()
+        }
 
     @property
     def db(self):
@@ -348,7 +355,7 @@ def mock_webhook():
     def _mock(result=WebhookResult.SUCCESS):
         runner = mock.Mock(return_value=result)
         with mock.patch(
-            'routemaster.webhooks.RequestsWebhookRunner',
+            'routemaster.app.App.get_webhook_runner',
             return_value=runner,
         ):
             yield runner

routemaster/feeds.py

@@ -1,4 +1,5 @@
 """Creation and fetching of feed data."""
+import threading
 from typing import Any, Dict, Callable, Optional
 
 import requests
@@ -20,6 +21,17 @@ class FeedNotFetched(Exception):
     pass
 
 
+_feed_sessions = threading.local()
+
+
+def _get_feed_session():
+    # We cache sessions per thread so that we can use `requests.Session`'s
+    # underlying `urllib3` connection pooling.
+    if not hasattr(_feed_sessions, 'session'):
+        _feed_sessions.session = requests.Session()
+    return _feed_sessions.session
+
+
 @dataclass
 class Feed:
     """A feed fetcher, able to retreive a feed and read keys out of it."""
@@ -38,7 +50,8 @@ def prefetch(
 
         url = template_url(self.url, self.state_machine, label)
 
-        response = requests.get(url)
+        session = _get_feed_session()
+        response = session.get(url)
         log_response(response)
         response.raise_for_status()
         self.data = response.json()

routemaster/state_machine/actions.py

@@ -9,10 +9,7 @@
 from routemaster.app import App
 from routemaster.utils import template_url
 from routemaster.config import State, Action, StateMachine
-from routemaster.webhooks import (
-    WebhookResult,
-    webhook_runner_for_state_machine,
-)
+from routemaster.webhooks import WebhookResult
 from routemaster.state_machine.types import LabelRef
 from routemaster.state_machine.utils import (
     choose_next_state,
@@ -58,7 +55,7 @@ def process_action(
         'label': label.name,
     }, sort_keys=True).encode('utf-8')
 
-    run_webhook = webhook_runner_for_state_machine(state_machine)
+    run_webhook = app.get_webhook_runner(state_machine)
 
     idempotency_token = _calculate_idempotency_token(label, latest_history)
 

routemaster/tests/test_layering.py

@@ -42,6 +42,7 @@
     ('app', 'db'),
     ('app', 'config'),
     ('app', 'logging'),
+    ('app', 'webhooks'),
 
     ('server', 'state_machine'),
     ('server', 'version'),

routemaster/tests/test_webhook_runner.py

@@ -5,7 +5,6 @@
 from routemaster.webhooks import (
     WebhookResult,
     RequestsWebhookRunner,
-    webhook_runner_for_state_machine,
 )
 
 
@@ -99,7 +98,7 @@ def test_requests_webhook_runner_for_state_machine_uses_webhook_config(app_confi
     )
 
     state_machine = app_config.config.state_machines['test_machine']
-    runner = webhook_runner_for_state_machine(state_machine)
+    runner = app_config.get_webhook_runner(state_machine)
     runner('http://example.com', 'application/test-data', b'\0\xff', '')
 
     last_request = httpretty.last_request()
@@ -117,7 +116,7 @@ def test_requests_webhook_runner_for_state_machine_does_not_apply_headers_for_no
     )
 
     state_machine = app_config.config.state_machines['test_machine']
-    runner = webhook_runner_for_state_machine(state_machine)
+    runner = app_config.get_webhook_runner(state_machine)
     runner('http://not-example.com', 'application/test-data', b'\0\xff', '')
 
     last_request = httpretty.last_request()

routemaster/webhooks.py

@@ -28,6 +28,8 @@ class RequestsWebhookRunner(object):
     """
 
     def __init__(self, webhook_configs: Iterable[Webhook]=()) -> None:
+        # Use a session so that we can take advantage of connection pooling in
+        # `urllib3`.
         self.session = requests.Session()
         self.webhook_configs = webhook_configs