Fix Google Maven proxy startup race in Copilot setup workflow

tools/google_maven_proxy.py

@@ -9,6 +9,7 @@
 import subprocess
 import sys
 import tarfile
+import threading
 import tempfile
 import zipfile
 from dataclasses import dataclass
@@ -429,7 +430,12 @@ def _download_from_google_maven(self, request_path: str, artifact_path: Path) ->
 
 
 class MirrorHandler(BaseHTTPRequestHandler):
-    mirror_index: MavenMirrorIndex
+    cache_dir: Path
+    donors: tuple[Donor, ...]
+    _mirror_index: MavenMirrorIndex | None = None
+    _config_generation = 0
+    _mirror_index_initializing_generation: int | None = None
+    _mirror_index_condition = threading.Condition()
     EMPTY_JAR_BYTES = (
         b"PK\x05\x06\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
     )
@@ -463,7 +469,7 @@ def _serve(self, send_body: bool) -> None:
             return
 
         try:
-            artifact_path = self.mirror_index.ensure_artifact(request_path)
+            artifact_path = self._get_mirror_index().ensure_artifact(request_path)
         except Exception as exc:
             self._send_text(
                 HTTPStatus.BAD_GATEWAY,
@@ -531,6 +537,50 @@ def _synthetic_response(request_path: str) -> tuple[str, bytes] | None:
 
         return None
 
+    @classmethod
+    def _get_mirror_index(cls) -> MavenMirrorIndex:
+        while True:
+            with cls._mirror_index_condition:
+                generation = cls._config_generation
+                if cls._mirror_index is not None:
+                    return cls._mirror_index
+                if cls._mirror_index_initializing_generation == generation:
+                    while (
+                        cls._mirror_index is None
+                        and cls._mirror_index_initializing_generation == generation
+                        and cls._config_generation == generation
+                    ):
+                        cls._mirror_index_condition.wait()
+                    if cls._config_generation != generation:
+                        continue
+                    if cls._mirror_index is not None:
+                        return cls._mirror_index
+                cls._mirror_index_initializing_generation = generation
+                cache_dir = cls.cache_dir
+                donors = cls.donors
+
+            try:
+                mirror_index = MavenMirrorIndex(cache_dir, donors)
+            except Exception:
+                with cls._mirror_index_condition:
+                    if cls._mirror_index_initializing_generation == generation:
+                        cls._mirror_index_initializing_generation = None
+                        cls._mirror_index_condition.notify_all()
+                raise
+
+            with cls._mirror_index_condition:
+                if cls._config_generation != generation:
+                    if cls._mirror_index_initializing_generation == generation:
+                        cls._mirror_index_initializing_generation = None
+                        cls._mirror_index_condition.notify_all()
+                    continue
+                if cls._mirror_index is None:
+                    cls._mirror_index = mirror_index
+                if cls._mirror_index_initializing_generation == generation:
+                    cls._mirror_index_initializing_generation = None
+                    cls._mirror_index_condition.notify_all()
+                return cls._mirror_index
+
 
 def main() -> int:
     parser = argparse.ArgumentParser(
@@ -551,15 +601,20 @@ def main() -> int:
     args = parser.parse_args()
 
     donors = parse_donors(args.donors)
-    mirror_index = MavenMirrorIndex(args.cache_dir, donors)
     print(f"Serving AndroidSA Google Maven proxy on http://{args.host}:{args.port}/", flush=True)
     print(
         f"Set ANDROIDSA_GOOGLE_MAVEN_URL=http://{args.host}:{args.port}/ or pass "
         f"-Pandroidsa.google.maven.url=http://{args.host}:{args.port}/ to Gradle.",
         flush=True,
     )
 
-    MirrorHandler.mirror_index = mirror_index
+    with MirrorHandler._mirror_index_condition:
+        MirrorHandler.cache_dir = args.cache_dir
+        MirrorHandler.donors = donors
+        MirrorHandler._config_generation += 1
+        MirrorHandler._mirror_index = None
+        MirrorHandler._mirror_index_initializing_generation = None
+        MirrorHandler._mirror_index_condition.notify_all()
     server = ThreadingHTTPServer((args.host, args.port), MirrorHandler)
     try:
         server.serve_forever()

tools/google_maven_proxy_test.py

@@ -0,0 +1,186 @@
+import sys
+import threading
+import unittest
+from pathlib import Path
+from unittest import mock
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import google_maven_proxy as proxy
+
+
+class MirrorHandlerInitializationTest(unittest.TestCase):
+    def setUp(self) -> None:
+        self.handler = proxy.MirrorHandler
+        with self.handler._mirror_index_condition:
+            self.handler.cache_dir = Path("/tmp/androidsa-google-maven-proxy-test")
+            self.handler.donors = (proxy.Donor("owner", "repo", "ref"),)
+            self.handler._config_generation += 1
+            self.handler._mirror_index = None
+            self.handler._mirror_index_initializing_generation = None
+            self.handler._mirror_index_condition.notify_all()
+
+    def test_competing_requests_share_one_initialization(self) -> None:
+        build_started = threading.Event()
+        second_request_waiting = threading.Event()
+        release_build = threading.Event()
+        call_count = 0
+        call_count_lock = threading.Lock()
+        built_index = object()
+        results: list[object] = []
+        errors: list[BaseException] = []
+
+        def build_index(*_args: object, **_kwargs: object) -> object:
+            nonlocal call_count
+            with call_count_lock:
+                call_count += 1
+            build_started.set()
+            release_build.wait(timeout=5)
+            return built_index
+
+        def worker() -> None:
+            try:
+                results.append(self.handler._get_mirror_index())
+            except BaseException as exc:  # pragma: no cover - assertion collects unexpected failures
+                errors.append(exc)
+
+        original_wait = self.handler._mirror_index_condition.wait
+
+        def wait_with_signal(timeout: float | None = None) -> bool:
+            second_request_waiting.set()
+            return original_wait(timeout)
+
+        with (
+            mock.patch.object(proxy, "MavenMirrorIndex", side_effect=build_index),
+            mock.patch.object(self.handler._mirror_index_condition, "wait", side_effect=wait_with_signal),
+        ):
+            first = threading.Thread(target=worker)
+            second = threading.Thread(target=worker)
+
+            first.start()
+            self.assertTrue(build_started.wait(timeout=5))
+            second.start()
+            self.assertTrue(second_request_waiting.wait(timeout=5))
+            with call_count_lock:
+                self.assertEqual(call_count, 1)
+
+            release_build.set()
+            first.join(timeout=5)
+            second.join(timeout=5)
+
+        self.assertFalse(errors)
+        self.assertEqual(results, [built_index, built_index])
+
+    def test_waiting_requests_retry_after_initialization_failure(self) -> None:
+        build_started = threading.Event()
+        waiting_request_blocked = threading.Event()
+        release_failure = threading.Event()
+        second_build_started = threading.Event()
+        call_count = 0
+        call_count_lock = threading.Lock()
+        fresh_index = object()
+        results: list[object] = []
+        errors: list[BaseException] = []
+
+        def build_index(*_args: object, **_kwargs: object) -> object:
+            nonlocal call_count
+            with call_count_lock:
+                call_count += 1
+                current_call = call_count
+            if current_call == 1:
+                build_started.set()
+                release_failure.wait(timeout=5)
+                raise RuntimeError("boom")
+            second_build_started.set()
+            return fresh_index
+
+        def worker(store_errors: bool) -> None:
+            try:
+                results.append(self.handler._get_mirror_index())
+            except BaseException as exc:
+                if store_errors:
+                    errors.append(exc)
+                else:  # pragma: no cover - assertion collects unexpected failures
+                    raise
+
+        original_wait = self.handler._mirror_index_condition.wait
+
+        def wait_with_signal(timeout: float | None = None) -> bool:
+            waiting_request_blocked.set()
+            return original_wait(timeout)
+
+        with (
+            mock.patch.object(proxy, "MavenMirrorIndex", side_effect=build_index),
+            mock.patch.object(self.handler._mirror_index_condition, "wait", side_effect=wait_with_signal),
+        ):
+            first = threading.Thread(target=worker, args=(True,))
+            second = threading.Thread(target=worker, args=(False,))
+
+            first.start()
+            self.assertTrue(build_started.wait(timeout=5))
+            second.start()
+            self.assertTrue(waiting_request_blocked.wait(timeout=5))
+            self.assertFalse(second_build_started.is_set())
+
+            release_failure.set()
+            self.assertTrue(second_build_started.wait(timeout=5))
+            first.join(timeout=5)
+            second.join(timeout=5)
+
+        self.assertEqual(len(errors), 1)
+        self.assertIsInstance(errors[0], RuntimeError)
+        self.assertEqual(str(errors[0]), "boom")
+        self.assertEqual(results, [fresh_index])
+        with call_count_lock:
+            self.assertEqual(call_count, 2)
+
+    def test_stale_initialization_is_discarded_after_reconfiguration(self) -> None:
+        first_build_started = threading.Event()
+        release_first_build = threading.Event()
+        call_count = 0
+        call_count_lock = threading.Lock()
+        stale_index = object()
+        fresh_index = object()
+        results: list[object] = []
+
+        def build_index(cache_dir: Path, _donors: tuple[proxy.Donor, ...]) -> object:
+            nonlocal call_count
+            with call_count_lock:
+                call_count += 1
+                current_call = call_count
+            if current_call == 1:
+                first_build_started.set()
+                release_first_build.wait(timeout=5)
+                return stale_index
+            self.assertEqual(cache_dir, Path("/tmp/androidsa-google-maven-proxy-test-fresh"))
+            return fresh_index
+
+        def worker() -> None:
+            results.append(self.handler._get_mirror_index())
+
+        with mock.patch.object(proxy, "MavenMirrorIndex", side_effect=build_index):
+            first = threading.Thread(target=worker)
+            second = threading.Thread(target=worker)
+
+            first.start()
+            self.assertTrue(first_build_started.wait(timeout=5))
+            with self.handler._mirror_index_condition:
+                self.handler.cache_dir = Path("/tmp/androidsa-google-maven-proxy-test-fresh")
+                self.handler.donors = (proxy.Donor("owner", "repo", "fresh"),)
+                self.handler._config_generation += 1
+                self.handler._mirror_index = None
+                self.handler._mirror_index_initializing_generation = None
+                self.handler._mirror_index_condition.notify_all()
+
+            second.start()
+            release_first_build.set()
+            first.join(timeout=5)
+            second.join(timeout=5)
+
+        self.assertEqual(results, [fresh_index, fresh_index])
+        with call_count_lock:
+            self.assertEqual(call_count, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()