Offload certificate key generation to worker threads

This commit is contained in:
2026-10-25 21:26:34 +00:00
parent 4c1bfa4e4d
commit 718d11b4c0
4 changed files with 117 additions and 24 deletions
+6 -11
View File
@@ -1,7 +1,6 @@
"""Key generation, detail, rename/delete and export.""" """Key generation, detail, rename/delete and export."""
from __future__ import annotations from __future__ import annotations
import asyncio
import base64 import base64
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING, Optional
@@ -18,6 +17,7 @@ from ..models import Certificate, Key, Request
from ..openapi import error_response, json_response, schema_ref from ..openapi import error_response, json_response, schema_ref
from ..service import ( from ..service import (
ValidationError, ValidationError,
generate_key as generate_private_key,
parse_key_generation, parse_key_generation,
parse_uuid, parse_uuid,
save_key, save_key,
@@ -30,9 +30,6 @@ from ..service import (
if TYPE_CHECKING: if TYPE_CHECKING:
from ..mixin import Kca from ..mixin import Kca
# Key generation is CPU-bound; this bounds concurrent RSA/DSA jobs.
_keygen_semaphore = asyncio.Semaphore(2)
async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str): async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str):
user = await authenticated_user(kca.oidc, ctx) user = await authenticated_user(kca.oidc, ctx)
@@ -85,13 +82,11 @@ def register(app: KayaApp, kca: "Kca") -> None:
await send_error(ctx, 400, str(exc)) await send_error(ctx, 400, str(exc))
return return
try: try:
async with _keygen_semaphore: generated = await generate_private_key(
generated = await asyncio.to_thread( params["key_type"],
key_core.generate_key, params["size"],
params["key_type"], params["curve"],
params["size"], )
params["curve"],
)
key = await save_key( key = await save_key(
vault, vault,
generated, generated,
@@ -1,7 +1,6 @@
"""Certificate signing requests: list, create, detail, export and transforms.""" """Certificate signing requests: list, create, detail, export and transforms."""
from __future__ import annotations from __future__ import annotations
import asyncio
import base64 import base64
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING, Optional
@@ -21,6 +20,7 @@ from ..models import Key, Request, Template
from ..openapi import error_response, json_response, schema_ref from ..openapi import error_response, json_response, schema_ref
from ..service import ( from ..service import (
ValidationError, ValidationError,
generate_key,
parse_extensions, parse_extensions,
parse_subject, parse_subject,
parse_uuid, parse_uuid,
@@ -36,8 +36,6 @@ from ..service import (
if TYPE_CHECKING: if TYPE_CHECKING:
from ..mixin import Kca from ..mixin import Kca
_keygen_semaphore = asyncio.Semaphore(2)
async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str): async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str):
user = await authenticated_user(kca.oidc, ctx) user = await authenticated_user(kca.oidc, ctx)
@@ -98,13 +96,11 @@ def register(app: KayaApp, kca: "Kca") -> None:
key_obj = await load_private_key(key_model) key_obj = await load_private_key(key_model)
elif isinstance(body.get("new_key"), dict): elif isinstance(body.get("new_key"), dict):
new_key = body["new_key"] new_key = body["new_key"]
async with _keygen_semaphore: generated = await generate_key(
generated = await asyncio.to_thread( str(new_key.get("key_type") or "rsa"),
key_core.generate_key, int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE),
str(new_key.get("key_type") or "rsa"), new_key.get("curve"),
int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE), )
new_key.get("curve"),
)
key_model = await save_key( key_model = await save_key(
vault, vault,
generated, generated,
+19 -3
View File
@@ -1,6 +1,7 @@
"""Business logic: validation, persistence, serialization and CA operations.""" """Business logic: validation, persistence, serialization and CA operations."""
from __future__ import annotations from __future__ import annotations
import asyncio
import hashlib import hashlib
import secrets import secrets
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
@@ -471,6 +472,20 @@ async def serialize_vault_detail(vault: Vault) -> Dict[str, Any]:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Key generation is CPU-bound; bound concurrent jobs across all routes.
_keygen_semaphore = asyncio.Semaphore(2)
async def generate_key(
key_type: str,
size: int = key_core.DEFAULT_KEY_SIZE,
curve: Optional[str] = None,
) -> key_core.PrivateKey:
"""Generate a private key without blocking the event loop."""
async with _keygen_semaphore:
return await asyncio.to_thread(key_core.generate_key, key_type, size, curve)
def parse_key_generation(payload: Mapping[str, Any]) -> Dict[str, Any]: def parse_key_generation(payload: Mapping[str, Any]) -> Dict[str, Any]:
key_type = str(payload.get("key_type") or payload.get("type") or "rsa").lower() key_type = str(payload.get("key_type") or payload.get("type") or "rsa").lower()
if key_type not in key_core.KEY_TYPES: if key_type not in key_core.KEY_TYPES:
@@ -708,7 +723,7 @@ async def issue_certificate(
key_obj = await load_private_key(key_model) key_obj = await load_private_key(key_model)
elif new_key_payload: elif new_key_payload:
try: try:
generated = key_core.generate_key( generated = await generate_key(
str(new_key_payload.get("key_type") or "rsa"), str(new_key_payload.get("key_type") or "rsa"),
int(new_key_payload.get("size") or key_core.DEFAULT_KEY_SIZE), int(new_key_payload.get("size") or key_core.DEFAULT_KEY_SIZE),
new_key_payload.get("curve"), new_key_payload.get("curve"),
@@ -873,13 +888,13 @@ async def preview_certificate(vault: Vault, payload: Mapping[str, Any]) -> Dict[
else: else:
new_key = payload.get("new_key") new_key = payload.get("new_key")
if isinstance(new_key, Mapping): if isinstance(new_key, Mapping):
generated = key_core.generate_key( generated = await generate_key(
str(new_key.get("key_type") or "ec"), str(new_key.get("key_type") or "ec"),
int(new_key.get("size") or 2048), int(new_key.get("size") or 2048),
new_key.get("curve") or "secp256r1", new_key.get("curve") or "secp256r1",
) )
else: else:
generated = key_core.generate_key("ec", curve="secp256r1") generated = await generate_key("ec", curve="secp256r1")
public_key = generated.public_key() public_key = generated.public_key()
extensions_spec = parse_extensions(payload.get("extensions")) extensions_spec = parse_extensions(payload.get("extensions"))
@@ -1446,6 +1461,7 @@ __all__ = [
"find_key", "find_key",
"find_request", "find_request",
"generate_crl", "generate_crl",
"generate_key",
"import_parsed_items", "import_parsed_items",
"iso", "iso",
"issue_certificate", "issue_certificate",
+86
View File
@@ -55,6 +55,92 @@ async def _cert(client, vault_id: str, key_id: str, name: str, ca: bool = False,
return response.json() 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): class RouteExtraTests(unittest.TestCase):
@async_test @async_test
async def test_csr_creation_and_signing(self) -> None: async def test_csr_creation_and_signing(self) -> None: