Coverage for app/services/deployment

1"""In-process pub/sub bridge between the Celery event listener and the SSE endpoint.

3The Celery event listener runs in a daemon thread and ingests RabbitMQ

4events synchronously. Each SSE request produces an asyncio coroutine

5that awaits fresh events scoped to one ``deployment_id``.

7Each SSE coroutine subscribes and gets back an ``asyncio.Queue``; the

8listener thread pushes each event into every queue for that

9``deployment_id`` via ``loop.call_soon_threadsafe`` (the push happens

10from a non-asyncio thread).

12A small per-deployment ring buffer of recent events lets a client that

13connects mid-stream be backfilled with what already happened.

15Notes:

17* Subscribers are confined to this process — fine for one backend

18 replica, would fan out incorrectly across N replicas.

19* Queues are bounded; on overflow the oldest entry is dropped and a

20 synthetic ``{"type": "task-overflow", ...}`` is pushed so the

21 frontend can render a "you missed entries" banner.

22* The recent buffer is bounded per deployment and cleared when the

23 deployment reaches a terminal state.

24"""

26from __future__ import annotations

28import asyncio

29import contextlib

30import logging

31import threading

32from collections import defaultdict, deque

33from typing import Any

35logger = logging.getLogger(__name__)

37# Maximum number of buffered events per subscriber. If a slow consumer

38# stalls past this, the oldest events are dropped and the frontend is told.

39_QUEUE_MAXSIZE = 200

41# Per-deployment recent-events buffer, used to backfill clients that

42# connect mid-stream. Sized for a "what's been happening lately" view;

43# the full transcript lives in the task row's ``logs`` column.

44_RECENT_MAX = 500

47class DeploymentPubSub:

48 """One-process, in-memory fan-out keyed by ``deployment_id``.

50 Construct once at app start and reuse — there's only ever one

51 instance per backend process. ``set_loop`` is called from the

52 FastAPI lifespan handler so cross-thread pushes go to the right

53 event loop.

54 """

56 def __init__(self) -> None:

57 self._subs: dict[str, list[asyncio.Queue[dict[str, Any]]]] = defaultdict(list)

58 # Per-deployment ring buffer of recent events, for mid-stream

59 # subscribers. Reset when a deployment hits a terminal event.

60 self._recent: dict[str, deque[dict[str, Any]]] = defaultdict(

61 lambda: deque(maxlen=_RECENT_MAX)

62 )

63 # Protects both maps, touched by the listener thread (publish)

64 # and the asyncio loop (subscribe/unsubscribe).

65 self._lock = threading.Lock()

66 self._loop: asyncio.AbstractEventLoop | None = None

68 # ------------------------------------------------------------------

69 # Loop binding

70 # ------------------------------------------------------------------

72 def set_loop(self, loop: asyncio.AbstractEventLoop) -> None:

73 """Bind the FastAPI event loop.

75 Called once during lifespan startup. ``publish`` cannot do

76 anything meaningful until this is set — the listener thread

77 would have no loop to schedule the put onto.

78 """

79 self._loop = loop

81 # ------------------------------------------------------------------

82 # Subscriber API (called from the asyncio loop)

83 # ------------------------------------------------------------------

85 def subscribe(self, deployment_id: str) -> asyncio.Queue[dict[str, Any]]:

86 """Register a new queue for ``deployment_id`` and return it.

88 The caller is responsible for matching every ``subscribe`` with

89 an ``unsubscribe`` (use ``try/finally`` around the consumer

90 loop) — leaked queues sit in memory until the process restarts.

91 """

92 queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=_QUEUE_MAXSIZE)

93 with self._lock:

94 self._subs[deployment_id].append(queue)

95 logger.debug("pubsub: subscribed to %s (now %d subs)", deployment_id, len(self._subs[deployment_id]))

96 return queue

98 def recent(self, deployment_id: str) -> list[dict[str, Any]]:

99 """Snapshot the recent-events ring buffer for backfill.

100

101 Used by the SSE endpoint right after subscribing so a client

102 connecting mid-stream sees what happened in the last few

103 minutes instead of an empty live tail until the worker emits

104 its next line.

105 """

106 with self._lock:

107 buf = self._recent.get(deployment_id)

108 return list(buf) if buf else []

109

110 def unsubscribe(self, deployment_id: str, queue: asyncio.Queue[dict[str, Any]]) -> None:

111 with self._lock:

112 subs = self._subs.get(deployment_id)

113 if not subs:

114 return

115 with contextlib.suppress(ValueError):

116 subs.remove(queue)

117 if not subs:

118 self._subs.pop(deployment_id, None)

119 logger.debug("pubsub: unsubscribed from %s", deployment_id)

120

121 # ------------------------------------------------------------------

122 # Publisher API (called from the listener thread)

123 # ------------------------------------------------------------------

124

125 def publish(self, deployment_id: str, event: dict[str, Any]) -> None:

126 """Push ``event`` to every subscriber for ``deployment_id``.

127

128 Threadsafe. If the loop hasn't been bound yet, the event is

129 dropped from the live fan-out (no SSE endpoint can be open yet).

130

131 Also appends to the recent-events ring buffer for backfill.

132 Terminal lifecycle events

133 (``task-succeeded``/``task-failed``/``task-revoked``) clear the

134 buffer after the fan-out.

135 """

136 loop = self._loop

137 # Always append to the recent buffer, even if the loop isn't up

138 # yet, so later subscribers still benefit.

139 with self._lock:

140 self._recent[deployment_id].append(event)

141 queues = list(self._subs.get(deployment_id, ()))

142

143 if loop is not None:

144 for queue in queues:

145 loop.call_soon_threadsafe(self._enqueue_or_drop, queue, event)

146

147 # Reset the buffer on terminal events so a follow-up run starts

148 # clean. Done after the live fan-out so connected subscribers

149 # still see the terminal frame.

150 if event.get("type") in ("task-succeeded", "task-failed", "task-revoked"):

151 with self._lock:

152 self._recent.pop(deployment_id, None)

153

154 @staticmethod

155 def _enqueue_or_drop(queue: asyncio.Queue[dict[str, Any]], event: dict[str, Any]) -> None:

156 """Push onto the queue; on full queues drop oldest + signal overflow.

157

158 Runs inside the asyncio loop thread (scheduled via

159 ``call_soon_threadsafe``), so manipulating the queue is safe

160 without an extra lock.

161 """

162 try:

163 queue.put_nowait(event)

164 except asyncio.QueueFull:

165 with contextlib.suppress(asyncio.QueueEmpty):

166 queue.get_nowait()

167 queue.put_nowait(

168 {

169 "type": "task-overflow",

170 "message": "live stream backpressure: dropped older events",

171 }

172 )

173

174

175# Singleton — imported directly by listener and SSE endpoint.

176pubsub = DeploymentPubSub()

Coverage for app/services/deployment_pubsub.py: 41.38%

58 statements