Repository: mercadona/rele
Pull request #243 (2023-04-24 13:46:32)

Add exponential backoff
### :tophat: What?

Add retry policy to the subscriptions (exponential backoff).

### :thinking: Why?

This feature allows to configure either a global (per project) retry policy or at subscription level, allowing us leveraging how  we want the subscriptions to behave (retry wise) when worker fails processing a message.

[More info about the feature here](https://cloud.google.com/pubsub/docs/handling-failures).

### :link: Related issue

#222 



Patch:
diff --git a/docs/settings.rst b/docs/settings.rst
index 3b5edb6..68acda1 100644
--- a/docs/settings.rst
+++ b/docs/settings.rst
@@ -29,7 +29,8 @@ Example::
         'ENCODER_PATH': 'rest_framework.utils.encoders.JSONEncoder',
         'ACK_DEADLINE': 120,
         'PUBLISHER_TIMEOUT': 3.0,
-        'FILTER_SUBS_BY': boolean_function
+        'FILTER_SUBS_BY': boolean_function,
+        'DEFAULT_RETRY_POLICY': RetryPolicy(10, 50),
     }
 
 ``GC_PROJECT_ID``
@@ -174,3 +175,15 @@ library behavior, please set this value to 10.
 
 Boolean function that applies a global filter on all subscriptions.
 For more information, please see `Filtering Messages section <https://mercadonarele.readthedocs.io/en/latest/guides/filters.html#global-filter>`_.
+
+
+``DEFAULT_RETRY_POLICY``
+----------------------------
+
+**Optional**
+
+A RetryPolicy object which must be instantiated with `minimum_backoff` and `maximum_backoff`, that specifies in seconds how Pub/Sub retries message delivery for all the subscriptions.
+
+If not set, the default retry policy is applied, meaning a minimum backoff of 10 seconds and a maximum backoff of 60 seconds.
+This generally implies that messages will be retried as soon as possible for healthy subscribers.
+RetryPolicy will be triggered on NACKs or acknowledgement deadline exceeded events for a given message.
diff --git a/rele/client.py b/rele/client.py
index e62b334..47d381b 100644
--- a/rele/client.py
+++ b/rele/client.py
@@ -8,6 +8,8 @@ from contextlib import suppress
 import google.auth
 from google.api_core import exceptions
 from google.cloud import pubsub_v1
+from google.protobuf import duration_pb2
+from google.pubsub_v1 import RetryPolicy as GCloudRetryPolicy
 
 from rele.middleware import run_middleware_hook
 
@@ -36,13 +38,21 @@ class Subscriber:
     :param gc_project_id: str :ref:`settings_project_id` .
     :param credentials: obj :meth:`~rele.config.Config.credentials`.
     :param default_ack_deadline: int Ack Deadline defined in settings
+    :param default_retry_policy: RetryPolicy Rele's RetryPolicy defined in settings
     """
 
-    def __init__(self, gc_project_id, credentials, default_ack_deadline=None):
+    def __init__(
+        self,
+        gc_project_id,
+        credentials,
+        default_ack_deadline=None,
+        default_retry_policy=None,
+    ):
         self._gc_project_id = gc_project_id
         self._ack_deadline = default_ack_deadline or DEFAULT_ACK_DEADLINE
         self.credentials = credentials if not USE_EMULATOR else None
         self._client = pubsub_v1.SubscriberClient(credentials=credentials)
+        self._retry_policy = default_retry_policy
 
     def create_subscription(self, subscription):
         """Handles creating the subscription when it does not exists.
@@ -52,8 +62,7 @@ class Subscriber:
         have a topic to subscribe to. Which means that the topic must be
         created manually before the worker is started.
 
-        :param subscription: str Subscription name
-        :param topic: str Topic name to subscribe
+        :param subscription: obj :class:`~rele.subscription.Subscription`.
         """
         subscription_path = self._client.subscription_path(
             self._gc_project_id, subscription.name
@@ -86,8 +95,25 @@ class Subscriber:
         if subscription.backend_filter_by:
             request["filter"] = subscription.backend_filter_by
 
+        retry_policy = subscription.retry_policy or self._retry_policy
+
+        if retry_policy:
+            request["retry_policy"] = self._build_gcloud_retry_policy(retry_policy)
+
         self._client.create_subscription(request=request)
 
+    def _build_gcloud_retry_policy(self, rele_retry_policy):
+        minimum_backoff = duration_pb2.Duration(
+            seconds=rele_retry_policy.minimum_backoff
+        )
+        maximum_backoff = duration_pb2.Duration(
+            seconds=rele_retry_policy.maximum_backoff
+        )
+
+        return GCloudRetryPolicy(
+            minimum_backoff=minimum_backoff, maximum_backoff=maximum_backoff
+        )
+
     def consume(self, subscription_name, callback, scheduler):
         """Begin listening to topic from the SubscriberClient.
 
