Offload certificate key generation to worker threads

server/tests/test_routes_extra.py

@@ -55,6 +55,92 @@ async def _cert(client, vault_id: str, key_id: str, name: str, ca: bool = False,
     return response.json()
 
 
+class KeyGenerationTests(unittest.TestCase):
+    @async_test
+    async def test_shared_generation_is_offloaded_and_bounded(self) -> None:
+        import asyncio
+        import threading
+        from unittest.mock import patch
+
+        from kca import service
+
+        active = 0
+        max_active = 0
+        worker_threads: set[int] = set()
+        lock = threading.Lock()
+        barrier = threading.Barrier(2)
+
+        def fake_generate(*args: object) -> object:
+            nonlocal active, max_active
+            with lock:
+                active += 1
+                max_active = max(max_active, active)
+                worker_threads.add(threading.get_ident())
+            try:
+                barrier.wait(timeout=5)
+            finally:
+                with lock:
+                    active -= 1
+            return object()
+
+        with (
+            patch.object(service, "_keygen_semaphore", asyncio.Semaphore(2)),
+            patch.object(service.key_core, "generate_key", new=fake_generate),
+        ):
+            results = await asyncio.gather(*(service.generate_key("rsa", 8192) for _ in range(6)))
+
+        self.assertEqual(len(results), 6)
+        self.assertEqual(max_active, 2)
+        self.assertNotIn(threading.get_ident(), worker_threads)
+
+    @async_test
+    async def test_certificate_generation_is_offloaded(self) -> None:
+        import threading
+        from unittest.mock import patch
+
+        from kca.core import keys as key_core
+
+        event_loop_thread = threading.get_ident()
+        generation_threads: list[int] = []
+        generated_key = key_core.generate_key("rsa", 1024)
+
+        def fake_generate(*args: object) -> object:
+            generation_threads.append(threading.get_ident())
+            return generated_key
+
+        async with api_client() as client:
+            with oidc_user("alice"):
+                vault_id = await _vault(client)
+                with patch.object(key_core, "generate_key", new=fake_generate):
+                    response = await client.post(
+                        f"/api/vaults/{vault_id}/certificates/preview",
+                        json={
+                            "subject": _subject("preview"),
+                            "new_key": {"key_type": "rsa", "size": 8192},
+                        },
+                    )
+                    assert response.status_code == 200, response.text
+
+                    response = await client.post(
+                        f"/api/vaults/{vault_id}/certificates/preview",
+                        json={"subject": _subject("default preview")},
+                    )
+                    assert response.status_code == 200, response.text
+
+                    response = await client.post(
+                        f"/api/vaults/{vault_id}/certificates",
+                        json={
+                            "name": "issued",
+                            "subject": _subject("issued"),
+                            "new_key": {"key_type": "rsa", "size": 8192, "name": "issued key"},
+                        },
+                    )
+                    assert response.status_code == 201, response.text
+
+        self.assertEqual(len(generation_threads), 3)
+        self.assertNotIn(event_loop_thread, generation_threads)
+
+
 class RouteExtraTests(unittest.TestCase):
     @async_test
     async def test_csr_creation_and_signing(self) -> None: