implement cancellation handling in CaptureService and Database, update job states accordingly

2026-09-05 06:54:10 +05:30
parent 9c1a301c1a
commit cab7a68b8b
7 changed files with 293 additions and 38 deletions
@@ -37,6 +37,38 @@ class MinimalSiteHandler(BaseHTTPRequestHandler):
        return


+class SlowResponseHandler(BaseHTTPRequestHandler):
+    slow_request_started = threading.Event()
+    release_slow_request = threading.Event()
+
+    def do_GET(self) -> None:  # noqa: N802
+        if self.path == "/":
+            body = b"<html><body><img src='/slow.bin'></body></html>"
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+            return
+        if self.path == "/slow.bin":
+            self.send_response(200)
+            self.send_header("Content-Type", "application/octet-stream")
+            self.send_header("Content-Length", "1")
+            self.end_headers()
+            self.slow_request_started.set()
+            self.release_slow_request.wait(timeout=30)
+            try:
+                self.wfile.write(b"x")
+            except BrokenPipeError:
+                pass
+            return
+        self.send_response(404)
+        self.end_headers()
+
+    def log_message(self, format: str, *args) -> None:  # type: ignore[override]
+        return
+
+
 def test_capture_service_creates_private_archive(tmp_path) -> None:
    server = ThreadingHTTPServer(("127.0.0.1", 0), MinimalSiteHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
@@ -89,3 +121,64 @@ def test_capture_service_creates_private_archive(tmp_path) -> None:
        names = archive.namelist()
    assert any(name.endswith("siteharbor-manifest.json") for name in names)
    assert any(name.endswith("site.css") for name in names)
+
+
+def test_capture_cancels_while_a_resource_is_stalled(tmp_path) -> None:
+    SlowResponseHandler.slow_request_started.clear()
+    SlowResponseHandler.release_slow_request.clear()
+    server = ThreadingHTTPServer(("127.0.0.1", 0), SlowResponseHandler)
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
+    port = server.server_address[1]
+
+    test_settings = replace(
+        settings,
+        data_dir=tmp_path,
+        database_path=tmp_path / "siteharbor.sqlite3",
+        work_dir=tmp_path / "work",
+        artifacts_dir=tmp_path / "artifacts",
+        reports_dir=tmp_path / "reports",
+        allow_private_networks=True,
+        allow_nonstandard_ports=True,
+        respect_robots=False,
+    )
+    test_settings.ensure_directories()
+    database = Database(test_settings.database_path)
+    database.initialize()
+    source_url = f"http://127.0.0.1:{port}/"
+    job = database.create_job(
+        source_url,
+        source_url,
+        {
+            "include_external_assets": False,
+            "max_pages": 10,
+            "max_depth": 2,
+            "max_bytes": 2_000_000,
+            "max_duration_seconds": 60,
+            "retention_days": 1,
+        },
+    )
+    claimed = database.claim_next("test-worker", lease_seconds=120)
+    assert claimed is not None
+    capture_thread = threading.Thread(
+        target=CaptureService(database, test_settings, "test-worker").run,
+        args=(claimed,),
+        daemon=True,
+    )
+    capture_thread.start()
+
+    try:
+        assert SlowResponseHandler.slow_request_started.wait(timeout=2)
+        cancelled = database.request_cancel(job["id"])
+        assert cancelled is not None
+        capture_thread.join(timeout=3)
+        assert not capture_thread.is_alive()
+    finally:
+        SlowResponseHandler.release_slow_request.set()
+        capture_thread.join(timeout=2)
+        server.shutdown()
+        server.server_close()
+
+    cancelled_job = database.get_job(job["id"])
+    assert cancelled_job is not None
+    assert cancelled_job["state"] == "cancelled"