feat: copy progress + cancel mid copy

Author: DeRuina

swift_browser_ui/ui/api.py

@@ -972,6 +972,87 @@ async def swift_replicate_container(
     )
 
 
+async def swift_replicate_status(request: aiohttp.web.Request) -> aiohttp.web.Response:
+    """Proxy replication status from upload-runner back to UI."""
+    session = await aiohttp_session.get_session(request)
+    request.app["Log"].info(
+        "API call for replication status from "
+        f"{request.remote}, sess: {session} :: {time.ctime()}"
+    )
+
+    project = request.query.get("project") or request.query.get("runner_project")
+    if not project:
+        raise aiohttp.web.HTTPBadRequest(reason="Missing ?project=<projectId>")
+
+    runner_id = await open_upload_runner_session(request, project=project)
+
+    job_id = request.match_info["job_id"]
+    path = f"/replicate/status/{job_id}"
+    signature = await sign(3600, path)
+
+    client = request.app["api_client"]
+    url = f"{setd['upload_internal_endpoint']}{path}"
+
+    async with client.get(
+        url,
+        cookies={"RUNNER_SESSION_ID": runner_id},
+        params=signature,
+        ssl=ssl_context,
+    ) as upstream:
+        body = await upstream.read()
+
+        # aiohttp.web.Response(content_type=...) cannot include charset
+        ctype = upstream.headers.get("Content-Type", "application/json")
+        if ";" in ctype:
+            ctype = ctype.split(";", 1)[0].strip()
+
+        return aiohttp.web.Response(
+            status=upstream.status,
+            body=body,
+            content_type=ctype,
+        )
+
+
+async def swift_replicate_cancel(request: aiohttp.web.Request) -> aiohttp.web.Response:
+    """Proxy replication cancel from UI to upload-runner."""
+    session = await aiohttp_session.get_session(request)
+    request.app["Log"].info(
+        "API call for replication cancel from "
+        f"{request.remote}, sess: {session} :: {time.ctime()}"
+    )
+
+    project = request.query.get("project") or request.query.get("runner_project")
+    if not project:
+        raise aiohttp.web.HTTPBadRequest(reason="Missing ?project=<projectId>")
+
+    runner_id = await open_upload_runner_session(request, project=project)
+
+    job_id = request.match_info["job_id"]
+    path = f"/replicate/cancel/{job_id}"
+    signature = await sign(3600, path)
+
+    client = request.app["api_client"]
+    url = f"{setd['upload_internal_endpoint']}{path}"
+
+    async with client.post(
+        url,
+        cookies={"RUNNER_SESSION_ID": runner_id},
+        params=signature,
+        ssl=ssl_context,
+    ) as upstream:
+        body = await upstream.read()
+
+        ctype = upstream.headers.get("Content-Type", "application/json")
+        if ";" in ctype:
+            ctype = ctype.split(";", 1)[0].strip()
+
+        return aiohttp.web.Response(
+            status=upstream.status,
+            body=body,
+            content_type=ctype,
+        )
+
+
 async def get_upload_session(
     request: aiohttp.web.Request,
 ) -> aiohttp.web.Response:

swift_browser_ui/ui/server.py

@@ -41,7 +41,9 @@
     swift_list_containers,
     swift_list_objects,
     swift_put_object,
+    swift_replicate_cancel,
     swift_replicate_container,
+    swift_replicate_status,
     swift_update_container_metadata,
 )
 from swift_browser_ui.ui.discover import handle_discover
@@ -305,6 +307,8 @@ async def on_prepare(
             aiohttp.web.post(
                 "/replicate/{project}/{container}", swift_replicate_container
             ),
+            aiohttp.web.get("/replicate/status/{job_id}", swift_replicate_status),
+            aiohttp.web.post("/replicate/cancel/{job_id}", swift_replicate_cancel),
         ]
     )
 

swift_browser_ui/upload/api.py

@@ -5,6 +5,7 @@
 import os
 import time
 import typing
+import uuid
 
 import aiohttp.web
 import msgpack
