Merge pull request #477 from minos-framework/issue-176-add-saga-runner-timeout

#176 - Add `timeout` to `SagaRunner.run`

Author: Sergio García Prado

packages/core/minos-microservice-common/minos/common/setup.py

@@ -4,6 +4,9 @@
 
 import logging
 import warnings
+from asyncio import (
+    shield,
+)
 from pathlib import (
     Path,
 )
@@ -95,7 +98,7 @@ async def setup(self) -> None:
         """
         if not self._already_setup:
             logger.debug(f"Setting up a {type(self).__name__!r} instance...")
-            await self._setup()
+            await shield(self._setup())
             self._already_setup = True
 
     async def _setup(self) -> None:
@@ -111,7 +114,7 @@ async def destroy(self) -> None:
         """
         if self._already_setup:
             logger.debug(f"Destroying a {type(self).__name__!r} instance...")
-            await self._destroy()
+            await shield(self._destroy())
             self._already_setup = False
 
     async def _destroy(self) -> None:

packages/core/minos-microservice-networks/tests/test_networks/test_brokers/test_collections/test_queues/test_database.py

@@ -75,9 +75,6 @@ async def test_aiter(self):
         ]
 
         queue = DatabaseBrokerQueue.from_config(self.config, operation_factory=self.operation_factory)
-        await queue.setup()
-        await queue.enqueue(messages[0])
-        await queue.enqueue(messages[1])
 
         with patch.object(
             MockedDatabaseClient,
@@ -90,6 +87,9 @@ async def test_aiter(self):
                 cycle([FakeAsyncIterator([(0,)])]),
             ),
         ):
+            await queue.setup()
+            await queue.enqueue(messages[0])
+            await queue.enqueue(messages[1])
 
             observed = list()
             async for message in queue:
@@ -110,11 +110,11 @@ async def test_dequeue_with_count(self):
             "fetch_all",
             return_value=FakeAsyncIterator([[1, messages[0].avro_bytes], [2, bytes()], [3, messages[1].avro_bytes]]),
         ):
-            async with DatabaseBrokerQueue.from_config(self.config, operation_factory=self.operation_factory) as queue:
-                queue._get_count = AsyncMock(side_effect=[3, 0])
+            queue = DatabaseBrokerQueue.from_config(self.config, operation_factory=self.operation_factory)
+            queue._get_count = AsyncMock(side_effect=[3, 0])
 
-                async with queue:
-                    observed = [await queue.dequeue(), await queue.dequeue()]
+            async with queue:
+                observed = [await queue.dequeue(), await queue.dequeue()]
 
         self.assertEqual(messages, observed)
 

packages/core/minos-microservice-saga/minos/saga/executions/runners.py

@@ -6,6 +6,10 @@
 
 import logging
 import warnings
+from asyncio import (
+    TimeoutError,
+    wait_for,
+)
 from functools import (
     reduce,
 )
@@ -96,6 +100,7 @@ async def run(
         pause_on_disk: bool = False,
         raise_on_error: bool = True,
         return_execution: bool = True,
+        timeout: Optional[float] = None,
         **kwargs,
     ) -> Union[UUID, SagaExecution]:
         """Perform a run of a ``Saga``.
@@ -114,73 +119,69 @@ async def run(
             but with ``Errored`` status.
         :param return_execution: If ``True`` the ``SagaExecution`` instance is returned. Otherwise, only the
             identifier (``UUID``) is returned.
+        :param timeout: Maximum execution time in seconds.
         :param kwargs: Additional named arguments.
         :return: This method does not return anything.
         """
         if isinstance(definition, SagaDecoratorWrapper):
             definition = definition.meta.definition
 
