Repository: mercadona/rele
Pull request #139 (2019-12-12 20:11:33)

Add subscription ThreadPoolExecutor
### :tophat: What?

Add custom ThreadPoolExecutor to each subscriber. It follows the pattern shown [ here.](https://googleapis.dev/python/pubsub/latest/_modules/google/cloud/pubsub_v1/subscriber/scheduler.html#ThreadScheduler) 

And also allows customer thread count per subscription. 

Note: Ill leave this as a draft since I strongly believe that code coverage can be increased. If this is the direction we want to go, then Ill add it. 

### :thinking: Why?

We want to be able to control the number of thread per subscription. 

### :link: Related issue

Addresses #135 



Patch:
diff --git a/docs/api/settings.rst b/docs/api/settings.rst
index 176177d..31c8d75 100644
--- a/docs/api/settings.rst
+++ b/docs/api/settings.rst
@@ -134,3 +134,13 @@ Timeout that the publishing result will wait on the future to publish successful
 
 `See Google PubSub documentation for more info
 <https://googleapis.dev/python/pubsub/1.1.0/publisher/api/futures.html?highlight=result#google.cloud.pubsub_v1.publisher.futures.Future.result>`_
+
+``THREADS_PER_SUBSCRIPTION``
+------------------
+
+**Optional**
+
+Default: 10
+
+Number of threads that will be consumed for each subscription.
+10 is the default for the Google Cloud PubSub library.
diff --git a/rele/client.py b/rele/client.py
index 749cf1f..cb364eb 100644
--- a/rele/client.py
+++ b/rele/client.py
@@ -57,17 +57,20 @@ class Subscriber:
                 ack_deadline_seconds=self._ack_deadline,
             )
 
-    def consume(self, subscription_name, callback):
+    def consume(self, subscription_name, callback, scheduler):
         """Begin listening to topic from the SubscriberClient.
         
         :param subscription_name: str Subscription name
         :param callback: Function which act on a topic message
+        :param scheduler: `Thread pool-based scheduler.<https://googleapis.dev/python/pubsub/latest/subscriber/api/scheduler.html?highlight=threadscheduler#google.cloud.pubsub_v1.subscriber.scheduler.ThreadScheduler>`_  # noqa
         :return: `Future <https://googleapis.github.io/google-cloud-python/latest/pubsub/subscriber/api/futures.html>`_  # noqa
         """
         subscription_path = self._client.subscription_path(
             self._gc_project_id, subscription_name
         )
-        return self._client.subscribe(subscription_path, callback=callback)
+        return self._client.subscribe(
+            subscription_path, callback=callback, scheduler=scheduler
+        )
 
 
 class Publisher:
diff --git a/rele/config.py b/rele/config.py
index 1d0c8c9..4c576a0 100644
--- a/rele/config.py
+++ b/rele/config.py
@@ -29,6 +29,7 @@ class Config:
         )
         self._encoder_path = setting.get("ENCODER_PATH", DEFAULT_ENCODER_PATH)
         self.publisher_timeout = setting.get("PUBLISHER_TIMEOUT", 3.0)
+        self.threads_per_subscription = setting.get("THREADS_PER_SUBSCRIPTION", 10)
 
     @property
     def encoder(self):
diff --git a/rele/management/commands/runrele.py b/rele/management/commands/runrele.py
index 25649d3..3c7e2e5 100644
--- a/rele/management/commands/runrele.py
+++ b/rele/management/commands/runrele.py
@@ -35,6 +35,7 @@ class Command(BaseCommand):
             self.config.gc_project_id,
             self.config.credentials,
             self.config.ack_deadline,
+            self.config.threads_per_subscription,
         )
 
         signal.signal(signal.SIGINT, signal.SIG_IGN)
diff --git a/rele/subscription.py b/rele/subscription.py
index 950d7e4..846618d 100644
--- a/rele/subscription.py
+++ b/rele/subscription.py
@@ -1,6 +1,7 @@
 import json
 import logging
 import time
+from concurrent import futures
 
 from .middleware import run_middleware_hook
 
diff --git a/rele/worker.py b/rele/worker.py
index 88fd2f7..567d1a3 100644
--- a/rele/worker.py
+++ b/rele/worker.py
@@ -1,6 +1,7 @@
 import logging
 import sys
 import time
+from concurrent import futures
 
 from .client import Subscriber
 from .middleware import run_middleware_hook
