railtracks.observability

23@dataclass
24class Event:
25    event_type: str
26    scope_type: str
27    scope_id: str
28    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
29    stamp: datetime = field(default_factory=Timestamp.now)
30    parent_scope_id: str | None = None
31    payload: dict[str, Any] = field(default_factory=dict)

14class Timestamp:
15    """Namespace helper for constructing `Event.stamp`. The field itself is a plain tz-aware UTC datetime."""
16
17    # Doing it this way to make potential changes easier
18    @staticmethod
19    def now() -> datetime:
20        return datetime.now(timezone.utc)

 38class Observer:
 39    def __init__(self) -> None:
 40        self._writers: dict[str, _Entry] = {}
 41        self._drops: dict[str, int] = {}  # dropped events per writer
 42        self._running = False
 43        self._pending_writers: list[Writer] = []
 44        self._start_lock: asyncio.Lock = asyncio.Lock()
 45
 46    # async context manager support added for now, this will become more clear
 47    # once we move to integrating with the other modules
 48    async def __aenter__(self) -> Observer:
 49        await self.start()
 50        return self
 51
 52    async def __aexit__(self, exc_type, exc, tb) -> None:
 53        await self.shutdown()
 54
 55    def configure_writers(self, writers: list[Writer]) -> None:
 56        """Register writers in bulk, to be brought up on the next `start()`.
 57
 58        Must be called before `start()` — raises `RuntimeError` if the observer
 59        is already running. Replaces any previously configured pending writers.
 60        """
 61        if self._running:
 62            raise RuntimeError(
 63                "configure_writers must be called before start(); use register() to add writers after."
 64            )
 65        self._pending_writers = list(writers)
 66
 67    async def start(self) -> None:
 68        """Bring up the observer.
 69
 70        This is idempotent safe to call multiple times.
 71        """
 72        if self._running:
 73            return
 74        async with self._start_lock:
 75            if self._running:
 76                return
 77            for i, writer in enumerate(self._pending_writers):
 78                await self._register_impl(writer, f"writer-{i}")
 79            self._running = True
 80
 81    async def shutdown(self) -> None:
 82        if not self._running:
 83            return
 84        self._running = False
 85        for name in list(self._writers.keys()):
 86            await self._teardown(name)
 87
 88    async def register(
 89        self,
 90        writer: Writer,
 91        name: str,
 92        maxsize: int = 10_000,
 93        policy: QueuePolicy = QueuePolicy.DROP_OLDEST,
 94    ) -> None:
 95        """Register a writer on a running observer. Post-start only.
 96
 97        For a pre-start batch, use `configure_writers()` and let `start()`
 98        register them.
 99        """