-        if response is not None:
-            return await self._load_and_run(
-                response=response,
-                autocommit=autocommit,
-                pause_on_disk=pause_on_disk,
-                raise_on_error=raise_on_error,
-                return_execution=return_execution,
-                **kwargs,
-            )
-
-        return await self._run_new(
-            definition=definition,
-            context=context,
-            user=user,
+        if response is None:
+            execution = await self._create(definition, context, user)
+        else:
+            execution = await self._load(response)
+
+        execution = await self._run(
+            execution,
+            timeout=timeout,
+            response=response,
             autocommit=autocommit,
             pause_on_disk=pause_on_disk,
             raise_on_error=raise_on_error,
-            return_execution=return_execution,
             **kwargs,
         )
+        if return_execution:
+            return execution
 
-    async def _run_new(
-        self, definition: Saga, context: Optional[SagaContext] = None, user: Optional[UUID] = None, **kwargs
-    ) -> Union[UUID, SagaExecution]:
+        return execution.uuid
+
+    @staticmethod
+    async def _create(definition: Saga, context: Optional[SagaContext], user: Optional[UUID]) -> SagaExecution:
         if REQUEST_USER_CONTEXT_VAR.get() is not None:
             if user is not None:
                 warnings.warn("The `user` Argument will be ignored in favor of the `user` ContextVar", RuntimeWarning)
             user = REQUEST_USER_CONTEXT_VAR.get()
 
-        execution = SagaExecution.from_definition(definition, context=context, user=user)
-        return await self._run(execution, **kwargs)
+        return SagaExecution.from_definition(definition, context=context, user=user)
 
-    async def _load_and_run(self, response: SagaResponse, **kwargs) -> Union[UUID, SagaExecution]:
-        execution = await self.storage.load(response.uuid)
-        return await self._run(execution, response=response, **kwargs)
+    async def _load(self, response: SagaResponse) -> Union[UUID, SagaExecution]:
+        return await self.storage.load(response.uuid)
 
-    async def _run(
-        self,
-        execution: SagaExecution,
-        pause_on_disk: bool = False,
-        raise_on_error: bool = True,
-        return_execution: bool = True,
-        **kwargs,
-    ) -> Union[UUID, SagaExecution]:
+    async def _run(self, execution: SagaExecution, raise_on_error: bool, **kwargs) -> SagaExecution:
         try:
-            if pause_on_disk:
-                await self._run_with_pause_on_disk(execution, **kwargs)
-            else:
-                await self._run_with_pause_on_memory(execution, **kwargs)
+            await self._run_with_timeout(execution, **kwargs)
         except SagaFailedExecutionException as exc:
             if raise_on_error:
                 raise exc
-            logger.exception(f"The execution identified by {execution.uuid!s} failed")
+            logger.exception(f"The execution identified by {execution.uuid!s} failed: {exc.exception!r}")
         finally:
             await self.storage.store(execution)
             self._update_request_headers(execution)
 
-        if return_execution:
-            return execution
+        return execution
 
-        return execution.uuid
+    async def _run_with_timeout(self, execution: SagaExecution, timeout: Optional[float], **kwargs) -> None:
+        future = self._run_with_pause(execution, **kwargs)
+        try:
+            return await wait_for(future, timeout=timeout)
+        except TimeoutError as exc:
+            raise SagaFailedExecutionException(exc)
+
+    async def _run_with_pause(self, execution: SagaExecution, pause_on_disk: bool, **kwargs) -> None:
+        if pause_on_disk:
+            await self._run_with_pause_on_disk(execution, **kwargs)
+        else:
+            await self._run_with_pause_on_memory(execution, **kwargs)
 
     @staticmethod
     def _update_request_headers(execution: SagaExecution) -> None:
@@ -197,9 +198,11 @@ def _update_request_headers(execution: SagaExecution) -> None:
         headers["related_services"] = ",".join(related_services)
 
     @staticmethod
-    async def _run_with_pause_on_disk(execution: SagaExecution, autocommit: bool = True, **kwargs) -> None:
+    async def _run_with_pause_on_disk(
+        execution: SagaExecution, response: Optional[SagaResponse] = None, autocommit: bool = True, **kwargs
+    ) -> None:
         try:
-            await execution.execute(autocommit=False, **kwargs)
+            await execution.execute(autocommit=False, response=response, **kwargs)
             if autocommit:
                 await execution.commit(**kwargs)
         except SagaPausedExecutionStepException:

packages/core/minos-microservice-saga/tests/test_saga/test_executions/test_runners.py

@@ -91,28 +91,62 @@ async def test_context_runner(self):
 
     async def test_run(self):
         expected = SagaExecution.from_definition(ADD_ORDER)
-        mock = AsyncMock(return_value=expected)
-        self.runner._run_new = mock
+        create_mock = AsyncMock()
+        run_mock = AsyncMock(return_value=expected)
+        self.runner._create = create_mock
+        self.runner._run = run_mock
 
         observed = await self.runner.run(ADD_ORDER)
         self.assertEqual(expected, observed)
 
     async def test_run_from_wrapper(self):
         expected = SagaExecution.from_definition(DeleteOrderSaga)
-        mock = AsyncMock(return_value=expected)
-        self.runner._run_new = mock
+        create_mock = AsyncMock()
+        run_mock = AsyncMock(return_value=expected)
+        self.runner._create = create_mock
+        self.runner._run = run_mock
 
         observed = await self.runner.run(DeleteOrderSaga)
         self.assertEqual(expected, observed)
 
     async def test_load_and_run(self):
         expected = SagaExecution.from_definition(ADD_ORDER)
-        mock = AsyncMock(return_value=expected)
-        self.runner._load_and_run = mock
+        load_mock = AsyncMock()
+        run_mock = AsyncMock(return_value=expected)
+        self.runner._load = load_mock
+        self.runner._run = run_mock
 
         observed = await self.runner.run(response=SagaResponse(uuid=expected.uuid))
         self.assertEqual(expected, observed)
 
+    async def test_run_with_timeout(self):
+        expected = SagaExecution.from_definition(ADD_ORDER)
+        create_mock = AsyncMock(return_value=expected)
+        store_mock = AsyncMock()
+        update_headers_mock = AsyncMock
+        run_mock = AsyncMock()
+        self.runner._create = create_mock
+        self.runner.storage.store = store_mock
+        self.runner._update_request_headers = update_headers_mock
+        self.runner._run_with_pause = run_mock
+
+        observed = await self.runner.run(ADD_ORDER, timeout=60)
+        self.assertEqual(expected, observed)
+
+    async def test_run_with_timeout_raises(self):
+        expected = SagaExecution.from_definition(ADD_ORDER)
+        create_mock = AsyncMock(return_value=expected)
+        store_mock = AsyncMock()
+        update_headers_mock = AsyncMock
+        run_mock = AsyncMock()
+        self.runner._create = create_mock
+        self.runner.storage.store = store_mock
+        self.runner._update_request_headers = update_headers_mock
+        self.runner._run_with_pause = run_mock
+
+        with self.assertRaises(SagaFailedExecutionException):
+            await self.runner.run(ADD_ORDER, timeout=0.0)
+
     async def test_run_with_pause_on_memory(self):
         self.broker_subscriber_builder.with_messages(
             [