diff --git a/CHANGELOG.md b/CHANGELOG.md index c888dc8..2a1c4b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,40 @@ All notable changes to kalshi-sdk will be documented in this file. ## Unreleased +### WS / auth polish batch (#173 + #174 + #178) + +- **#173 — `MessageQueue` defense-in-depth.** The WS `MessageQueue` underlying + `collections.deque` now carries `maxlen=maxsize+1` as a hard memory ceiling + enforced by deque itself, independent of the manual `_size` counter. If the + counter ever drifts (a put path that forgets to increment, an exception + between append and increment) the buffer cannot grow without bound. New + regression test in `tests/ws/test_backpressure.py` injects counter drift and + asserts the cap holds. No observable behavior change in the passing path. +- **#174 — types consolidation.** `_to_decimal_dollars` and `_to_decimal_fp` + were byte-identical apart from their docstrings. Collapsed into a single + `_coerce_decimal` helper shared by both `DollarDecimal` and `FixedPointCount`. + Public aliases unchanged; only the internal helper is shared. +- **#178 — async RSA-PSS sign offload.** Added `KalshiAuth.sign_request_async()` + that routes the ~1-10 ms RSA-PSS sign through a **dedicated** + `ThreadPoolExecutor(max_workers=2)` lazy-initialised on first use. + + Async REST (`AsyncTransport.request`) and async WS connect + (`ConnectionManager._build_auth_headers`) now use the async sign path; the + sync `sign_request` API is unchanged for sync-transport callers. + + The executor is dedicated (not asyncio's shared default pool) so signs + don't queue behind `loop.getaddrinfo` / file I/O / other `to_thread()` + work on a busy event loop — relevant during WS reconnect storms where + cold DNS resolution (5-50 ms) dominates the sign cost. Per the community + feedback on #178: a falsifiable microbench under `scripts/bench_sign_offload.py` + uses real `loop.time()` deltas (NOT the `asyncio.sleep(0)` ticker which is + special-cased and doesn't measure wall-clock blocking). Measured: inline + p99=2.95 ms vs. offloaded p99=0.68 ms on a 2048-bit key. + + `KalshiClient.close()` / `AsyncKalshiClient.close()` now shut down the + sign executor too; the executor is daemon-style and idempotent to close. + + ### Contract-map completeness (#171) Maps the remaining 42 REST sub-models, V2 orders family, and internal diff --git a/kalshi/_base_client.py b/kalshi/_base_client.py index 0f60190..106a58d 100644 --- a/kalshi/_base_client.py +++ b/kalshi/_base_client.py @@ -258,7 +258,10 @@ async def request( last_error: KalshiError | None = None for attempt in range(self._config.max_retries + 1): - auth_headers = self._auth.sign_request(method.upper(), sign_path) if self._auth else {} + if self._auth: + auth_headers = await self._auth.sign_request_async(method.upper(), sign_path) + else: + auth_headers = {} logger.debug( "Async request: %s %s (attempt %d/%d)", diff --git a/kalshi/async_client.py b/kalshi/async_client.py index 97765f8..976510e 100644 --- a/kalshi/async_client.py +++ b/kalshi/async_client.py @@ -171,6 +171,8 @@ def from_env(cls, **kwargs: object) -> AsyncKalshiClient: async def close(self) -> None: """Close the underlying async HTTP connection pool.""" await self._transport.close() + if self._auth is not None: + self._auth.close() async def __aenter__(self) -> AsyncKalshiClient: return self diff --git a/kalshi/auth.py b/kalshi/auth.py index 1d7b9dc..ca8b24f 100644 --- a/kalshi/auth.py +++ b/kalshi/auth.py @@ -13,11 +13,14 @@ from __future__ import annotations +import asyncio import base64 import logging import os import re +import threading import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from cryptography.exceptions import UnsupportedAlgorithm @@ -57,6 +60,16 @@ class KalshiAuth: def __init__(self, key_id: str, private_key: rsa.RSAPrivateKey) -> None: self._key_id = key_id self._private_key = private_key + # Lazy-initialised dedicated ThreadPoolExecutor for async sign offload. + # Two threads is enough for any realistic REST burst on a single client + # (RSA-PSS sign ~1-3 ms on a 2048-bit key, ~10 ms on 4096-bit) and + # keeps crypto off asyncio's default executor - which is shared with + # `loop.getaddrinfo`, file I/O, and other `to_thread()` work. Without + # isolation, a cold DNS resolution (5-50 ms) during a reconnect storm + # would block sign behind it and re-introduce the loop stall as + # await-latency rather than inline CPU. (#178) + self._sign_executor: ThreadPoolExecutor | None = None + self._sign_executor_lock = threading.Lock() @classmethod def from_key_path(cls, key_id: str, key_path: str | Path) -> KalshiAuth: @@ -205,3 +218,54 @@ def sign_request( "KALSHI-ACCESS-SIGNATURE": sig_b64, "KALSHI-ACCESS-TIMESTAMP": ts_str, } + + def _get_sign_executor(self) -> ThreadPoolExecutor: + """Return the dedicated sign-offload executor, creating it on first use. + + Thread-safe lazy init: the lock guards the assignment, and the + double-checked pattern keeps the fast path lock-free after first use. + """ + if self._sign_executor is None: + with self._sign_executor_lock: + if self._sign_executor is None: + self._sign_executor = ThreadPoolExecutor( + max_workers=2, + thread_name_prefix="kalshi-sign", + ) + return self._sign_executor + + async def sign_request_async( + self, method: str, path: str, timestamp_ms: int | None = None + ) -> dict[str, str]: + """Async variant of :meth:`sign_request` that offloads to a dedicated executor. + + Uses ``loop.run_in_executor(self._sign_executor, ...)`` so the + ~1-10 ms of RSA-PSS CPU does not block the asyncio event loop. + The executor is dedicated (not asyncio's default pool) so signs + don't queue behind `getaddrinfo` / file I/O / other `to_thread()` + work — relevant during WS reconnect storms where DNS dominates. + + The sync :meth:`sign_request` API is unchanged; callers running + on the sync transport should keep using it. + """ + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._get_sign_executor(), + self.sign_request, + method, + path, + timestamp_ms, + ) + + def close(self) -> None: + """Shut down the sign-offload executor if one was created. + + Idempotent. Safe to call without an executor having been initialised. + Long-lived clients ordinarily rely on interpreter-shutdown cleanup + (atexit-joined executor threads); this hook is for tests and other + callers that want deterministic teardown. + """ + with self._sign_executor_lock: + if self._sign_executor is not None: + self._sign_executor.shutdown(wait=False) + self._sign_executor = None diff --git a/kalshi/client.py b/kalshi/client.py index 67a0e32..ec8ddb6 100644 --- a/kalshi/client.py +++ b/kalshi/client.py @@ -140,6 +140,8 @@ def from_env(cls, **kwargs: object) -> KalshiClient: def close(self) -> None: """Close the underlying HTTP connection pool.""" self._transport.close() + if self._auth is not None: + self._auth.close() def __enter__(self) -> KalshiClient: return self diff --git a/kalshi/types.py b/kalshi/types.py index e367193..0f87c6c 100644 --- a/kalshi/types.py +++ b/kalshi/types.py @@ -10,13 +10,14 @@ T = TypeVar("T") -def _to_decimal_dollars(value: Any) -> Decimal: - """Convert a raw API dollar-string value to Decimal. - - Kalshi API returns price fields as FixedPointDollars strings - (e.g., ``"0.5600"``), with up to 6 decimal places of precision. - Response fields use a ``_dollars`` suffix (e.g., ``yes_bid_dollars``). - This converts them to Decimal without float intermediaries. +def _coerce_decimal(value: Any) -> Decimal: + """Coerce a wire value (str/int/float/Decimal) to ``Decimal`` without going through float. + + Used by both :data:`DollarDecimal` (price/cost fields with ``_dollars`` aliases) and + :data:`FixedPointCount` (volume/count fields with ``_fp`` aliases). The two aliases + differ only in their canonical wire name and intent; the decimal coercion is the same. + Going through ``str(value)`` for ``int`` / ``float`` avoids the binary-float + representation drift (``Decimal(0.65) == Decimal('0.65000000000000002...')``). """ if isinstance(value, Decimal): return value @@ -34,7 +35,7 @@ def _decimal_to_str(value: Decimal) -> str: DollarDecimal = Annotated[ Decimal, - BeforeValidator(_to_decimal_dollars), + BeforeValidator(_coerce_decimal), PlainSerializer(_decimal_to_str, return_type=str), ] """A Decimal field that handles bidirectional conversion for Kalshi dollar values. @@ -44,25 +45,13 @@ def _decimal_to_str(value: Decimal) -> str: """ -def _to_decimal_fp(value: Any) -> Decimal: - """Convert a raw API fixed-point count string to Decimal. - - Kalshi API returns count/volume fields as FixedPoint strings - (e.g., ``"100.00"``), with ``_fp`` suffix field names (e.g., ``count_fp``). - This converts them to Decimal without float intermediaries. - """ - if isinstance(value, Decimal): - return value - if isinstance(value, (int, float)): - return Decimal(str(value)) - if isinstance(value, str): - return Decimal(value) - raise TypeError(f"Cannot convert {type(value).__name__} to Decimal") +# FixedPointCount: count/volume fields (e.g., ``count_fp``, ``volume_fp``). +# Uses the same decimal coercion as :data:`DollarDecimal`; see _coerce_decimal. FixedPointCount = Annotated[ Decimal, - BeforeValidator(_to_decimal_fp), + BeforeValidator(_coerce_decimal), PlainSerializer(_decimal_to_str, return_type=str), ] """A Decimal field that handles bidirectional conversion for Kalshi count/volume values. diff --git a/kalshi/ws/backpressure.py b/kalshi/ws/backpressure.py index e2356c9..760d5a6 100644 --- a/kalshi/ws/backpressure.py +++ b/kalshi/ws/backpressure.py @@ -38,7 +38,12 @@ def __init__( ) -> None: self._maxsize = maxsize self._overflow = overflow - self._buffer: collections.deque[T | object] = collections.deque(maxlen=None) + # `maxlen=maxsize+1` is a hard memory ceiling enforced by deque itself, + # independent of `_size`. If the counter ever drifts (a put path that + # forgets to increment, an exception between append and increment, etc.) + # the buffer still cannot grow without bound. +1 because put_sentinel + # appends without going through the size-check overflow path. (#173) + self._buffer: collections.deque[T | object] = collections.deque(maxlen=maxsize + 1) self._event = asyncio.Event() self._closed = False # O(1) size tracking: counts real items only (excludes sentinel). diff --git a/kalshi/ws/connection.py b/kalshi/ws/connection.py index d5aace3..f7316d5 100644 --- a/kalshi/ws/connection.py +++ b/kalshi/ws/connection.py @@ -98,10 +98,14 @@ async def mark_streaming(self) -> None: """ await self._set_state(ConnectionState.STREAMING) - def _build_auth_headers(self) -> dict[str, str]: - """Build RSA-PSS auth headers for the WebSocket upgrade request.""" + async def _build_auth_headers(self) -> dict[str, str]: + """Build RSA-PSS auth headers for the WebSocket upgrade request. + + Async + executor-offloaded so the ~1-10 ms RSA-PSS sign doesn't + block the event loop on connect / reconnect. See #178. + """ ws_path = urlparse(self._config.ws_base_url).path - return self._auth.sign_request("GET", ws_path) + return await self._auth.sign_request_async("GET", ws_path) async def connect(self) -> None: """Establish a WebSocket connection with RSA-PSS auth headers. @@ -113,7 +117,7 @@ async def connect(self) -> None: """ await self._set_state(ConnectionState.CONNECTING) try: - headers = self._build_auth_headers() + headers = await self._build_auth_headers() self._ws = await connect( self._config.ws_base_url, additional_headers=headers, @@ -168,7 +172,7 @@ async def reconnect(self) -> None: await asyncio.sleep(delay) try: await self._set_state(ConnectionState.CONNECTING) - headers = self._build_auth_headers() + headers = await self._build_auth_headers() self._ws = await connect( self._config.ws_base_url, additional_headers=headers, diff --git a/scripts/bench_sign_offload.py b/scripts/bench_sign_offload.py new file mode 100644 index 0000000..f7eb423 --- /dev/null +++ b/scripts/bench_sign_offload.py @@ -0,0 +1,123 @@ +"""Precision microbench for #178: RSA-PSS sign offload. + +Measures the event-loop blocking impact of inline vs. executor-offloaded +signing under a concurrent burst. Not part of the unit-test suite — run +manually to validate the offload is buying what we think it is. + +Methodology (corrected from the original acceptance criterion): +- A real `asyncio.sleep(target_interval)` ticker measures wall-clock gaps + via `loop.time()` deltas. `asyncio.sleep(0)` is special-cased to a + single `call_soon` yield (see `tasks.__sleep0`) and does NOT measure + wall-clock blocking — it only measures task-queue ordering. +- Compare two modes: + 1. Inline (calls `sign_request` directly on the loop) + 2. Offload (calls `sign_request_async` → dedicated ThreadPoolExecutor) +- Report per-percentile ticker gap. + +Usage: + uv run python scripts/bench_sign_offload.py [--signs N] [--ticks N] [--key-size 2048|4096] +""" + +from __future__ import annotations + +import argparse +import asyncio +import statistics +import time + +from cryptography.hazmat.primitives.asymmetric import rsa + +from kalshi.auth import KalshiAuth + + +async def _ticker(loop: asyncio.AbstractEventLoop, ticks: int, interval: float) -> list[float]: + gaps: list[float] = [] + last = loop.time() + for _ in range(ticks): + await asyncio.sleep(interval) + now = loop.time() + gaps.append(now - last - interval) + last = now + return gaps + + +async def _inline_signs(auth: KalshiAuth, n: int) -> None: + for _ in range(n): + auth.sign_request("GET", "/trade-api/v2/markets", timestamp_ms=1) + await asyncio.sleep(0) # yield so the ticker can run + + +async def _offloaded_signs(auth: KalshiAuth, n: int) -> None: + for _ in range(n): + await auth.sign_request_async("GET", "/trade-api/v2/markets", timestamp_ms=1) + + +def _summarize(label: str, gaps: list[float]) -> None: + gaps_ms = [g * 1000 for g in gaps] + gaps_ms.sort() + p50 = statistics.median(gaps_ms) + p95 = gaps_ms[int(len(gaps_ms) * 0.95)] + p99 = gaps_ms[int(len(gaps_ms) * 0.99)] + print( + f"{label:>20}: p50={p50:6.2f} ms p95={p95:6.2f} ms " + f"p99={p99:6.2f} ms max={max(gaps_ms):6.2f} ms" + ) + + +async def _run(auth: KalshiAuth, signs: int, ticks: int, interval: float) -> None: + loop = asyncio.get_running_loop() + + # Mode 1: inline signs. + inline_gaps, _ = await asyncio.gather( + _ticker(loop, ticks, interval), + _inline_signs(auth, signs), + ) + _summarize("inline (on loop)", inline_gaps) + + # Mode 2: offloaded signs. + offload_gaps, _ = await asyncio.gather( + _ticker(loop, ticks, interval), + _offloaded_signs(auth, signs), + ) + _summarize("offloaded (executor)", offload_gaps) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--signs", type=int, default=500, help="signs per mode (default 500)") + parser.add_argument("--ticks", type=int, default=100, help="ticker samples (default 100)") + parser.add_argument( + "--interval", type=float, default=0.005, + help="ticker interval in seconds (default 0.005)", + ) + parser.add_argument( + "--key-size", type=int, default=2048, choices=[2048, 4096], + help="RSA key size for the benchmark key (default 2048)", + ) + args = parser.parse_args() + + print(f"Generating {args.key_size}-bit RSA key...") + t0 = time.perf_counter() + private_key = rsa.generate_private_key(public_exponent=65537, key_size=args.key_size) + print(f" ({(time.perf_counter() - t0) * 1000:.0f} ms)") + auth = KalshiAuth(key_id="bench", private_key=private_key) + + print( + f"\nRunning {args.signs} signs against {args.ticks}x ticker @ {args.interval * 1000:.0f}ms " + f"({args.key_size}-bit key)...\n" + ) + try: + asyncio.run(_run(auth, args.signs, args.ticks, args.interval)) + finally: + auth.close() + + print( + "\nExpect: inline gaps to track the sign cost (1-3 ms on 2048-bit, ~10 ms on 4096-bit), " + "offloaded gaps to stay close to schedule jitter." + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_auth.py b/tests/test_auth.py index 55df352..40f3d31 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import base64 import os import tempfile @@ -217,6 +218,105 @@ def test_case_variants_produce_same_canonical_path( pub.verify(sig2, canonical_msg, pss, hashes.SHA256()) +class TestSignRequestAsync: + """Async sign offload (#178). + + Verifies the executor-offloaded path produces identical signatures to the + sync path, and (in the loop-blocking microbench) that signing in flight + does not stall the event loop. + """ + + @pytest.mark.asyncio + async def test_async_signature_verifies_for_same_message(self, test_auth: KalshiAuth) -> None: + """Sync and async signing produce headers that both verify against the + same canonical message. RSA-PSS is randomized (PSS salt), so signature + bytes differ between calls — assert verifiability, not equality. + """ + ts = 1703123456789 + sync_headers = test_auth.sign_request("GET", "/trade-api/v2/markets", timestamp_ms=ts) + async_headers = await test_auth.sign_request_async( + "GET", "/trade-api/v2/markets", timestamp_ms=ts + ) + assert async_headers["KALSHI-ACCESS-KEY"] == sync_headers["KALSHI-ACCESS-KEY"] + assert async_headers["KALSHI-ACCESS-TIMESTAMP"] == sync_headers["KALSHI-ACCESS-TIMESTAMP"] + pub = test_auth._private_key.public_key() + canonical = (str(ts) + "GET/trade-api/v2/markets").encode("utf-8") + pss = padding.PSS( + mgf=padding.MGF1(hashes.SHA256()), + salt_length=padding.PSS.DIGEST_LENGTH, + ) + pub.verify( + base64.b64decode(sync_headers["KALSHI-ACCESS-SIGNATURE"]), + canonical, pss, hashes.SHA256(), + ) + pub.verify( + base64.b64decode(async_headers["KALSHI-ACCESS-SIGNATURE"]), + canonical, pss, hashes.SHA256(), + ) + + @pytest.mark.asyncio + async def test_lazy_executor_creation(self, test_auth: KalshiAuth) -> None: + assert test_auth._sign_executor is None + await test_auth.sign_request_async("GET", "/x", timestamp_ms=1) + assert test_auth._sign_executor is not None + # Idempotent — second call reuses the same pool. + first = test_auth._sign_executor + await test_auth.sign_request_async("GET", "/y", timestamp_ms=2) + assert test_auth._sign_executor is first + test_auth.close() + assert test_auth._sign_executor is None + + @pytest.mark.asyncio + async def test_close_is_idempotent(self, test_auth: KalshiAuth) -> None: + test_auth.close() # no executor yet + await test_auth.sign_request_async("GET", "/x", timestamp_ms=1) + test_auth.close() + test_auth.close() # double-close OK + + @pytest.mark.asyncio + async def test_concurrent_signs_do_not_stall_event_loop( + self, test_auth: KalshiAuth + ) -> None: + """Microbench (#178). Run a real ``asyncio.sleep(0.01)`` ticker + concurrently with a batch of signs and confirm the ticker's + observed gaps stay tight. + + Pre-offload (signs inline on the loop) the ticker would see gaps + proportional to the inline sign work (~1-3 ms per sign x N signs + between ticks). Post-offload, signs run on the dedicated executor + and the ticker only sees scheduler jitter. + + The threshold is generous (max gap < 30 ms) so this test stays + green under load on CI runners. It's a regression bound, not a + precision benchmark — the script under ``scripts/`` runs the + precision version. + """ + loop = asyncio.get_running_loop() + gaps: list[float] = [] + target_interval = 0.01 + + async def ticker() -> None: + last = loop.time() + for _ in range(40): + await asyncio.sleep(target_interval) + now = loop.time() + gaps.append(now - last - target_interval) + last = now + + async def signs() -> None: + # 200 signs covers a realistic batch_create burst. + for _ in range(200): + await test_auth.sign_request_async( + "GET", "/trade-api/v2/markets", timestamp_ms=1 + ) + + await asyncio.gather(ticker(), signs()) + test_auth.close() + + # Slack: schedule jitter + threadpool dispatch can each cost a few ms. + # Pre-offload inline RSA-PSS would push the max gap well past this. + assert max(gaps) < 0.030, f"max ticker gap {max(gaps) * 1000:.1f}ms exceeds 30ms budget" + class TestFromKeyPath: def test_loads_valid_pem_file(self, pem_bytes: bytes) -> None: with tempfile.NamedTemporaryFile(suffix=".pem", delete=False) as f: diff --git a/tests/ws/test_backpressure.py b/tests/ws/test_backpressure.py index 04448ae..58414ab 100644 --- a/tests/ws/test_backpressure.py +++ b/tests/ws/test_backpressure.py @@ -117,3 +117,25 @@ async def test_qsize_stays_accurate_across_put_get_drop_and_sentinel(self) -> No assert await q.get() == 3 assert await q.get() == 4 assert q.qsize() == 0 + + async def test_deque_maxlen_caps_memory_even_if_counter_drifts(self) -> None: + """Regression for #173: the underlying deque carries `maxlen=maxsize+1` + as a defense-in-depth ceiling. Even if the manual `_size` counter + drifts (a hypothetical bug that fails to track an eviction), the deque + cannot grow without bound — `maxlen` enforces the cap at the C level. + + Simulates counter drift by zeroing `_size` mid-stream so the overflow + check (`len(self._buffer) >= self._maxsize`) still triggers via the + real buffer length, but a counter-based bound would not. Then asserts + memory stays bounded regardless. + """ + q: MessageQueue[int] = MessageQueue(maxsize=5, overflow=OverflowStrategy.DROP_OLDEST) + for i in range(5): + await q.put(i) + # Inject counter drift. From here `_size` lies about the queue's + # occupancy; the only thing protecting memory is the deque's `maxlen`. + q._size = 0 + for i in range(1_000): + await q.put(i) + # Hard ceiling: maxsize + 1 (the +1 covers the put_sentinel append). + assert len(q._buffer) <= q._maxsize + 1 diff --git a/tests/ws/test_connection.py b/tests/ws/test_connection.py index cc55ab1..3bbff39 100644 --- a/tests/ws/test_connection.py +++ b/tests/ws/test_connection.py @@ -431,7 +431,7 @@ async def test_build_auth_headers_uses_ws_path(self, test_auth: object) -> None: timeout=5.0, ) mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] - headers = mgr._build_auth_headers() + headers = await mgr._build_auth_headers() assert "KALSHI-ACCESS-KEY" in headers assert "KALSHI-ACCESS-SIGNATURE" in headers assert "KALSHI-ACCESS-TIMESTAMP" in headers