stream: registre de connexions ZSET et livraison par XREAD bloquant

* stream:conns:{principal} : ZSET (membre = id de connexion, score = dernier
  heartbeat) au lieu d'un compteur INCR/DECR qui restait gonflé 1 h quand un worker
  mourait sans DECR (4029 injustifiés). Les membres muets depuis 90 s sont élagués à
  chaque comptage ; le score est rafraîchi toutes les 30 s ; ZREM à la déconnexion.
* Fan-out : un XREAD BLOCK (1 s) par socket abonnée sur filings:stream, via le client
  redis.asyncio, au lieu d'un XRANGE toutes les 250 ms par connexion. Le lecteur a sa
  propre connexion Redis fermée avec lui (une lecture bloquante annulée peut laisser une
  réponse tardive sur la connexion — doublons à la ré-abonnement) + événement d'arrêt.
* Plus aucun appel Redis synchrone dans la coroutine : ZSET/hello via le client async,
  accounting.charge (compteur sync) via asyncio.to_thread.
* Tests : élagage des connexions fantômes, rafraîchissement du score au heartbeat,
  livraison bloquante + unsubscribe/resume sans doublon. conftest.make_user : e-mail
  unique par uuid (id(counter) provoquait des collisions EMAIL_TAKEN).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Simon-Pierre Boucher committed 18 days ago (Sep 6, 2026) parent 6391e90

4 changed files +234 −47

modified hfmarketdata/api/stream/broker.py +21 −3

		`@@ -2,7 +2,8 @@`
2	2
3	3	* `publish(event)` — allocates a monotonic `seq` (INCR), appends the event to the Redis stream
4	4	`filings:stream` with id `<seq>-0` (MAXLEN 1000, exact) and publishes it on the `filings` channel.
5		−* `read_since(seq)` — entries with seq > `seq` (used for `resume_token` and for polling delivery).
	5	+* `read_since(seq)` — entries with seq > `seq` (one-shot, e.g. tests / tooling);
	6	+* `read_block(seq)` — same, waiting up to `block_ms` for new entries (`XREAD BLOCK`, what the WebSocket uses).
6	7
7	8	`HFMD_REDIS_URL=fakeredis://` swaps in fakeredis (one shared in-process server for sync and async clients),
8	9	which is what the tests use.
		`@@ -91,8 +92,12 @@ def current_seq() -> int:`
91	92	return int(v) if v else 0
92	93
93	94
94		−async def read_since(ar, seq: int, count: int = 200) -> list[dict[str, Any]]:
95		− entries = await ar.xrange(STREAM, min=f"{int(seq) + 1}-0", max="+", count=count)
	95	+async def current_seq_async(ar) -> int:
	96	+ v = await ar.get(SEQ_KEY)
	97	+ return int(v) if v else 0
	98	+
	99	+
	100	+def _decode(entries) -> list[dict[str, Any]]:
96	101	out = []
97	102	for _id, fields in entries:
98	103	data = fields.get("data") if isinstance(fields, dict) else None
		`@@ -101,6 +106,19 @@ async def read_since(ar, seq: int, count: int = 200) -> list[dict[str, Any]]:`
101	106	return out
102	107
103	108
	109	+async def read_since(ar, seq: int, count: int = 200) -> list[dict[str, Any]]:
	110	+ return _decode(await ar.xrange(STREAM, min=f"{int(seq) + 1}-0", max="+", count=count))
	111	+
	112	+
	113	+async def read_block(ar, seq: int, *, block_ms: int = 1000, count: int = 200) -> list[dict[str, Any]]:
	114	+ """Entries with seq > `seq`, waiting up to `block_ms` for new ones (`XREAD BLOCK`) — [] on timeout."""
	115	+ res = await ar.xread({STREAM: f"{int(seq)}-0"}, count=count, block=block_ms)
	116	+ out: list[dict[str, Any]] = []
	117	+ for _stream, entries in (res or []):
	118	+ out.extend(_decode(entries))
	119	+ return out
	120	+
	121	+