@@ -98,9 +99,33 @@ async def handle_replicate_container(
     await replicator.a_ensure_container()
     await replicator.a_ensure_container(segmented=True)
 
-    asyncio.ensure_future(replicator.a_copy_from_container())
+    job_id = uuid.uuid4().hex[:12]
 
-    return aiohttp.web.Response(status=202)
+    # create job record
+    request.app["replication_jobs"][job_id] = {
+        "state": "running",  # running | finished | failed | cancelled
+        "done": 0,
+        "total": 0,
+        "error": "",
+        "cancel": False,
+    }
+
+    async def runner():
+        try:
+            await replicator.a_copy_from_container(job_id=job_id, app=request.app)
+            request.app["replication_jobs"][job_id]["state"] = "finished"
+        except asyncio.CancelledError:
+            request.app["replication_jobs"][job_id]["state"] = "cancelled"
+            return
+        except Exception as e:
+            request.app["replication_jobs"][job_id]["state"] = "failed"
+            request.app["replication_jobs"][job_id]["error"] = str(e)
+            return
+
+    task = asyncio.create_task(runner())
+    request.app["replication_jobs"][job_id]["task"] = task
+
+    return aiohttp.web.json_response({"job_id": job_id}, status=202)
 
 
 async def handle_replicate_object(request: aiohttp.web.Request) -> aiohttp.web.Response:
@@ -138,6 +163,38 @@ async def handle_replicate_object(request: aiohttp.web.Request) -> aiohttp.web.R
     return aiohttp.web.Response(status=202)
 
 
+async def handle_replicate_status(request):
+    """Handle a request for getting replication job status."""
+    job_id = request.match_info["job_id"]
+    job = request.app["replication_jobs"].get(job_id)
+    if not job:
+        raise aiohttp.web.HTTPNotFound(reason="Job not found")
+
+    return aiohttp.web.json_response(
+        {
+            "state": job["state"],
+            "done": job["done"],
+            "total": job["total"],
+            "error": job.get("error", ""),
+        }
+    )
+
+
+async def handle_replicate_cancel(request):
+    """Handle a request for cancelling a replication job."""
+    job_id = request.match_info["job_id"]
+    job = request.app["replication_jobs"].get(job_id)
+    if not job:
+        raise aiohttp.web.HTTPNotFound(reason="Job not found")
+
+    job["cancel"] = True
+    task = job.get("task")
+    if task and not task.done():
+        task.cancel()
+
+    return aiohttp.web.json_response({"ok": True})
+
+
 async def handle_post_object_chunk(request: aiohttp.web.Request) -> aiohttp.web.Response:
     """Handle a request for posting an object chunk."""
     if "from_object" in request.query.keys():

swift_browser_ui/upload/replicate.py

@@ -1,5 +1,6 @@
 """Container and object replication handlers using aiohttp."""
 
+import asyncio
 import logging
 import os
 import ssl
@@ -26,6 +27,11 @@
 REPL_TIMEOUT = 16384
 
 
+def _is_cancelled(app, job_id: str) -> bool:
+    job = app["replication_jobs"].get(job_id) if app else None
+    return (job is None) or job.get("cancel", False)
+
+
 class ObjectReplicationProxy:
     """A class for replicating objects."""
 
@@ -90,8 +96,11 @@ async def a_ensure_container(self, segmented: bool = False) -> None:
 
         LOGGER.info(f"Created container '{container}'.")
 
-    async def a_sync_object_segments(self, manifest: str) -> str:
+    async def a_sync_object_segments(self, manifest: str, job_id: str, app) -> str:
         """Get object segments."""
+        if _is_cancelled(app, job_id):
+            raise asyncio.CancelledError()
+
         async with self.client.get(
             common.generate_download_url(
                 self.source_host, container=manifest.split("/")[0]
@@ -123,6 +132,8 @@ def filter_with_prefix(segment: str) -> bool:
         LOGGER.debug(f"Got following segments: {segments}")
 
         for segment in segments:
+            if _is_cancelled(app, job_id):
+                raise asyncio.CancelledError()
             from_url = common.generate_download_url(
                 self.source_host, container=manifest.split("/")[0], object_name=segment
             )
@@ -152,6 +163,8 @@ def filter_with_prefix(segment: str) -> bool:
                         reason="ETag missing, maybe segments file empty"
                     )
 
+                if _is_cancelled(app, job_id):
+                    raise asyncio.CancelledError()
                 to_url = common.generate_download_url(
                     self.host, container=f"{self.container}_segments", object_name=segment
                 )
@@ -175,8 +188,10 @@ def filter_with_prefix(segment: str) -> bool:
         )
         return new_manifest
 
-    async def a_copy_object(self, object_name: str) -> None:
+    async def a_copy_object(self, object_name: str, job_id: str, app) -> None:
         """Copy an object from a location."""
+        if _is_cancelled(app, job_id):
+            raise asyncio.CancelledError()
         # Get the object stream handle
         async with self.client.get(
             common.generate_download_url(
@@ -216,6 +231,8 @@ async def a_copy_object(self, object_name: str) -> None:
                     raise aiohttp.web.HTTPUnprocessableEntity(
                         reason="ETag missing, maybe segments file empty"
                     )
+                if _is_cancelled(app, job_id):
+                    raise asyncio.CancelledError()
                 async with self.client.put(
                     common.generate_download_url(self.host, self.container, object_name),
                     data=resp_g.content.iter_chunked(65564),
@@ -235,13 +252,20 @@ async def a_copy_object(self, object_name: str) -> None:
                 # segmented upload
                 LOGGER.debug(f"Copying object {object_name} in segments.")
 
+                if _is_cancelled(app, job_id):
+                    raise asyncio.CancelledError()
                 manifest = await self.a_sync_object_segments(
-                    resp_g.headers["X-Object-Manifest"]
+                    resp_g.headers["X-Object-Manifest"],
+                    job_id=job_id,
+                    app=app,
                 )
 
                 LOGGER.debug("Uploading manifest")
                 # Add manifest headers
                 headers["X-Object-Manifest"] = manifest
+
+                if _is_cancelled(app, job_id):
+                    raise asyncio.CancelledError()
                 # Create manifest file
                 async with self.client.put(
                     common.generate_download_url(
@@ -262,44 +286,80 @@ async def a_copy_single_object(self, object_name: str) -> None:
         """Only copy a single object."""
         await self.a_copy_object(object_name)
 
-    async def a_get_container_page(self, marker: str = "") -> list[str]:
-        """Get a single page of objects from a container."""
+    async def a_get_container_count(self) -> int:
+        """Get total object count in a container."""
+        async with self.client.head(
+            common.generate_download_url(
+                self.source_host, container=self.source_container
+            ),
+            headers={"X-Auth-Token": self.token},
+            timeout=ClientTimeout(total=REPL_TIMEOUT),
+            ssl=ssl_context,
+        ) as resp:
+            if resp.status >= 400:
+                raise aiohttp.web.HTTPBadRequest(
+                    reason="Could not HEAD source container."
+                )
+
+            try:
+                return int(resp.headers.get("X-Container-Object-Count", "0"))
+            except ValueError:
+                return 0
+
+    async def a_get_container_page(
+        self, marker: str = "", limit: int = 10000
+    ) -> list[str]:
+        """Get one page of object names from a container."""
+        params = {"limit": str(limit)}
+        if marker:
+            params["marker"] = marker
+
         async with self.client.get(
             common.generate_download_url(
-                self.source_host,
-                container=self.source_container,
+                self.source_host, container=self.source_container
             ),
             headers={"X-Auth-Token": self.token},
-            params={"marker": marker} if marker else None,
+            params=params,
             timeout=ClientTimeout(total=REPL_TIMEOUT),
             ssl=ssl_context,
         ) as resp:
             if resp.status >= 400:
-                LOGGER.debug(f"Container fetch failed with status {resp.status}")
                 raise aiohttp.web.HTTPBadRequest(
                     reason="Could not fetch source container."
                 )
 
             if resp.status == 200:
-                ret = await resp.text()
-                return ret.rstrip().lstrip().split("\n")
+                text = await resp.text()
+                return [x for x in text.strip().split("\n") if x]
 
             return []
 
-    async def a_copy_from_container(self) -> None:
-        """Copy objects from a source container."""
+    async def a_copy_from_container(self, job_id: str, app) -> None:
+        """Copy objects from a source container with live progress + cancel."""
         LOGGER.debug(f"Fetching objects from container {self.source_container}")
-        container_url = common.generate_download_url(
-            self.source_host, container=self.source_container
-        )
-        LOGGER.debug(f"Container url: {container_url}")
 
-        # Page through all the objects in a container
-        to_copy: list[str] = []
-        page = await self.a_get_container_page()
-        while page:
-            to_copy = to_copy + page
-            page = await self.a_get_container_page(to_copy[-1])
+        job = app["replication_jobs"].get(job_id)
+        if job is not None:
+            job["total"] = await self.a_get_container_count()
+            job["done"] = 0
+
+        marker = ""
+
+        while True:
+            job = app["replication_jobs"].get(job_id)
+            if job is None or job.get("cancel"):
+                raise asyncio.CancelledError()
+
+            page = await self.a_get_container_page(marker=marker)
+            if not page:
+                break
+
+            for obj in page:
+                job = app["replication_jobs"].get(job_id)
+                if job is None or job.get("cancel"):
+                    raise asyncio.CancelledError()
+
+                await self.a_copy_object(obj, job_id=job_id, app=app)
 
-        for obj in to_copy:
-            await self.a_copy_object(obj)
+                job["done"] += 1
+            marker = page[-1]