@@ -18,10 +19,18 @@ class Worker:
     :param subscriptions: list :class:`~rele.subscription.Subscription`
     """
 
-    def __init__(self, subscriptions, gc_project_id, credentials, default_ack_deadline):
+    def __init__(
+        self,
+        subscriptions,
+        gc_project_id,
+        credentials,
+        default_ack_deadline,
+        threads_per_subscription,
+    ):
         self._subscriber = Subscriber(gc_project_id, credentials, default_ack_deadline)
         self._futures = []
         self._subscriptions = subscriptions
+        self.threads_per_subscription = threads_per_subscription
 
     def setup(self):
         """Create the subscriptions on a Google PubSub topic.
@@ -44,9 +53,17 @@ class Worker:
         """
         run_middleware_hook("pre_worker_start")
         for subscription in self._subscriptions:
+            executor_kwargs = {
+                "thread_name_prefix": "ThreadPoolExecutor-ThreadScheduler"
+            }
+            scheduler = futures.ThreadPoolExecutor(
+                max_workers=self.threads_per_subscription, **executor_kwargs
+            )
             self._futures.append(
                 self._subscriber.consume(
-                    subscription_name=subscription.name, callback=Callback(subscription)
+                    subscription_name=subscription.name,
+                    callback=Callback(subscription),
+                    scheduler=scheduler,
                 )
             )
         run_middleware_hook("post_worker_start")


Test patch:
diff --git a/tests/commands/test_runrele.py b/tests/commands/test_runrele.py
index df9385a..e6c6c38 100644
--- a/tests/commands/test_runrele.py
+++ b/tests/commands/test_runrele.py
@@ -20,7 +20,7 @@ class TestRunReleCommand:
     def test_calls_worker_start_and_setup_when_runrele(self, mock_worker):
         call_command("runrele")
 
-        mock_worker.assert_called_with([], "SOME-PROJECT-ID", ANY, 60)
+        mock_worker.assert_called_with([], "SOME-PROJECT-ID", ANY, 60, 10)
         mock_worker.return_value.run_forever.assert_called_once_with()
 
     def test_prints_warning_when_conn_max_age_not_set_to_zero(
@@ -35,5 +35,5 @@ class TestRunReleCommand:
             "This may result in slots for database connections to "
             "be exhausted." in err
         )
-        mock_worker.assert_called_with([], "SOME-PROJECT-ID", ANY, 60)
+        mock_worker.assert_called_with([], "SOME-PROJECT-ID", ANY, 60, 10)
         mock_worker.return_value.run_forever.assert_called_once_with()
diff --git a/tests/test_worker.py b/tests/test_worker.py
index c304e76..33507ad 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -1,3 +1,4 @@
+from concurrent import futures
 from unittest.mock import ANY, patch
 
 import pytest
@@ -12,9 +13,15 @@ def sub_stub(data, **kwargs):
 
 
 @pytest.fixture
-def worker(project_id, credentials):
+def worker(project_id, credentials, config):
     subscriptions = (sub_stub,)
-    return Worker(subscriptions, project_id, credentials, 60)
+    return Worker(
+        subscriptions,
+        project_id,
+        credentials,
+        default_ack_deadline=60,
+        threads_per_subscription=10,
+    )
 
 
 @pytest.fixture
@@ -36,8 +43,10 @@ class TestWorker:
         worker.start()
 
         mock_consume.assert_called_once_with(
-            subscription_name="rele-some-cool-topic", callback=ANY
+            subscription_name="rele-some-cool-topic", callback=ANY, scheduler=ANY
         )
+        scheduler = mock_consume.call_args_list[0][1]["scheduler"]
+        assert isinstance(scheduler, futures.ThreadPoolExecutor)
 
     def test_setup_creates_subscription_when_topic_given(
         self, mock_create_subscription, worker
@@ -58,8 +67,10 @@ class TestWorker:
         subscription = "rele-some-cool-topic"
         mock_create_subscription.assert_called_once_with(subscription, topic)
         mock_consume.assert_called_once_with(
-            subscription_name="rele-some-cool-topic", callback=ANY
+            subscription_name="rele-some-cool-topic", callback=ANY, scheduler=ANY
         )
+        scheduler = mock_consume.call_args_list[0][1]["scheduler"]
+        assert isinstance(scheduler, futures.ThreadPoolExecutor)
         mock_wait_forever.assert_called_once()
 
     @patch.object(Worker, "_wait_forever")
@@ -87,7 +98,13 @@ class TestWorker:
     ):
         subscriptions = (sub_stub,)
         custom_ack_deadline = 234
-        worker = Worker(subscriptions, project_id, credentials, custom_ack_deadline)
+        worker = Worker(
+            subscriptions,
+            project_id,
+            credentials,
+            custom_ack_deadline,
+            threads_per_subscription=10,
+        )
         worker.setup()
 
         assert worker._subscriber._ack_deadline == custom_ack_deadline