104	122	def read_since_sync(seq: int, count: int = 200) -> list[dict[str, Any]]:
105	123	entries = get_redis().xrange(STREAM, min=f"{int(seq) + 1}-0", max="+", count=count)
106	124	return [json.loads(f["data"]) for _id, f in entries if "data" in f]

modified hfmarketdata/api/stream/routes.py +150 −43

		`@@ -11,6 +11,12 @@ Protocol (AsyncAPI 3 spec in docs/asyncapi.yaml):`
11	11	* server → `{"type":"heartbeat","ts":…,"seq":…}` every 20 s; client `{"action":"ping"}` → `{"type":"pong"}`;
12	12	* each delivered `filing` message is charged 1 row through `stream.accounting.charge(principal, 1)`.
13	13
	14	+Plumbing: the principal comes from the rate-limit middleware (`scope["state"]`, see auth.py); live sockets per
	15	+principal are a Redis ZSET `stream:conns:{principal}` (member = connection id, score = last heartbeat; members
	16	+silent for 90 s are pruned, so a worker dying without cleanup never leaves phantom slots); delivery is one
	17	+blocking `XREAD` per subscribed socket on `filings:stream` (no polling). Every Redis call inside the coroutine
	18	+goes through the async client; the sync accounting counter runs in a worker thread.
	19	+