diff --git a/rele/config.py b/rele/config.py
index 8dd62e6..319acb1 100644
--- a/rele/config.py
+++ b/rele/config.py
@@ -42,6 +42,7 @@ class Config:
         self.threads_per_subscription = setting.get("THREADS_PER_SUBSCRIPTION", 2)
         self.filter_by = setting.get("FILTER_SUBS_BY")
         self._credentials = None
+        self.retry_policy = setting.get("DEFAULT_RETRY_POLICY")
 
     @property
     def encoder(self):
diff --git a/rele/retry_policy.py b/rele/retry_policy.py
new file mode 100644
index 0000000..7c136ec
--- /dev/null
+++ b/rele/retry_policy.py
@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class RetryPolicy:
+    """A RetryPolicy object encapsulates the validation rules.
+
+    Defines retry policy settings and ensures the values are correct.
+    If provided values are wrong, a ValidationError is raised.
+
+    :param minimum_backoff: int Accepts values greater than 0
+    :param maximum_backoff: int Accepts values greater than minimum_backoff.
+    """
+
+    minimum_backoff: int
+    maximum_backoff: int
+
+    def __init__(self, minimum_backoff, maximum_backoff):
+        self._guard_against_wrong_parameters(minimum_backoff, maximum_backoff)
+
+        self.minimum_backoff = minimum_backoff
+        self.maximum_backoff = maximum_backoff
+
+    def _guard_against_wrong_parameters(self, minimum_backoff, maximum_backoff):
+        if minimum_backoff == 0:
+            raise ValueError("minimum_backoff must be greater than 0")
+
+        if maximum_backoff == 0:
+            raise ValueError("maximum_backoff must be greater than 0")
+
+        if minimum_backoff > maximum_backoff:
+            raise ValueError("minimum_backoff should be less than maximum_backoff.")
diff --git a/rele/subscription.py b/rele/subscription.py
index 078ae78..7b02145 100644
--- a/rele/subscription.py
+++ b/rele/subscription.py
@@ -35,16 +35,26 @@ class Subscription:
     """
 
     def __init__(
-        self, func, topic, prefix="", suffix="", filter_by=None, backend_filter_by=None
+        self,
+        func,
+        topic,
+        prefix="",
+        suffix="",
+        filter_by=None,
+        backend_filter_by=None,
+        retry_policy=None,
     ):
+        self._validate_filter_by(filter_by)
+
         self._func = func
         self.topic = topic
         self._prefix = prefix
         self._suffix = suffix
         self._filters = self._init_filters(filter_by)
         self.backend_filter_by = backend_filter_by
+        self.retry_policy = retry_policy
 
-    def _init_filters(self, filter_by):
+    def _validate_filter_by(self, filter_by):
         if filter_by and not (
             callable(filter_by)
             or (
@@ -54,6 +64,7 @@ class Subscription:
         ):
             raise ValueError("Filter_by must be a callable or a list of callables.")
 
+    def _init_filters(self, filter_by):
         if isinstance(filter_by, Iterable):
             return filter_by
         elif filter_by:
@@ -145,7 +156,14 @@ class Callback:
             run_middleware_hook("post_process_message")
 
 
-def sub(topic, prefix=None, suffix=None, filter_by=None, backend_filter_by=None):
+def sub(
+    topic,
+    prefix=None,
+    suffix=None,
+    filter_by=None,
+    backend_filter_by=None,
+    retry_policy=None,
+):
     """Decorator function that makes declaring a PubSub Subscription simple.
 
     The Subscriber returned will automatically create and name
@@ -190,6 +208,7 @@ def sub(topic, prefix=None, suffix=None, filter_by=None, backend_filter_by=None)
     :param filter_by: Union[function, list] An optional function or tuple of
                       functions that filters the messages to be processed by
                       the sub regarding their attributes.