100        if not self._running:
101            raise RuntimeError(
102                "Observer is not running; call start() first, or use configure_writers() "
103                "for pre-start batch registration."
104            )
105        await self._register_impl(writer, name, maxsize=maxsize, policy=policy)
106
107    async def _register_impl(
108        self,
109        writer: Writer,
110        name: str,
111        maxsize: int = 10_000,
112        policy: QueuePolicy = QueuePolicy.DROP_OLDEST,
113    ) -> None:
114        """Internal registration path used by both public `register` (post-start)
115        and `start()` (registering pending writers during startup). Skips the
116        `_running` check so `start()` can register pending writers before
117        flipping the flag."""
118        if name in self._writers:
119            raise ValueError(f"Writer {name!r} is already registered.")
120        await writer.start()
121        queue: asyncio.Queue[_QueueItem] = asyncio.Queue(
122            maxsize=maxsize
123        )  # Each writer has its own queue
124        task = asyncio.create_task(
125            self._consumer_loop(name, writer, queue),
126            name=f"observer-consumer:{name}",
127        )
128        self._writers[name] = _Entry(
129            writer=writer, queue=queue, task=task, policy=policy
130        )
131        self._drops[name] = 0
132
133    async def unregister(self, name: str) -> None:
134        if name not in self._writers:
135            raise KeyError(f"No writer registered as {name!r}.")
136        await self._teardown(name)
137
138    async def publish(self, event: Event) -> None:
139        """Fan the event out to every registered writer's queue.
140
141        `async` on this method is contract-enforcement, the body doesn't `await` anything.
142        requiring callers to be inside a coroutine means they're on the same running loop
143
144        Args:
145            event: The event to publish.
146        """
147        if not self._running:
148            raise RuntimeError("Observer is not running.")
149        for name, entry in self._writers.items():
150            try:
151                entry.queue.put_nowait(event)
152            except asyncio.QueueFull:
153                self._handle_full_queue(name, entry, event)
154
155    def _handle_full_queue(self, name: str, entry: _Entry, event: Event) -> None:
156        match entry.policy:
157            case QueuePolicy.DROP_OLDEST:
158                self._drop_oldest(name, entry, event)
159
160    def _drop_oldest(self, name: str, entry: _Entry, event: Event) -> None:
161        try:
162            entry.queue.get_nowait()
163        except asyncio.QueueEmpty:
164            pass
165        entry.queue.put_nowait(event)
166        self._drops[name] += 1
167        logger.warning(
168            "observability writer %r queue full; dropped oldest event "
169            "(policy=%s, total drops for this writer: %d)",
170            name,
171            entry.policy.value,
172            self._drops[name],
173        )
174
175    async def _teardown(self, name: str) -> None:
176        entry = self._writers.pop(name)
177        self._drops.pop(name, None)
178        _enqueue_end(entry.queue)
179        await entry.task
180        await entry.writer.shutdown()
181
182    async def _consumer_loop(
183        self, name: str, writer: Writer, queue: asyncio.Queue[_QueueItem]
184    ) -> None:
185        while True:
186            item = await queue.get()
187            if isinstance(item, _EndOfStream):
188                return
189            try:
190                await writer.write(item)
191            except Exception as exc:
192                logger.warning(
193                    "observability writer %r failed on event %s: %s",
194                    name,
195                    item.event_id,
196                    exc,
197                )

15class QueuePolicy(Enum):
16    """How a writer's queue behaves when it's full at publish time."""
17
18    DROP_OLDEST = "drop_oldest"

 9class Writer(Protocol):
10    async def start(self) -> None: ...
11    async def write(self, event: Event) -> None: ...
12    async def shutdown(self) -> None: ...

12class JsonlWriter:
13    def __init__(self, directory: Path):
14        self._directory = directory
15        self._files: dict[str, TextIO] = {}
16
17    async def start(self) -> None:
18        self._directory.mkdir(parents=True, exist_ok=True)
19
20    async def write(self, event: Event) -> None:
21        handle = self._files.get(event.scope_id)
22        if handle is None:
23            _check_safe_scope_id(event.scope_id)
24            handle = (self._directory / f"{event.scope_id}.jsonl").open(
25                "a", encoding="utf-8"
26            )
27            self._files[event.scope_id] = handle
28        handle.write(_serialize(event) + "\n")
29        handle.flush()
30
31    async def shutdown(self) -> None:
32        for handle in self._files.values():
33            handle.flush()
34            handle.close()
35        self._files.clear()

12def configure_writers(writers: list[Writer]) -> None:
13    """Set the writers to register on the singleton Observer on first start().
14
15    Delegates to `observer.configure_writers`. Must be called before the
16    observer has started; raises `RuntimeError` otherwise.
17    """
18    observer.configure_writers(writers)

10async def publish_event(event: Event) -> None:
11    """Convenience wrapper to publish an Event via the process-wide singleton Observer."""
12    await configure.observer.publish(event)

21async def ensure_started() -> Observer:
22    """Start the singleton observer if not already started, return it."""
23    await observer.start()
24    return observer

27async def shutdown() -> None:
28    """Drain per-writer queues and stop the singleton Observer's consumer tasks.
29
30    Safe to call when the observer isn't running.
31    """
32    await observer.shutdown()