14	20	Author: Simon-Pierre Boucher <contact@spboucher.ai>
15	21	"""
16	22	from __future__ import annotations
		`@@ -19,6 +25,7 @@ import asyncio`
19	25	import json
20	26	import logging
21	27	import time
	28	+import uuid
22	29	from datetime import datetime, timezone
23	30	from typing import Any
24	31
		`@@ -37,9 +44,11 @@ log = logging.getLogger("hfmarketdata.stream")`
37	44	router = APIRouter(tags=["stream"])
38	45
39	46	HEARTBEAT_SECONDS = 20.0
40		−POLL_SECONDS = 0.25
41	47	MAX_CONNECTIONS_PER_KEY = 5
42		−CONN_TTL_SECONDS = 3600
	48	+CONN_TTL_SECONDS = 3600 # safety TTL of the ZSET itself
	49	+CONN_HEARTBEAT_SECONDS = 30.0 # how often a live socket refreshes its ZSET score
	50	+CONN_STALE_SECONDS = 90.0 # members older than this are pruned at every count
	51	+READ_BLOCK_MS = 1000 # XREAD BLOCK per subscribed socket
43	52	CLOSE_AUTH, CLOSE_LIMIT, CLOSE_PROTOCOL, CLOSE_UNAUTHENTICATED = 4001, 4029, 4400, 4401
44	53	CHANNELS = ("filings",)
45	54
		`@@ -116,64 +125,95 @@ async def stream_ws(ws: WebSocket, api_key: str \| None = Query(None)):`
116	125	await ws.close(code=CLOSE_AUTH, reason="AUTH_REQUIRED")
117	126	return
118	127	log.debug("stream connection %s tier=%s", principal, tier)
119		− r = broker.get_redis()
120		− conn_key = f"stream:conns:{principal}"
121		− n = int(r.incr(conn_key))
122		− r.expire(conn_key, CONN_TTL_SECONDS)
	128	+ ar = broker.get_async_redis()
	129	+ conn_id = uuid.uuid4().hex
	130	+ try:
	131	+ n = await register_connection(ar, principal, conn_id)
	132	+ except Exception as e: # Redis down: refuse rather than stream unaccounted
	133	+ log.warning("stream: redis unavailable (%s)", e)
	134	+ await ws.send_json(_err("INTERNAL_ERROR", "The stream is temporarily unavailable. Please retry."))
	135	+ await ws.close(code=CLOSE_PROTOCOL, reason="stream unavailable")
	136	+ return
123	137	if n > MAX_CONNECTIONS_PER_KEY:
124		− r.decr(conn_key)
	138	+ await release_connection(ar, principal, conn_id)
125	139	await ws.send_json(_err("STREAM_CONNECTION_LIMIT", CODES["STREAM_CONNECTION_LIMIT"], limit=MAX_CONNECTIONS_PER_KEY))
126	140	await ws.close(code=CLOSE_LIMIT, reason="STREAM_CONNECTION_LIMIT")
127	141	return
128		− ar = broker.get_async_redis()
129	142	sub = Subscription()
130		− last_seq = broker.current_seq()
	143	+ last_seq = await broker.current_seq_async(ar)
131	144	inbox: asyncio.Queue = asyncio.Queue()
	145	+ reader_task: asyncio.Task \| None = None
	146	+ reader_stop: asyncio.Event \| None = None
132	147
133	148	async def receiver():
134	149	try:
135	150	while True:
136	151	msg = await ws.receive_text()
137		− await inbox.put(msg)
138		− except WebSocketDisconnect:
139		− await inbox.put(None)
140		− except Exception:
141		− await inbox.put(None)
	152	+ await inbox.put(("msg", msg))
	153	+ except (WebSocketDisconnect, Exception):
	154	+ await inbox.put(("closed", None))
	155	+
	156	+ def stop_reader() -> None:
	157	+ nonlocal reader_task, reader_stop
	158	+ if reader_stop is not None:
	159	+ reader_stop.set()
	160	+ if reader_task is not None:
	161	+ reader_task.cancel()
	162	+ reader_task = reader_stop = None
	163	+
	164	+ def restart_reader(from_seq: int) -> None:
	165	+ nonlocal reader_task, reader_stop
	166	+ stop_reader()
	167	+ reader_stop = asyncio.Event()
	168	+ reader_task = asyncio.create_task(_reader(from_seq, inbox, reader_stop))
142	169
143	170	recv_task = asyncio.create_task(receiver())
144	171	try:
145	172	await ws.send_json({"type": "hello", "seq": last_seq, "heartbeat_seconds": HEARTBEAT_SECONDS,
146	173	"max_connections_per_key": MAX_CONNECTIONS_PER_KEY, "buffer": broker.MAXLEN, "ts": _now()})
147		− last_beat = time.monotonic()
	174	+ last_beat = last_touch = time.monotonic()
148	175	while True:
	176	+ wait = max(0.0, HEARTBEAT_SECONDS - (time.monotonic() - last_beat))
149	177	try:
150		− raw = await asyncio.wait_for(inbox.get(), timeout=POLL_SECONDS)
151		− if raw is None:
152		− break
153		− seq_box = [last_seq]
154		− await _handle(ws, raw, sub, ar, principal, seq_box)
155		− last_seq = seq_box[0]
	178	+ kind, payload = await asyncio.wait_for(inbox.get(), timeout=wait)
156	179	except asyncio.TimeoutError:
157		− pass
158		− if sub.channel:
159		− events = await broker.read_since(ar, last_seq)
160		− for ev in events:
161		− last_seq = max(last_seq, int(ev.get("seq", last_seq)))
162		− if sub.matches(ev):
163		− await ws.send_json(ev)
164		− accounting.charge(principal, accounting.REDUCED_RATE_ROWS_PER_MESSAGE)
165		− if time.monotonic() - last_beat >= HEARTBEAT_SECONDS:
166	180	await ws.send_json({"type": "heartbeat", "ts": _now(), "seq": last_seq})
167		− r.expire(conn_key, CONN_TTL_SECONDS)
168	181	last_beat = time.monotonic()
	182	+ if last_beat - last_touch >= CONN_HEARTBEAT_SECONDS:
	183	+ await touch_connection(ar, principal, conn_id)
	184	+ last_touch = last_beat
	185	+ continue
	186	+ if kind == "closed":
	187	+ break
	188	+ if kind == "event":
	189	+ ev = payload
	190	+ last_seq = max(last_seq, int(ev.get("seq", last_seq)))
	191	+ if sub.channel and sub.matches(ev):
	192	+ await ws.send_json(ev)
	193	+ await asyncio.to_thread(accounting.charge, principal, accounting.REDUCED_RATE_ROWS_PER_MESSAGE)
	194	+ continue
	195	+ seq_box = [last_seq]
	196	+ action = await _handle(ws, payload, sub, seq_box)
	197	+ if action == "subscribe":
	198	+ last_seq = seq_box[0]
	199	+ restart_reader(last_seq)
	200	+ elif action == "unsubscribe":
	201	+ stop_reader()
169	202	except WebSocketDisconnect:
170	203	pass
171	204	except Exception as e: # pragma: no cover
172	205	log.exception("stream error: %s", e)
173	206	finally:
174	207	recv_task.cancel()
	208	+ pending = reader_task
	209	+ stop_reader()
	210	+ if pending is not None:
	211	+ try:
	212	+ await pending
	213	+ except (asyncio.CancelledError, Exception):
	214	+ pass
175	215	try:
176		− r.decr(conn_key)
	216	+ await release_connection(ar, principal, conn_id)
177	217	except Exception:
178	218	pass
179	219	try:
		`@@ -187,30 +227,96 @@ async def stream_ws(ws: WebSocket, api_key: str \| None = Query(None)):`
187	227	pass
188	228
189	229
190		−async def _handle(ws: WebSocket, raw: str, sub: Subscription, ar, principal: str, seq_box: list[int]) -> None:
	230	+# ------------------------------------------------------------------------------------ connection registry
	231	+def conn_key(principal: str) -> str:
	232	+ return f"stream:conns:{principal}"
	233	+
	234	+
	235	+async def register_connection(ar, principal: str, conn_id: str) -> int:
	236	+ """Add the socket to the per-principal ZSET (score = heartbeat timestamp), prune members silent for more than
	237	+ `CONN_STALE_SECONDS` (worker killed without releasing) and return the live count — the 6th gets 4029."""
	238	+ key, now = conn_key(principal), time.time()
	239	+ await ar.zremrangebyscore(key, "-inf", now - CONN_STALE_SECONDS)
	240	+ await ar.zadd(key, {conn_id: now})
	241	+ await ar.expire(key, CONN_TTL_SECONDS)
	242	+ return int(await ar.zcard(key))
	243	+
	244	+
	245	+async def touch_connection(ar, principal: str, conn_id: str) -> None:
	246	+ await ar.zadd(key := conn_key(principal), {conn_id: time.time()})
	247	+ await ar.expire(key, CONN_TTL_SECONDS)
	248	+
	249	+
	250	+async def release_connection(ar, principal: str, conn_id: str) -> None:
	251	+ await ar.zrem(conn_key(principal), conn_id)
	252	+
	253	+
	254	+async def live_connections(ar, principal: str) -> int:
	255	+ """Live sockets of a principal (stale members pruned) — for tests / the admin console."""
	256	+ key = conn_key(principal)
	257	+ await ar.zremrangebyscore(key, "-inf", time.time() - CONN_STALE_SECONDS)
	258	+ return int(await ar.zcard(key))
	259	+
	260	+
	261	+# ------------------------------------------------------------------------------------ delivery
	262	+async def _reader(from_seq: int, inbox: asyncio.Queue, stop: asyncio.Event) -> None:
	263	+ """Blocking `XREAD` loop (one per subscribed socket, no polling): every entry after `from_seq` is queued
	264	+ to the connection loop, which filters and delivers. Redis hiccups back off one second and resume.
	265	+
	266	+ The reader owns a dedicated Redis connection, closed with it: a blocking read cancelled mid-flight can leave a
	267	+ late reply on its connection (fakeredis does, redis-py disconnects), and a shared connection would hand that
	268	+ reply to the next command — duplicates on resubscribe. `stop` is also checked on every turn."""
	269	+ last = int(from_seq)
	270	+ ar = broker.get_async_redis()
	271	+ try:
	272	+ while not stop.is_set():
	273	+ try:
	274	+ events = await broker.read_block(ar, last, block_ms=READ_BLOCK_MS)
	275	+ except asyncio.CancelledError:
	276	+ raise
	277	+ except Exception as e:
	278	+ if stop.is_set():
	279	+ return
	280	+ log.warning("stream reader: %s", e)
	281	+ await asyncio.sleep(1.0)
	282	+ continue
	283	+ if stop.is_set():
	284	+ return
	285	+ for ev in events:
	286	+ last = max(last, int(ev.get("seq", last)))
	287	+ await inbox.put(("event", ev))
	288	+ finally:
	289	+ try:
	290	+ await ar.aclose() if hasattr(ar, "aclose") else await ar.close()
	291	+ except Exception:
	292	+ pass
	293	+
	294	+
	295	+async def _handle(ws: WebSocket, raw: str, sub: Subscription, seq_box: list[int]) -> str \| None:
	296	+ """Process one client message; returns the action applied ('subscribe' \| 'unsubscribe' \| 'ping' \| None)."""
191	297	try:
192	298	msg = json.loads(raw)
193	299	except json.JSONDecodeError:
194	300	await ws.send_json(_err("VALIDATION_ERROR", "messages must be JSON objects"))
195		− return
	301	+ return None
196	302	if not isinstance(msg, dict):
197	303	await ws.send_json(_err("VALIDATION_ERROR", "messages must be JSON objects"))
198		− return
	304	+ return None
199	305	action = str(msg.get("action") or msg.get("type") or "").lower()
200	306	if action == "ping":
201	307	await ws.send_json({"type": "pong", "ts": _now(), "seq": seq_box[0]})
202		− return
	308	+ return "ping"
203	309	if action == "unsubscribe":
204	310	sub.channel = None
205	311	await ws.send_json({"type": "unsubscribed", "ts": _now()})
206		− return
	312	+ return "unsubscribe"
207	313	if action != "subscribe":
208	314	await ws.send_json(_err("VALIDATION_ERROR", f"unknown action '{action}' (subscribe, unsubscribe, ping)"))
209		− return
	315	+ return None
210	316	channel = str(msg.get("channel") or "filings").lower()
211	317	if channel not in CHANNELS:
212	318	await ws.send_json(_err("NOT_FOUND", f"unknown channel '{channel}'", channels=list(CHANNELS)))
213		− return
	319	+ return None
214	320	tickers = msg.get("tickers", "all")
215	321	forms = msg.get("forms")
216	322	if isinstance(tickers, str) and tickers.lower() == "all":
		`@@ -219,21 +325,22 @@ async def _handle(ws: WebSocket, raw: str, sub: Subscription, ar, principal: str`
219	325	sub.tickers = {t.upper() for t in tickers}
220	326	else:
221	327	await ws.send_json(_err("VALIDATION_ERROR", "tickers must be a non-empty list of strings or \"all\""))
222		− return
	328	+ return None
223	329	if forms in (None, "all"):
224	330	sub.forms = None
225	331	elif isinstance(forms, list) and all(isinstance(f, str) for f in forms) and forms:
226	332	sub.forms = {f.upper() for f in forms}
227	333	else:
228	334	await ws.send_json(_err("VALIDATION_ERROR", "forms must be a non-empty list of strings or omitted"))
229		− return
230		− sub.channel = channel
	335	+ return None
231	336	resume = msg.get("resume_token")
232	337	if resume is not None:
233	338	try:
234	339	seq_box[0] = int(resume)
235	340	except (TypeError, ValueError):
236	341	await ws.send_json(_err("VALIDATION_ERROR", "resume_token must be an integer (last seq received)"))
237		− return
	342	+ return None
	343	+ sub.channel = channel
238	344	await ws.send_json({"type": "subscribed", "channel": channel, "tickers": sorted(sub.tickers) if sub.tickers else "all",
239	345	"forms": sorted(sub.forms) if sub.forms else "all", "resume_from": seq_box[0], "ts": _now()})
	346	+ return "subscribe"

modified tests/conftest.py +2 −1

		`@@ -71,7 +71,8 @@ def make_user(app):`
71	71	def _make(email: str \| None = None, *, tier: str = "free", role: str = "user", password: str = "correct-horse-battery",
72	72	with_key: bool = True):
73	73	counter["n"] += 1
74		− email = email or f"user{counter['n']}-{os.getpid()}-{id(counter)}@example.com"
	74	+ import uuid
	75	+ email = email or f"user{counter['n']}-{os.getpid()}-{uuid.uuid4().hex[:10]}@example.com" # id(counter) collided
75	76	with session() as s:
76	77	u = service.create_user(s, email, "Test User", password=password, tier=tier, role=role, status="active", actor="test")
77	78	u.email_verified_at = service.now()

modified tests/test_stream.py +61 −0

		`@@ -196,6 +196,67 @@ def test_connection_limit_per_key(keys, client):`
196	196	assert ws.receive_json()["type"] == "hello"
197	197
198	198
	199	+def test_phantom_connections_are_pruned(keys, client):
	200	+ """A worker killed mid-stream never releases its slot: members whose heartbeat is older than 90 s must not
	201	+ count (the old INCR/DECR counter stayed inflated for an hour and produced unjustified 4029s)."""
	202	+ import time
	203	+
	204	+ from stream import broker, routes
	205	+ principal = f"key:{_key_id(keys[1])}"
	206	+ r = broker.get_redis()
	207	+ key = routes.conn_key(principal)
	208	+ r.delete(key)
	209	+ stale = time.time() - routes.CONN_STALE_SECONDS - 5
	210	+ r.zadd(key, {f"dead-{i}": stale for i in range(routes.MAX_CONNECTIONS_PER_KEY)}) # 5 phantom slots
	211	+ r.zadd(key, {"alive": time.time()}) # 1 real one elsewhere
	212	+ with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:
	213	+ assert ws.receive_json()["type"] == "hello" # phantoms pruned
	214	+ members = r.zrange(key, 0, -1)
	215	+ assert len(members) == 2 and "alive" in members and not any(m.startswith("dead-") for m in members)
	216	+ assert r.zrange(key, 0, -1) == ["alive"] # own slot released
	217	+ r.delete(key)
	218	+
	219	+
	220	+def test_heartbeat_refreshes_connection_score(keys, client, monkeypatch):
	221	+ import time
	222	+
	223	+ from stream import broker, routes
	224	+ monkeypatch.setattr(routes, "HEARTBEAT_SECONDS", 0.2)
	225	+ monkeypatch.setattr(routes, "CONN_HEARTBEAT_SECONDS", 0.0)
	226	+ principal = f"key:{_key_id(keys[1])}"
	227	+ r = broker.get_redis()
	228	+ with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:
	229	+ ws.receive_json()
	230	+ conn_id = r.zrange(routes.conn_key(principal), 0, -1)[0]
	231	+ s0 = r.zscore(routes.conn_key(principal), conn_id)
	232	+ assert ws.receive_json()["type"] == "heartbeat"
	233	+ assert ws.receive_json()["type"] == "heartbeat"
	234	+ assert r.zscore(routes.conn_key(principal), conn_id) >= s0 and time.time() - s0 < 5
	235	+
	236	+
	237	+def test_delivery_is_blocking_read_not_polling(keys, client):
	238	+ """Events published after subscribe arrive through XREAD BLOCK; events before a resume token are skipped;
	239	+ unsubscribe stops delivery (no reader), resubscribe with a token replays."""
	240	+ from stream import broker
	241	+ with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
	242	+ ws.receive_json()
	243	+ ws.send_json({"action": "subscribe", "tickers": ["NVDA"]})
	244	+ ws.receive_json()
	245	+ s1 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b1"})
	246	+ assert ws.receive_json()["accn"] == "b1"
	247	+ ws.send_json({"action": "unsubscribe"})
	248	+ assert ws.receive_json()["type"] == "unsubscribed"
	249	+ s2 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b2"})
	250	+ ws.send_json({"action": "ping"})
	251	+ assert ws.receive_json()["type"] == "pong" # nothing delivered while unsubscribed
	252	+ ws.send_json({"action": "subscribe", "tickers": ["NVDA"], "resume_token": s1})
	253	+ assert ws.receive_json()["resume_from"] == s1
	254	+ assert ws.receive_json()["accn"] == "b2" and s2 > s1
	255	+ s3 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b3"})
	256	+ got = ws.receive_json()
	257	+ assert got["accn"] == "b3" and got["seq"] == s3
	258	+
	259	+
199	260	def test_accounting_hook_interface(app):
200	261	from stream import accounting
201	262	calls = []
202	263