+    :param retry_policy: obj :class:`~rele.retry_policy.RetryPolicy`
     :return: :class:`~rele.subscription.Subscription`
     """
 
@@ -214,6 +233,7 @@ def sub(topic, prefix=None, suffix=None, filter_by=None, backend_filter_by=None)
             suffix=suffix,
             filter_by=filter_by,
             backend_filter_by=backend_filter_by,
+            retry_policy=retry_policy,
         )
 
     return decorator
diff --git a/rele/worker.py b/rele/worker.py
index 639dae5..82cd1c6 100644
--- a/rele/worker.py
+++ b/rele/worker.py
@@ -29,8 +29,11 @@ class Worker:
         credentials=None,
         default_ack_deadline=None,
         threads_per_subscription=None,
+        default_retry_policy=None,
     ):
-        self._subscriber = Subscriber(gc_project_id, credentials, default_ack_deadline)
+        self._subscriber = Subscriber(
+            gc_project_id, credentials, default_ack_deadline, default_retry_policy
+        )
         self._futures = {}
         self._subscriptions = subscriptions
         self.threads_per_subscription = threads_per_subscription
@@ -152,6 +155,7 @@ def create_and_run(subs, config):
         config.credentials,
         config.ack_deadline,
         config.threads_per_subscription,
+        config.retry_policy,
     )
 
     # to allow killing runrele worker via ctrl+c


Test patch:
diff --git a/tests/commands/test_runrele.py b/tests/commands/test_runrele.py
index 06128fe..ab03e00 100644
--- a/tests/commands/test_runrele.py
+++ b/tests/commands/test_runrele.py
@@ -22,7 +22,7 @@ class TestRunReleCommand:
     def test_calls_worker_start_and_setup_when_runrele(self, mock_worker):
         call_command("runrele")
 
-        mock_worker.assert_called_with([], "rele-test", ANY, 60, 2)
+        mock_worker.assert_called_with([], "rele-test", ANY, 60, 2, None)
         mock_worker.return_value.run_forever.assert_called_once_with()
 
     def test_prints_warning_when_conn_max_age_not_set_to_zero(
@@ -37,5 +37,5 @@ class TestRunReleCommand:
             "This may result in slots for database connections to "
             "be exhausted." in err
         )
-        mock_worker.assert_called_with([], "rele-test", ANY, 60, 2)
+        mock_worker.assert_called_with([], "rele-test", ANY, 60, 2, None)
         mock_worker.return_value.run_forever.assert_called_once_with()
diff --git a/tests/conftest.py b/tests/conftest.py
index d3bfb61..2e7464b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -12,6 +12,7 @@ from rele import Publisher
 from rele.client import Subscriber
 from rele.config import Config
 from rele.middleware import register_middleware
+from rele.retry_policy import RetryPolicy
 
 
 @pytest.fixture
@@ -31,6 +32,19 @@ def config(project_id):
     )
 
 
+@pytest.fixture
+def config_with_retry_policy(project_id):
+    return Config(
+        {
+            "APP_NAME": "rele",
+            "SUB_PREFIX": "rele",
+            "GC_CREDENTIALS_PATH": "tests/dummy-pub-sub-credentials.json",
+            "MIDDLEWARE": ["rele.contrib.LoggingMiddleware"],
+            "DEFAULT_RETRY_POLICY": RetryPolicy(5, 30),
+        }
+    )
+
+
 @pytest.fixture
 def subscriber(project_id, config):
     return Subscriber(config.gc_project_id, config.credentials, 60)
diff --git a/tests/test_client.py b/tests/test_client.py
index 542ade8..53581bd 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -6,8 +6,12 @@ from unittest.mock import ANY, patch
 
 import pytest
 from google.api_core import exceptions
+from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
+from google.protobuf import duration_pb2
 
+from rele import Subscriber
+from rele.retry_policy import RetryPolicy
 from rele.subscription import Subscription
 
 
@@ -286,3 +290,68 @@ class TestSubscriber:
                 "filter": backend_filter_by,
             }
         )
+
+    @patch.object(SubscriberClient, "create_subscription")
+    def test_creates_subscription_with_retry_policy_when_provided(
+        self, _mocked_client, project_id, subscriber
+    ):
+        expected_subscription = (
+            f"projects/{project_id}/subscriptions/" f"{project_id}-test-topic"
+        )
+        expected_topic = f"projects/{project_id}/topics/" f"{project_id}-test-topic"
+        expected_retry_policy = pubsub_v1.types.RetryPolicy(
+            minimum_backoff=duration_pb2.Duration(seconds=10),
+            maximum_backoff=duration_pb2.Duration(seconds=50),
+        )
+
+        subscriber.create_subscription(
+            Subscription(
+                None,
+                topic=f"{project_id}-test-topic",
+                retry_policy=RetryPolicy(10, 50),
+            )
+        )
+
+        _mocked_client.assert_called_once_with(
+            request={
+                "ack_deadline_seconds": 60,
+                "name": expected_subscription,
+                "topic": expected_topic,
+                "retry_policy": expected_retry_policy,
+            }
+        )
+
+    @patch.object(SubscriberClient, "create_subscription")
+    def test_default_retry_policy_is_applied_when_not_explicitly_provided(
+        self, _mocked_client, project_id, config_with_retry_policy
+    ):
+        subscriber = Subscriber(
+            config_with_retry_policy.gc_project_id,
+            config_with_retry_policy.credentials,
+            60,
+            config_with_retry_policy.retry_policy,
+        )
+        expected_subscription = (
+            f"projects/{project_id}/subscriptions/" f"{project_id}-test-topic"
+        )
+        expected_topic = f"projects/{project_id}/topics/" f"{project_id}-test-topic"
+        expected_retry_policy = pubsub_v1.types.RetryPolicy(
+            minimum_backoff=duration_pb2.Duration(seconds=5),
+            maximum_backoff=duration_pb2.Duration(seconds=30),
+        )
+
+        subscriber.create_subscription(
+            Subscription(
+                None,
+                topic=f"{project_id}-test-topic",
+            )
+        )
+
+        _mocked_client.assert_called_once_with(
+            request={
+                "ack_deadline_seconds": 60,
+                "name": expected_subscription,
+                "topic": expected_topic,
+                "retry_policy": expected_retry_policy,
+            }
+        )
diff --git a/tests/test_retry_policy.py b/tests/test_retry_policy.py
new file mode 100644
index 0000000..72bcef3
--- /dev/null
+++ b/tests/test_retry_policy.py
@@ -0,0 +1,19 @@
+import pytest
+
+from rele.retry_policy import RetryPolicy
+
+
+class TestRetryPolicy:
+    @pytest.mark.parametrize(
+        "minimum_backoff, maximum_backoff",
+        [
+            (0, 0),
+            (0, 1),
+            (10, 1),
+        ],
+    )
+    def test_value_error_is_raised_instantiating_with_wrong_values(
+        self, minimum_backoff, maximum_backoff
+    ):
+        with pytest.raises(ValueError):
+            RetryPolicy(minimum_backoff, maximum_backoff)
diff --git a/tests/test_subscription.py b/tests/test_subscription.py
index 6b1f800..a236faf 100644
--- a/tests/test_subscription.py
+++ b/tests/test_subscription.py
@@ -9,6 +9,7 @@ from google.protobuf import timestamp_pb2
 
 from rele import Callback, Subscription, sub
 from rele.middleware import register_middleware
+from rele.retry_policy import RetryPolicy
 
 logger = logging.getLogger(__name__)
 
@@ -112,7 +113,7 @@ class TestSubscription:
 
         assert response is None
 
-    def test_raises_error_when_filter_by_is_not_valid(self, caplog):
+    def test_raises_error_when_filter_by_is_not_valid(self):
         Subscription(
             func=lambda x: None, topic="topic", prefix="rele", filter_by=lambda x: True
         )
@@ -361,3 +362,12 @@ class TestDecorator:
             "Subscription function tests.test_subscription.<lambda> is outside a subs "
             "module that will not be discovered." in caplog.text
         )
+
+    def test_retry_policy_is_applied_when_specified(self):
+        subscription = sub(
+            topic="topic",
+            prefix="rele",
+            retry_policy=RetryPolicy(1, 10),
+        )(lambda data, **kwargs: None)
+
+        assert subscription.retry_policy == RetryPolicy(1, 10)
diff --git a/tests/test_worker.py b/tests/test_worker.py
index 94d610d..bd3b1d5 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -8,6 +8,7 @@ from google.cloud.pubsub_v1.subscriber.scheduler import ThreadScheduler
 
 from rele import Subscriber, Worker, sub
 from rele.middleware import register_middleware
+from rele.retry_policy import RetryPolicy
 from rele.subscription import Callback
 from rele.worker import create_and_run
 
@@ -26,6 +27,7 @@ def worker(config):
         config.credentials,
         default_ack_deadline=60,
         threads_per_subscription=10,
+        default_retry_policy=config.retry_policy,
     )
 
 
@@ -156,9 +158,18 @@ class TestCreateAndRun:
         with patch("rele.worker.Worker", autospec=True) as p:
             yield p
 
-    def test_waits_forever_when_called_with_config_and_subs(self, config, mock_worker):
+    def test_waits_forever_when_called_with_config_and_subs(
+        self, config_with_retry_policy, mock_worker
+    ):
         subscriptions = (sub_stub,)
-        create_and_run(subscriptions, config)
+        create_and_run(subscriptions, config_with_retry_policy)
 
-        mock_worker.assert_called_with(subscriptions, "rele-test", ANY, 60, 2)
+        mock_worker.assert_called_with(
+            subscriptions,
+            "rele-test",
+            ANY,
+            60,
+            2,
+            RetryPolicy(5, 30),
+        )
         mock_worker.return_value.run_forever.assert_called_once_with()