✨ Aiohttp long running tasks client (ITISFoundation#3290)

sanderegg · web-flow · commit c604d7ceaf39 · 2022-08-30T21:12:48.000+02:00
diff --git a/packages/service-library/src/servicelib/aiohttp/long_running_tasks/_server.py b/packages/service-library/src/servicelib/aiohttp/long_running_tasks/_server.py
@@ -1,17 +1,26 @@
+import asyncio
 import logging
 from functools import wraps
-from typing import AsyncGenerator, Callable
+from typing import Any, AsyncGenerator, Callable
 
 from aiohttp import web
 from pydantic import PositiveFloat
-
-from ...long_running_tasks._task import TasksManager
+from servicelib.json_serialization import json_dumps
+
+from ...long_running_tasks._models import TaskGet
+from ...long_running_tasks._task import (
+    TaskContext,
+    TaskProtocol,
+    TasksManager,
+    start_task,
+)
 from ..typing_extension import Handler
 from ._constants import (
     APP_LONG_RUNNING_TASKS_MANAGER_KEY,
     MINUTE,
     RQT_LONG_RUNNING_TASKS_CONTEXT_KEY,
 )
+from ._dependencies import create_task_name_from_request, get_tasks_manager
 from ._error_handlers import base_long_running_error_handler
 from ._routes import routes
 
@@ -31,6 +40,49 @@ async def _wrap(request: web.Request):
     return _wrap
 
 
+async def start_long_running_task(
+    request: web.Request,
+    task: TaskProtocol,
+    *,
+    task_context: TaskContext,
+    **task_kwargs: Any,
+) -> web.Response:
+    task_manager = get_tasks_manager(request.app)
+    task_name = create_task_name_from_request(request)
+    task_id = None
+    try:
+        task_id = start_task(
+            task_manager,
+            task,
+            task_context=task_context,
+            task_name=task_name,
+            **task_kwargs,
+        )
+        status_url = request.app.router["get_task_status"].url_for(task_id=task_id)
+        result_url = request.app.router["get_task_result"].url_for(task_id=task_id)
+        abort_url = request.app.router["cancel_and_delete_task"].url_for(
+            task_id=task_id
+        )
+        task_get = TaskGet(
+            task_id=task_id,
+            task_name=task_name,
+            status_href=f"{status_url}",
+            result_href=f"{result_url}",
+            abort_href=f"{abort_url}",
+        )
+        return web.json_response(
+            data={"data": task_get},
+            status=web.HTTPAccepted.status_code,
+            dumps=json_dumps,
+        )
+    except asyncio.CancelledError:
+        # cancel the task, the client has disconnected
+        if task_id:
+            task_manager = get_tasks_manager(request.app)
+            await task_manager.cancel_task(task_id, with_task_context=None)
+        raise
+
+
 def setup(
     app: web.Application,
     *,
diff --git a/packages/service-library/src/servicelib/aiohttp/long_running_tasks/client.py b/packages/service-library/src/servicelib/aiohttp/long_running_tasks/client.py
@@ -0,0 +1,155 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Coroutine, Final, Optional
+
+from aiohttp import ClientConnectionError, ClientSession, web
+from pydantic import Json
+from tenacity import TryAgain, retry
+from tenacity._asyncio import AsyncRetrying
+from tenacity.retry import retry_if_exception_type
+from tenacity.stop import stop_after_delay
+from tenacity.wait import wait_random_exponential
+from yarl import URL
+
+from ..rest_responses import unwrap_envelope
+from .server import TaskGet, TaskId, TaskProgress, TaskStatus
+
+RequestBody = Json
+
+_MINUTE: Final[int] = 60
+_HOUR: Final[int] = 60 * _MINUTE
+_DEFAULT_POLL_INTERVAL_S: Final[float] = 1
+_DEFAULT_AIOHTTP_RETRY_POLICY = dict(
+    retry=retry_if_exception_type(ClientConnectionError),
+    wait=wait_random_exponential(max=20),
+    stop=stop_after_delay(30),
+    reraise=True,
+)
+
+
+@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+async def _start(
+    session: ClientSession, url: URL, json: Optional[RequestBody]
+) -> TaskGet:
+    async with session.post(url, json=json) as response:
+        response.raise_for_status()
+        data, error = unwrap_envelope(await response.json())
+    assert not error  # nosec
+    assert data is not None  # nosec
+    task = TaskGet.parse_obj(data)
+    return task
+
+
+@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+async def _wait_for_completion(
+    session: ClientSession,
+    task_id: TaskId,
+    status_url: URL,
+    wait_timeout_s: int,
+) -> AsyncGenerator[TaskProgress, None]:
+    try:
+
+        async for attempt in AsyncRetrying(
+            stop=stop_after_delay(wait_timeout_s),
+            reraise=True,
+            retry=retry_if_exception_type(TryAgain),
+        ):
+            with attempt:
+                async with session.get(status_url) as response:
+                    response.raise_for_status()
+                    data, error = unwrap_envelope(await response.json())
+                    assert not error  # nosec
+                    assert data is not None  # nosec
+                task_status = TaskStatus.parse_obj(data)
+                yield task_status.task_progress
+                if not task_status.done:
+                    await asyncio.sleep(
+                        float(
+                            response.headers.get(
+                                "retry-after", _DEFAULT_POLL_INTERVAL_S
+                            )
+                        )
+                    )
+                    raise TryAgain(
+                        f"{task_id=}, {task_status.started=} has "
+                        f"status: '{task_status.task_progress.message}'"
+                        f" {task_status.task_progress.percent}%"
+                    )
+    except TryAgain as exc:
+        # this is a timeout
+        raise asyncio.TimeoutError(
+            f"Long running task {task_id}, calling to {status_url} timed-out after {wait_timeout_s} seconds"
+        ) from exc
+
+
+@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+async def _task_result(session: ClientSession, result_url: URL) -> Any:
+    async with session.get(result_url) as response:
+        response.raise_for_status()
+        if response.status != web.HTTPNoContent.status_code:
+            data, error = unwrap_envelope(await response.json())
+            assert not error  # nosec
+            assert data  # nosec
+            return data
+
+
+@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+async def _abort_task(session: ClientSession, abort_url: URL) -> None:
+    async with session.delete(abort_url) as response:
+        response.raise_for_status()
+        data, error = unwrap_envelope(await response.json())
+        assert not error  # nosec
+        assert not data  # nosec
+
+
+@dataclass(frozen=True)
+class LRTask:
+    progress: TaskProgress
+    _result: Optional[Coroutine[Any, Any, Any]] = None
+
+    def done(self) -> bool:
+        return self._result is not None
+
+    async def result(self) -> Any:
+        if not self._result:
+            raise ValueError("No result ready!")
+        return await self._result
+
+
+async def long_running_task_request(
+    session: ClientSession,
+    url: URL,
+    json: Optional[RequestBody] = None,
+    client_timeout: int = 1 * _HOUR,
+) -> AsyncGenerator[LRTask, None]:
+    """Will use the passed `ClientSession` to call an oSparc long
+    running task `url` passing `json` as request body.
+    NOTE: this follows the usual aiohttp client syntax, and will raise the same errors
+
+    Raises:
+        [https://docs.aiohttp.org/en/stable/client_reference.html#hierarchy-of-exceptions]
+    """
+    task = None
+    try:
+        task = await _start(session, url, json)
+        last_progress = None
+        async for task_progress in _wait_for_completion(
+            session,
+            task.task_id,
+            url.with_path(task.status_href, encoded=True),
+            client_timeout,
+        ):
+            last_progress = task_progress
+            yield LRTask(progress=task_progress)
+        assert last_progress  # nosec
+        yield LRTask(
+            progress=last_progress,
+            _result=_task_result(
+                session, url.with_path(task.result_href, encoded=True)
+            ),
+        )
+
+    except (asyncio.CancelledError, asyncio.TimeoutError):
+        if task:
+            await _abort_task(session, url.with_path(task.abort_href, encoded=True))
+        raise
diff --git a/packages/service-library/src/servicelib/aiohttp/long_running_tasks/server.py b/packages/service-library/src/servicelib/aiohttp/long_running_tasks/server.py
@@ -13,17 +13,16 @@
     TaskProtocol,
     TasksManager,
     TaskStatus,
-    start_task,
 )
 from ._dependencies import create_task_name_from_request, get_tasks_manager
 from ._routes import TaskGet
-from ._server import setup
+from ._server import setup, start_long_running_task
 
 __all__: tuple[str, ...] = (
     "create_task_name_from_request",
     "get_tasks_manager",
     "setup",
-    "start_task",
+    "start_long_running_task",
     "TaskAlreadyRunningError",
     "TaskCancelledError",
     "TaskId",
diff --git a/packages/service-library/src/servicelib/long_running_tasks/_task.py b/packages/service-library/src/servicelib/long_running_tasks/_task.py
@@ -422,4 +422,5 @@ async def _progress_task(progress: TaskProgress, handler: TaskProtocol):
         task_progress=task_progress,
         task_context=task_context or {},
     )
+
     return tracked_task.task_id
diff --git a/packages/service-library/tests/aiohttp/long_running_tasks/conftest.py b/packages/service-library/tests/aiohttp/long_running_tasks/conftest.py
@@ -3,7 +3,6 @@
 # pylint: disable=unused-variable
 
 import asyncio
-import urllib.parse
 from typing import Awaitable, Callable
 
 import pytest
@@ -13,12 +12,8 @@
 from pydantic import BaseModel, parse_obj_as
 from pytest_simcore.helpers.utils_assert import assert_status
 from servicelib.aiohttp import long_running_tasks
-from servicelib.aiohttp.long_running_tasks.server import (
-    TaskId,
-    create_task_name_from_request,
-)
+from servicelib.aiohttp.long_running_tasks.server import TaskId
 from servicelib.aiohttp.requests_validation import parse_request_query_parameters_as
-from servicelib.json_serialization import json_dumps
 from servicelib.long_running_tasks._task import TaskContext
 from tenacity._asyncio import AsyncRetrying
 from tenacity.retry import retry_if_exception_type
@@ -68,26 +63,15 @@ class _LongTaskQueryParams(BaseModel):
         fail: bool = False
 
     @routes.post("/long_running_task:start", name=long_running_task_entrypoint)
-    async def generate_list_strings(request: web.Request):
-        task_manager = long_running_tasks.server.get_tasks_manager(request.app)
+    async def generate_list_strings(request: web.Request) -> web.Response:
         query_params = parse_request_query_parameters_as(_LongTaskQueryParams, request)
-        assert task_manager, "task manager is not initiated!"
-        task_name = create_task_name_from_request(request)
-        task_id = long_running_tasks.server.start_task(
-            task_manager,
+        return await long_running_tasks.server.start_long_running_task(
+            request,
             _string_list_task,
             num_strings=query_params.num_strings,
             sleep_time=query_params.sleep_time,
             fail=query_params.fail,
             task_context=task_context,
-            task_name=task_name,
-        )
-        assert task_id
-        assert task_id.startswith(urllib.parse.quote(task_name, safe=""))
-        return web.json_response(
-            data={"data": task_id},
-            status=web.HTTPAccepted.status_code,
-            dumps=json_dumps,
         )
 
     return routes
@@ -108,8 +92,8 @@ async def _caller(client: TestClient, **query_kwargs) -> TaskId:
         data, error = await assert_status(resp, web.HTTPAccepted)
         assert data
         assert not error
-        task_id = parse_obj_as(long_running_tasks.server.TaskId, data)
-        return task_id
+        task_get = parse_obj_as(long_running_tasks.server.TaskGet, data)
+        return task_get.task_id
 
     return _caller
 
diff --git a/packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks.py b/packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks.py
@@ -18,8 +18,6 @@
 from aiohttp import web
 from aiohttp.test_utils import TestClient
 from pydantic import parse_obj_as
-
-# TESTS
 from pytest_simcore.helpers.utils_assert import assert_status
 from servicelib.aiohttp import long_running_tasks
 from servicelib.aiohttp.long_running_tasks.server import TaskGet, TaskId
diff --git a/packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks_client.py b/packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks_client.py
diff --git a/services/web/server/src/simcore_service_webserver/long_running_tasks.py b/services/web/server/src/simcore_service_webserver/long_running_tasks.py
diff --git a/services/web/server/src/simcore_service_webserver/projects/projects_handlers_crud.py b/services/web/server/src/simcore_service_webserver/projects/projects_handlers_crud.py

Original file line number	Diff line number	Diff line change
`@@ -422,4 +422,5 @@ async def _progress_task(progress: TaskProgress, handler: TaskProtocol):`
`422`	`422`	`task_progress=task_progress,`
`423`	`423`	`task_context=task_context or {},`
`424`	`424`	`)`
	`425`	`+`
`425`	`426`	`return tracked_task.task_id`