Split into files

Author: Denis-Averin

scripts/check-urls.py

@@ -1,17 +1,18 @@
 import contextlib
 import fileinput
+import signal
 import os
 import re
-import subprocess
 import sys
 import threading
-import time
 import typing
 import urllib.parse
-from queue import Queue, Empty
 
 from github_job_summary import JobSummary
 from subdomains import Subdomains
+from curl_wrapper import EXIT_CODES as CURL_EXIT_CODES
+from curl_wrapper import CurlWrapper
+from url_checker import UrlChecker
 
 """
 Read file names from stdin (feed from git ls-files)
@@ -20,33 +21,16 @@
 Check them with CURL
 """
 
-# To avoid 403 responses
-USER_AGENT = "Googlebot/2.1 (+http://www.google.com/bot.html)"
-
-CONNECT_TIMEOUT_SEC = 5
-MAX_TIME_SEC = 10
 JOIN_TIMEOUT_SEC = 120
 
-
-class Curl:
-    """
-    See: https://curl.se/libcurl/c/libcurl-errors.html
-    """
-
-    CURL_STDERR_HTTP_RE = re.compile(r"^curl: \(22\) The requested URL returned error: (?P<http_code>\d+)")
-    OK = 0
-    COULDNT_RESOLVE_HOST = 6
-    HTTP_RETURNED_ERROR = 22
-
-
 CURL_EXIT_CODES_AND_HTTP_CODES = {
-    "https://api.aspose.cloud/connect/token": (Curl.HTTP_RETURNED_ERROR, 400),
-    "https://api.aspose.cloud/v3.0": (Curl.HTTP_RETURNED_ERROR, 404),
-    "https://api.aspose.cloud/v4.0": (Curl.HTTP_RETURNED_ERROR, 404),
-    "https://api.aspose.cloud/v4.0/": (Curl.HTTP_RETURNED_ERROR, 404),
-    "https://id.aspose.cloud/connect/token": (Curl.HTTP_RETURNED_ERROR, 400),
+    "https://api.aspose.cloud/connect/token": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 400),
+    "https://api.aspose.cloud/v3.0": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 404),
+    "https://api.aspose.cloud/v4.0": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 404),
+    "https://api.aspose.cloud/v4.0/": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 404),
+    "https://id.aspose.cloud/connect/token": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 400),
     # TODO: Temporary fix
