Offload certificate key generation to worker threads
This commit is contained in:
@@ -1,7 +1,6 @@
|
||||
"""Key generation, detail, rename/delete and export."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
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 ..service import (
|
||||
ValidationError,
|
||||
generate_key as generate_private_key,
|
||||
parse_key_generation,
|
||||
parse_uuid,
|
||||
save_key,
|
||||
@@ -30,9 +30,6 @@ from ..service import (
|
||||
if TYPE_CHECKING:
|
||||
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):
|
||||
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))
|
||||
return
|
||||
try:
|
||||
async with _keygen_semaphore:
|
||||
generated = await asyncio.to_thread(
|
||||
key_core.generate_key,
|
||||
params["key_type"],
|
||||
params["size"],
|
||||
params["curve"],
|
||||
)
|
||||
generated = await generate_private_key(
|
||||
params["key_type"],
|
||||
params["size"],
|
||||
params["curve"],
|
||||
)
|
||||
key = await save_key(
|
||||
vault,
|
||||
generated,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
"""Certificate signing requests: list, create, detail, export and transforms."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from datetime import datetime, timezone
|
||||
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 ..service import (
|
||||
ValidationError,
|
||||
generate_key,
|
||||
parse_extensions,
|
||||
parse_subject,
|
||||
parse_uuid,
|
||||
@@ -36,8 +36,6 @@ from ..service import (
|
||||
if TYPE_CHECKING:
|
||||
from ..mixin import Kca
|
||||
|
||||
_keygen_semaphore = asyncio.Semaphore(2)
|
||||
|
||||
|
||||
async def _vault(kca: "Kca", ctx: HttpContext, vault_id: str):
|
||||
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)
|
||||
elif isinstance(body.get("new_key"), dict):
|
||||
new_key = body["new_key"]
|
||||
async with _keygen_semaphore:
|
||||
generated = await asyncio.to_thread(
|
||||
key_core.generate_key,
|
||||
str(new_key.get("key_type") or "rsa"),
|
||||
int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE),
|
||||
new_key.get("curve"),
|
||||
)
|
||||
generated = await generate_key(
|
||||
str(new_key.get("key_type") or "rsa"),
|
||||
int(new_key.get("size") or key_core.DEFAULT_KEY_SIZE),
|
||||
new_key.get("curve"),
|
||||
)
|
||||
key_model = await save_key(
|
||||
vault,
|
||||
generated,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Business logic: validation, persistence, serialization and CA operations."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import secrets
|
||||
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]:
|
||||
key_type = str(payload.get("key_type") or payload.get("type") or "rsa").lower()
|
||||
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)
|
||||
elif new_key_payload:
|
||||
try:
|
||||
generated = key_core.generate_key(
|
||||
generated = await generate_key(
|
||||
str(new_key_payload.get("key_type") or "rsa"),
|
||||
int(new_key_payload.get("size") or key_core.DEFAULT_KEY_SIZE),
|
||||
new_key_payload.get("curve"),
|
||||
@@ -873,13 +888,13 @@ async def preview_certificate(vault: Vault, payload: Mapping[str, Any]) -> Dict[
|
||||
else:
|
||||
new_key = payload.get("new_key")
|
||||
if isinstance(new_key, Mapping):
|
||||
generated = key_core.generate_key(
|
||||
generated = await generate_key(
|
||||
str(new_key.get("key_type") or "ec"),
|
||||
int(new_key.get("size") or 2048),
|
||||
new_key.get("curve") or "secp256r1",
|
||||
)
|
||||
else:
|
||||
generated = key_core.generate_key("ec", curve="secp256r1")
|
||||
generated = await generate_key("ec", curve="secp256r1")
|
||||
public_key = generated.public_key()
|
||||
|
||||
extensions_spec = parse_extensions(payload.get("extensions"))
|
||||
@@ -1446,6 +1461,7 @@ __all__ = [
|
||||
"find_key",
|
||||
"find_request",
|
||||
"generate_crl",
|
||||
"generate_key",
|
||||
"import_parsed_items",
|
||||
"iso",
|
||||
"issue_certificate",
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user