Mitigate hanging requests from api-server to storage (#7918)

Author: bisgaard-itis

services/api-server/src/simcore_service_api_server/api/routes/files.py

@@ -11,22 +11,18 @@
 from fastapi.exceptions import HTTPException
 from fastapi_pagination.api import create_page
 from models_library.api_schemas_storage.storage_schemas import (
-    ETag,
     FileUploadCompletionBody,
-    LinkType,
 )
 from models_library.basic_types import SHA256Str
 from models_library.projects_nodes_io import NodeID
-from pydantic import AnyUrl, ByteSize, PositiveInt, TypeAdapter, ValidationError
+from pydantic import PositiveInt, ValidationError
 from servicelib.fastapi.requests_decorators import cancel_on_disconnect
+from servicelib.logging_utils import log_context
 from simcore_sdk.node_ports_common.constants import SIMCORE_LOCATION
 from simcore_sdk.node_ports_common.file_io_utils import UploadableFileObject
 from simcore_sdk.node_ports_common.filemanager import (
     UploadedFile,
     UploadedFolder,
-    abort_upload,
-    complete_file_upload,
-    get_upload_links_from_s3,
 )
 from simcore_sdk.node_ports_common.filemanager import upload_path as storage_upload_path
 from starlette.datastructures import URL
@@ -297,23 +293,21 @@ async def get_upload_links(
     client_file: UserFileToProgramJob | UserFile,
     user_id: Annotated[PositiveInt, Depends(get_current_user_id)],
     webserver_api: Annotated[AuthSession, Depends(get_webserver_session)],
+    storage_client: Annotated[StorageApi, Depends(get_api_client(StorageApi))],
 ):
     """Get upload links for uploading a file to storage"""
     assert request  # nosec
     file_meta = await _create_domain_file(
         webserver_api=webserver_api, file_id=None, client_file=client_file
     )
-    _, upload_links = await get_upload_links_from_s3(
-        user_id=user_id,
-        store_name=None,
-        store_id=SIMCORE_LOCATION,
-        s3_object=file_meta.storage_file_id,
-        client_session=None,
-        link_type=LinkType.PRESIGNED,
-        file_size=ByteSize(client_file.filesize),
-        is_directory=False,
-        sha256_checksum=file_meta.sha256_checksum,
-    )
+
+    with log_context(
+        logger=_logger, level=logging.DEBUG, msg=f"Getting upload links for {file_meta}"
+    ):
+        upload_links = await storage_client.get_file_upload_links(
+            user_id=user_id, file=file_meta, client_file=client_file
+        )
+
     completion_url: URL = request.url_for(
         "complete_multipart_upload", file_id=file_meta.id
     )
@@ -420,12 +414,7 @@ async def abort_multipart_upload(
     file = await _create_domain_file(
         webserver_api=webserver_api, file_id=file_id, client_file=client_file
     )
-    abort_link: URL = await storage_client.create_abort_upload_link(
-        file=file, query={"user_id": str(user_id)}
-    )
-    await abort_upload(
-        abort_upload_link=TypeAdapter(AnyUrl).validate_python(str(abort_link))
-    )
+    await storage_client.abort_file_upload(user_id=user_id, file=file)
 
 
 @router.post(
@@ -449,13 +438,8 @@ async def complete_multipart_upload(
     file = await _create_domain_file(
         webserver_api=webserver_api, file_id=file_id, client_file=client_file
     )
-    complete_link: URL = await storage_client.create_complete_upload_link(
-        file=file, query={"user_id": str(user_id)}
-    )
-
-    e_tag: ETag | None = await complete_file_upload(
-        uploaded_parts=uploaded_parts.parts,
-        upload_completion_link=TypeAdapter(AnyUrl).validate_python(f"{complete_link}"),
+    e_tag = await storage_client.complete_file_upload(
+        user_id=user_id, file=file, uploaded_parts=uploaded_parts.parts
     )
     assert e_tag is not None  # nosec
 

services/api-server/src/simcore_service_api_server/exceptions/backend_errors.py

@@ -17,6 +17,11 @@ def named_fields(cls) -> set[str]:
         )
 
 
+class BackendTimeoutError(BaseBackEndError):
+    msg_template = "Backend request timed out"
+    status_code = status.HTTP_504_GATEWAY_TIMEOUT
+
+
 class InvalidInputError(BaseBackEndError):
     msg_template = "Invalid input"
     status_code = status.HTTP_422_UNPROCESSABLE_ENTITY

services/api-server/src/simcore_service_api_server/services_http/storage.py

@@ -1,35 +1,57 @@
 import logging
 import re
 import urllib.parse
+from datetime import timedelta
 from functools import partial
 from mimetypes import guess_type
-from typing import Literal
+from typing import Final, Literal
 from uuid import UUID
 
 from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
+from httpx import QueryParams
 from models_library.api_schemas_storage.storage_schemas import (
+    ETag,
     FileMetaDataArray,
 )
 from models_library.api_schemas_storage.storage_schemas import (
     FileMetaDataGet as StorageFileMetaData,
 )
 from models_library.api_schemas_storage.storage_schemas import (
+    FileUploadCompleteFutureResponse,
+    FileUploadCompleteResponse,
+    FileUploadCompleteState,
+    FileUploadCompletionBody,
     FileUploadSchema,
+    LinkType,
     PresignedLink,
+    UploadedPart,
 )
 from models_library.basic_types import SHA256Str
 from models_library.generics import Envelope
 from models_library.rest_pagination import PageLimitInt, PageOffsetInt
 from pydantic import AnyUrl
 from settings_library.tracing import TracingSettings
-from starlette.datastructures import URL
+from simcore_service_api_server.exceptions.backend_errors import BackendTimeoutError
+from simcore_service_api_server.models.schemas.files import UserFile
+from simcore_service_api_server.models.schemas.jobs import UserFileToProgramJob
+from tenacity import (
+    AsyncRetrying,
+    TryAgain,
+    before_sleep_log,
+    retry_if_exception_type,
+    stop_after_delay,
+    wait_fixed,
+)
 
 from ..core.settings import StorageSettings
 from ..exceptions.service_errors_utils import service_exception_mapper
 from ..models.domain.files import File
 from ..utils.client_base import BaseServiceClientApi, setup_client_instance
 
+_POLL_TIMEOUT: Final[timedelta] = timedelta(minutes=10)
+
+
 _logger = logging.getLogger(__name__)
 
 _exception_mapper = partial(service_exception_mapper, service_name="Storage")
@@ -157,41 +179,81 @@ async def delete_file(self, *, user_id: int, quoted_storage_file_id: str) -> Non
         response.raise_for_status()
 
     @_exception_mapper(http_status_map={})
-    async def get_upload_links(
-        self, *, user_id: int, file_id: UUID, file_name: str
+    async def get_file_upload_links(
+        self, *, user_id: int, file: File, client_file: UserFileToProgramJob | UserFile
     ) -> FileUploadSchema:
-        object_path = urllib.parse.quote_plus(f"api/{file_id}/{file_name}")
+
+        query_params = QueryParams(
+            user_id=f"{user_id}",
+            link_type=LinkType.PRESIGNED.value,
+            file_size=int(client_file.filesize),
+            is_directory="false",
+            sha256_checksum=f"{client_file.sha256_checksum}",
+        )
 
         # complete_upload_file
         response = await self.client.put(
-            f"/locations/{self.SIMCORE_S3_ID}/files/{object_path}",
-            params={"user_id": user_id, "file_size": 0},
+            f"/locations/{self.SIMCORE_S3_ID}/files/{file.quoted_storage_file_id}",
+            params=query_params,
         )
         response.raise_for_status()
 
         enveloped_data = Envelope[FileUploadSchema].model_validate_json(response.text)
         assert enveloped_data.data  # nosec
         return enveloped_data.data
 
-    async def create_complete_upload_link(
-        self, *, file: File, query: dict[str, str] | None = None
-    ) -> URL:
-        url = URL(
-            f"{self.client.base_url}locations/{self.SIMCORE_S3_ID}/files/{file.quoted_storage_file_id}:complete"
+    @_exception_mapper(http_status_map={})
+    async def complete_file_upload(
+        self, *, user_id: int, file: File, uploaded_parts: list[UploadedPart]
+    ) -> ETag:
+
+        response = await self.client.post(
+            f"/locations/{self.SIMCORE_S3_ID}/files/{file.storage_file_id}:complete",
+            params={"user_id": f"{user_id}"},
+            json=jsonable_encoder(FileUploadCompletionBody(parts=uploaded_parts)),
         )
-        if query is not None:
-            url = url.include_query_params(**query)
-        return url
-
-    async def create_abort_upload_link(
-        self, *, file: File, query: dict[str, str] | None = None
-    ) -> URL:
-        url = URL(
-            f"{self.client.base_url}locations/{self.SIMCORE_S3_ID}/files/{file.quoted_storage_file_id}:abort"
+        response.raise_for_status()
+        file_upload_complete_response = Envelope[
+            FileUploadCompleteResponse
+        ].model_validate_json(response.text)
+        assert file_upload_complete_response.data  # nosec
+        state_url = f"{file_upload_complete_response.data.links.state}"
+        try:
+            async for attempt in AsyncRetrying(
+                reraise=True,
+                wait=wait_fixed(1),
+                stop=stop_after_delay(_POLL_TIMEOUT),
+                retry=retry_if_exception_type(TryAgain),
+                before_sleep=before_sleep_log(_logger, logging.DEBUG),
+            ):
+                with attempt:
+                    resp = await self.client.post(state_url)
+                    resp.raise_for_status()
+                    future_enveloped = Envelope[
+                        FileUploadCompleteFutureResponse
+                    ].model_validate_json(resp.text)
+                    assert future_enveloped.data  # nosec
+                    if future_enveloped.data.state == FileUploadCompleteState.NOK:
+                        raise TryAgain()
+
+                    assert future_enveloped.data.e_tag  # nosec
+                    _logger.debug(
+                        "multipart upload completed in %s, received %s",
+                        attempt.retry_state.retry_object.statistics,
+                        f"{future_enveloped.data.e_tag=}",
+                    )
+                    return future_enveloped.data.e_tag
+        except TryAgain as exc:
+            raise BackendTimeoutError() from exc
+        raise BackendTimeoutError()
+
+    @_exception_mapper(http_status_map={})
+    async def abort_file_upload(self, *, user_id: int, file: File) -> None:
+        response = await self.client.post(
+            f"/locations/{self.SIMCORE_S3_ID}/files/{file.quoted_storage_file_id}:abort",
+            params={"user_id": f"{user_id}"},
         )
-        if query is not None:
-            url = url.include_query_params(**query)
-        return url
+        response.raise_for_status()
 
     @_exception_mapper(http_status_map={})
     async def create_soft_link(

services/api-server/tests/unit/conftest.py

@@ -5,12 +5,14 @@
 # pylint: disable=broad-exception-caught
 
 import json
+import re
 import subprocess
 from collections.abc import AsyncIterator, Callable, Iterator
 from copy import deepcopy
 from pathlib import Path
 from typing import Any
 from unittest.mock import MagicMock
+from urllib.parse import urlparse, urlunparse
 
 import aiohttp.test_utils
 import httpx
@@ -21,13 +23,19 @@
 from faker import Faker
 from fastapi import FastAPI, status
 from fastapi.encoders import jsonable_encoder
-from httpx import ASGITransport
+from httpx import ASGITransport, Request, Response
 from models_library.api_schemas_long_running_tasks.tasks import (
     TaskGet,
     TaskProgress,
     TaskStatus,
 )
-from models_library.api_schemas_storage.storage_schemas import HealthCheck
+from models_library.api_schemas_storage.storage_schemas import (
+    FileUploadCompleteFutureResponse,
+    FileUploadCompleteResponse,
+    FileUploadCompleteState,
+    FileUploadSchema,
+    HealthCheck,
+)
 from models_library.api_schemas_webserver.projects import ProjectGet
 from models_library.app_diagnostics import AppStatusCheck
 from models_library.generics import Envelope
@@ -422,6 +430,77 @@ def mocked_storage_rest_api_base(
             ).model_dump(mode="json"),
         )
 
+        assert (
+            openapi["paths"]["/v0/locations/{location_id}/files/{file_id}"]["put"][
+                "operationId"
+            ]
+            == "upload_file_v0_locations__location_id__files__file_id__put"
+        )
+        respx_mock.put(
+            re.compile(r"^http://[a-z\-_]*storage:[0-9]+/v0/locations/[0-9]+/files.+$"),
+            name="upload_file_v0_locations__location_id__files__file_id__put",
+        ).respond(
+            status.HTTP_200_OK,
+            json=Envelope[FileUploadSchema](
+                data=FileUploadSchema.model_json_schema()["examples"][0]
+            ).model_dump(mode="json"),
+        )
+
+        # Add mocks for completion and abort endpoints
+        def generate_future_link(request: Request, **kwargs):
+            parsed_url = urlparse(f"{request.url}")
+            stripped_url = urlunparse(
+                (parsed_url.scheme, parsed_url.netloc, parsed_url.path, "", "", "")
+            )
+
+            payload = FileUploadCompleteResponse.model_validate(
+                {
+                    "links": {
+                        "state": stripped_url
+                        + ":complete/futures/"
+                        + str(faker.uuid4())
+                    },
+                },
+            )
+            return Response(
+                status_code=status.HTTP_200_OK,
+                json=jsonable_encoder(
+                    Envelope[FileUploadCompleteResponse](data=payload)
+                ),
+            )
+
+        respx_mock.post(
+            re.compile(
+                r"^http://[a-z\-_]*storage:[0-9]+/v0/locations/[0-9]+/files/.+complete(?:\?.*)?$"
+            ),
+            name="complete_upload_file_v0_locations__location_id__files__file_id__complete_post",
+        ).side_effect = generate_future_link
+
+        respx_mock.post(
+            re.compile(
+                r"^http://[a-z\-_]*storage:[0-9]+/v0/locations/[0-9]+/files/.+complete/futures/.+"
+            )
+        ).respond(
+            status_code=status.HTTP_200_OK,
+            json=jsonable_encoder(
+                Envelope[FileUploadCompleteFutureResponse](
+                    data=FileUploadCompleteFutureResponse(
+                        state=FileUploadCompleteState.OK,
+                        e_tag="07d1c1a4-b073-4be7-b022-f405d90e99aa",
+                    )
+                )
+            ),
+        )
+
+        respx_mock.post(
+            re.compile(
+                r"^http://[a-z\-_]*storage:[0-9]+/v0/locations/[0-9]+/files/.+:abort(?:\?.*)?$"
+            ),
+            name="abort_upload_file_v0_locations__location_id__files__file_id__abort_post",
+        ).respond(
+            status.HTTP_204_NO_CONTENT,
+        )
+
         # SEE https://github.com/pcrespov/sandbox-python/blob/f650aad57aced304aac9d0ad56c00723d2274ad0/respx-lib/test_disable_mock.py
         if not services_mocks_enabled:
             respx_mock.stop()