-    "https://dashboard.aspose.cloud/applications": (Curl.HTTP_RETURNED_ERROR, 404),
+    "https://dashboard.aspose.cloud/applications": (CURL_EXIT_CODES.HTTP_RETURNED_ERROR, 404),
 }
 
 REGEX_TO_IGNORE: list[re.Pattern[str]] = [
@@ -170,69 +154,17 @@ def text_extractor(files: list[str]) -> typing.Generator[tuple[str, str], None,
                     raise
 
 
-class Task:
-    _proc: subprocess.Popen[bytes]
-    _stderr: str | None
-
-    def __init__(self, url: str):
-        self.url = url
-        self._proc = subprocess.Popen(
-            [
-                "curl",
-                "-sSf",
-                "--output",
-                "-",
-                "--connect-timeout",
-                str(CONNECT_TIMEOUT_SEC),
-                "--max-time",
-                str(MAX_TIME_SEC),
-                "--user-agent",
-                USER_AGENT,
-                self.url,
-            ],
-            stdout=open(os.devnull, "w"),
-            stderr=subprocess.PIPE,
-        )
-        self._stderr = None
-        self._started = time.time()
-
-    @property
-    def running(self) -> bool:
-        return self._proc.poll() is None
-
-    @property
-    def ret_code(self) -> int:
-        assert not self.running
-        return self._proc.returncode
-
-    @property
-    def stderr(self) -> str:
-        assert not self.running
-        if self._stderr is None:
-            self._stderr = self._proc.stderr.read().decode()
-        return self._stderr
-
-    @property
-    def age(self) -> float:
-        return time.time() - self._started
-
-
-def create_new_task(url: str) -> Task:
-    # print("Create task:", url)
-    return Task(url)
-
-
-def process_finished_task(task: Task) -> None:
+def process_finished_task(task) -> None:
     # print("Finish task:", task.url)
     expected_ret_code, expected_http_code = CURL_EXIT_CODES_AND_HTTP_CODES.get(task.url, (0, None))
     if task.ret_code == 0 or task.ret_code == expected_ret_code:
         print("OK:", "'%s' %.2fs" % (task.url, task.age))
         JOB_SUMMARY.add_success(task.url)
         return
 
-    if task.ret_code == Curl.HTTP_RETURNED_ERROR and expected_http_code:
+    if task.ret_code == CURL_EXIT_CODES.HTTP_RETURNED_ERROR and expected_http_code:
         # Try parse stderr for HTTP code
-        match = Curl.CURL_STDERR_HTTP_RE.match(task.stderr)
+        match = CurlWrapper.CURL_STDERR_HTTP_RE.match(task.stderr)
         assert match, "Unexpected output: %s" % task.stderr
         http_code = int(match.groupdict()["http_code"])
         if http_code == expected_http_code:
@@ -247,56 +179,31 @@ def process_finished_task(task: Task) -> None:
     JOB_SUMMARY.add_error(f"Broken URL '{task.url}': {task.stderr}Files: {EXTRACTED_URLS_WITH_FILES[task.url]}")
 
 
-WORKER_QUEUE: Queue[str | None] = Queue()
-
-
-def url_checker(num_workers: int = 8) -> None:
-    next_report_age_sec = 5
-    workers: list[Task | None] = [None for _ in range(num_workers)]
-
-    queue_is_empty = False
-
-    while not queue_is_empty or any(workers):
-        for i, task in enumerate(workers):
-            if task is None:
-                continue
-            if not task.running:
-                process_finished_task(task)
-                workers[i] = None
-            elif task.age > next_report_age_sec:
-                print("Long request: '%s' %.2fs" % (task.url, task.age))
-                next_report_age_sec += 3
-
-        if not queue_is_empty:
-            for i in (i for (i, w) in enumerate(workers) if w is None):
-                # Avoid blocking forever if the queue is currently empty
-                try:
-                    item = WORKER_QUEUE.get_nowait()
-                except Empty:
-                    break
-                if item is None:
-                    queue_is_empty = True
-                    print("--- url queue is over ---")
-                    break
-                url = item
-                workers[i] = create_new_task(url)
-        time.sleep(0.2)
-    print("Worker finished")
-
-
 JOB_SUMMARY = JobSummary(os.environ.get("GITHUB_STEP_SUMMARY", "step_summary.md"))
 JOB_SUMMARY.add_header("Test all URLs")
 
 
 def main(files: list[str]) -> int:
-    checker = threading.Thread(target=url_checker, daemon=True)
+    url_checker = UrlChecker(
+        on_finish=process_finished_task,
+    )
+
+    # Setup signal handlers for graceful shutdown
+    def _handle_signal(_sig: int, _frame: typing.Any) -> None:
+        url_checker.stop()
+
+    with contextlib.suppress(Exception):
+        signal.signal(signal.SIGINT, _handle_signal)
+        signal.signal(signal.SIGTERM, _handle_signal)
+
+    checker = threading.Thread(target=url_checker.run, daemon=True)
     checker.start()
 
     for filename, text in text_extractor(files):
         for url in url_extractor(text, filename):
             # print("In:", url)
-            WORKER_QUEUE.put_nowait(url)
-    WORKER_QUEUE.put_nowait(None)
+            url_checker.add_url(url)
+    url_checker.close()
     checker.join(timeout=JOIN_TIMEOUT_SEC)
     if checker.is_alive():
         print(

scripts/curl_wrapper.py

@@ -0,0 +1,86 @@
+import contextlib
+import os
+import re
+import subprocess
+import time
+from typing import Optional
+
+# To avoid 403 responses (default); caller may override per instance
+DEFAULT_USER_AGENT = "Googlebot/2.1 (+http://www.google.com/bot.html)"
+
+
+class EXIT_CODES:
+    OK = 0
+    COULDNT_RESOLVE_HOST = 6
+    HTTP_RETURNED_ERROR = 22
+
+
+class CurlWrapper:
+    """
+    Encapsulates a single curl execution with timeouts and helpers.
+    See: https://curl.se/libcurl/c/libcurl-errors.html
+    """
+
+    CURL_STDERR_HTTP_RE = re.compile(r"^curl: \(22\) The requested URL returned error: (?P<http_code>\d+)")
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        user_agent: str = DEFAULT_USER_AGENT,
+        connect_timeout: int = 5,
+        max_time: int = 10,
+    ) -> None:
+        self.url = url
+        self._stderr: Optional[str] = None
+        self._started = time.time()
+        self._proc = subprocess.Popen(
+            [
+                "curl",
+                "-sSf",
+                "--output",
+                "-",
+                "--connect-timeout",
+                str(connect_timeout),
+                "--max-time",
+                str(max_time),
+                "--user-agent",
+                user_agent,
+                self.url,
+            ],
+            stdout=open(os.devnull, "w"),
+            stderr=subprocess.PIPE,
+        )
+
+    @property
+    def running(self) -> bool:
+        return self._proc.poll() is None
+
+    @property
+    def ret_code(self) -> int:
+        assert not self.running
+        return self._proc.returncode
+
+    @property
+    def stderr(self) -> str:
+        assert not self.running
+        if self._stderr is None:
+            assert self._proc.stderr is not None
+            self._stderr = self._proc.stderr.read().decode()
+        return self._stderr
+
+    @property
+    def age(self) -> float:
+        return time.time() - self._started
+
+    def terminate(self, timeout: float | None = None) -> None:
+        try:
+            self._proc.terminate()
+            if timeout is not None:
+                self._proc.wait(timeout=timeout)
+        except Exception:
+            pass
+
+    def kill(self) -> None:
+        with contextlib.suppress(Exception):
+            self._proc.kill()

scripts/url_checker.py

@@ -0,0 +1,82 @@
+import contextlib
+import time
+from queue import Queue, Empty
+from typing import Callable, Optional
+
+from curl_wrapper import CurlWrapper
+
+
+class UrlChecker:
+    def __init__(
+        self,
+        *,
+        num_workers: int = 8,
+        hard_kill_sec: int = 15,
+        on_finish: Optional[Callable[[CurlWrapper], None]] = None,
+        worker_factory: Optional[Callable[[str], CurlWrapper]] = None,
+    ) -> None:
+        self.num_workers = num_workers
+        self.hard_kill_sec = hard_kill_sec
+        self.on_finish = on_finish
+        self.worker_factory = worker_factory or (lambda url: CurlWrapper(url))
+
+        self.queue: Queue[str | None] = Queue()
+        self.workers: list[CurlWrapper | None] = [None for _ in range(self.num_workers)]
+        self.stop_event = False
+        self.next_report_age_sec = 5
+
+    def add_url(self, url: str) -> None:
+        self.queue.put_nowait(url)
+
+    def close(self) -> None:
+        self.queue.put_nowait(None)
+
+    def stop(self) -> None:
+        self.stop_event = True
+        with contextlib.suppress(Exception):
+            self.queue.put_nowait(None)
+
+    def run(self) -> None:
+        queue_is_empty = False
+        while not queue_is_empty or any(self.workers):
+            # Graceful stop: cancel running curls
+            if self.stop_event:
+                queue_is_empty = True
+                for t in self.workers:
+                    if t is not None and t.running:
+                        t.terminate(timeout=1)
+                        if t.running:
+                            t.kill()
+
+            # Tick workers
+            for i, task in enumerate(self.workers):
+                if task is None:
+                    continue
+                if not task.running:
+                    if self.on_finish is not None:
+                        self.on_finish(task)
+                    self.workers[i] = None
+                elif task.age > self.next_report_age_sec:
+                    print("Long request: '%s' %.2fs" % (task.url, task.age))
+                    self.next_report_age_sec += 3
+                    if task.age > self.hard_kill_sec:
+                        task.terminate(timeout=2)
+                        if task.running:
+                            task.kill()
+                        print("Killed long request: '%s' %.2fs" % (task.url, task.age))
+
+            # Fill idle workers
+            if not queue_is_empty:
+                for i in (i for (i, w) in enumerate(self.workers) if w is None):
+                    try:
+                        item = self.queue.get_nowait()
+                    except Empty:
+                        break
+                    if item is None:
+                        queue_is_empty = True
+                        print("--- url queue is over ---")
+                        break
+                    url = item
+                    self.workers[i] = self.worker_factory(url)
+            time.sleep(0.2)
+        print("Worker finished")