|
| 1 | +"""The one background event loop that runs every periodic observability flush. |
| 2 | +
|
| 3 | +Before this module each sink started its own: the log adapter a |
| 4 | +``new_event_loop`` + ``run_forever`` thread, the metrics and traces adapters a |
| 5 | +task on whatever loop was running or else an ``asyncio.run`` thread, the sizing |
| 6 | +sink a hand-made loop on a thread, and the Segment client a dedicated thread and |
| 7 | +loop it waited up to 5 s to come up. Up to five loops, each with its own |
| 8 | +cadence and its own (or no) shutdown, so no single bound covered the final |
| 9 | +flushes at exit. |
| 10 | +
|
| 11 | +Now one daemon thread (``atlan-observability-loop``) runs one loop, started on |
| 12 | +first use. Long-lived work is started with :meth:`ObservabilityLoop.start`, |
| 13 | +which returns a :class:`BackgroundTask` handle. :meth:`ObservabilityLoop.shutdown` |
| 14 | +runs once at interpreter exit (``atexit``) and stops everything within one |
| 15 | +bound, ``OBSERVABILITY_SHUTDOWN_TIMEOUT_SECONDS``: |
| 16 | +
|
| 17 | +1. started tasks are cancelled once each, in ascending ``stop_order`` stages: |
| 18 | + a stage is cancelled only after the one before it has finished. The |
| 19 | + periodic flushes (stage 0) make their final flush in their |
| 20 | + ``CancelledError`` handler, which for metrics drains the Segment queue; |
| 21 | + the Segment worker (stage 1) then sends its last batch and closes the HTTP |
| 22 | + client those drains used; |
| 23 | +2. every other task on the loop (a batch flush ``add_record`` started there) is |
| 24 | + awaited, never cancelled, because cancelling one deletes its half-written |
| 25 | + file and loses the batch; |
| 26 | +3. the loop stops and the thread is joined with what is left of the bound. |
| 27 | +
|
| 28 | +After ``fork()`` the child gets a fresh lock and no loop; the next |
| 29 | +:meth:`ObservabilityLoop.start` starts the child's own thread. Callbacks |
| 30 | +registered with :func:`register_after_fork` (the sinks restart their |
| 31 | +periodic flushes, the Segment client drops its worker) run in the child |
| 32 | +after that reset. |
| 33 | +
|
| 34 | +Two pieces of background work deliberately stay off this loop. The console |
| 35 | +writer (``console_sink``) is a plain thread because its whole job is a blocking |
| 36 | +``write()`` to stdout; on this loop a stalled pipe would stall every flush. The |
| 37 | +OTLP log exporter's batch thread belongs to the OpenTelemetry SDK. |
| 38 | +""" |
| 39 | + |
| 40 | +from __future__ import annotations |
| 41 | + |
| 42 | +import asyncio |
| 43 | +import atexit |
| 44 | +import concurrent.futures |
| 45 | +import logging |
| 46 | +import os |
| 47 | +import threading |
| 48 | +import time |
| 49 | +from collections.abc import Callable, Coroutine |
| 50 | +from typing import Any |
| 51 | + |
| 52 | +from application_sdk.constants import OBSERVABILITY_SHUTDOWN_TIMEOUT_SECONDS |
| 53 | + |
| 54 | +THREAD_NAME = "atlan-observability-loop" |
| 55 | + |
| 56 | + |
| 57 | +class BackgroundTask: |
| 58 | + """Handle to a long-lived coroutine running on the observability loop. |
| 59 | +
|
| 60 | + Thread-safe. :meth:`cancel` may be called from any thread, any number of |
| 61 | + times: the task is cancelled at most once, on the loop's own thread, so a |
| 62 | + second cancel can never interrupt the final flush the first one started. |
| 63 | + """ |
| 64 | + |
| 65 | + def __init__( |
| 66 | + self, name: str, loop: asyncio.AbstractEventLoop, stop_order: int = 0 |
| 67 | + ) -> None: |
| 68 | + self.name = name |
| 69 | + self.loop = loop |
| 70 | + self.stop_order = stop_order |
| 71 | + self._task: asyncio.Task[Any] | None = None |
| 72 | + # Set by the task's first step. Until then a cancel is deferred: a task |
| 73 | + # cancelled before its first step never enters the coroutine, so its |
| 74 | + # ``CancelledError`` handler (the final flush) would never run. |
| 75 | + self._started = False |
| 76 | + self._cancel_requested = False |
| 77 | + self._cancel_sent = False |
| 78 | + self._done = threading.Event() |
| 79 | + |
| 80 | + def done(self) -> bool: |
| 81 | + """Whether the coroutine has finished, or will never run.""" |
| 82 | + return self._done.is_set() |
| 83 | + |
| 84 | + def wait(self, timeout: float) -> bool: |
| 85 | + """Wait up to *timeout* seconds for the coroutine to finish. |
| 86 | +
|
| 87 | + Returns ``True`` if it finished. Never call from the loop's thread. |
| 88 | + """ |
| 89 | + return self._done.wait(timeout) |
| 90 | + |
| 91 | + def cancel(self) -> None: |
| 92 | + """Ask the coroutine to stop. Returns at once; use :meth:`wait` to wait.""" |
| 93 | + self._cancel_requested = True |
| 94 | + try: |
| 95 | + self.loop.call_soon_threadsafe(self._cancel_on_loop) |
| 96 | + except RuntimeError: |
| 97 | + # The loop is closed, so the coroutine can never run again. |
| 98 | + self._done.set() |
| 99 | + |
| 100 | + async def _run(self, coro: Coroutine[Any, Any, Any]) -> Any: |
| 101 | + """The task body: mark the start, then run *coro*.""" |
| 102 | + self._started = True |
| 103 | + return await coro |
| 104 | + |
| 105 | + def _attach(self, task: asyncio.Task[Any]) -> None: |
| 106 | + """Bind the running task. Runs on the loop's thread.""" |
| 107 | + self._task = task |
| 108 | + task.add_done_callback(lambda _t: self._done.set()) |
| 109 | + if self._cancel_requested: |
| 110 | + self._cancel_on_loop() |
| 111 | + |
| 112 | + def _cancel_on_loop(self) -> None: |
| 113 | + """Cancel the task once. Runs on the loop's thread only.""" |
| 114 | + if self._task is None or self._cancel_sent: |
| 115 | + return |
| 116 | + if not self._started: |
| 117 | + # Retry after the task's first step, which is already scheduled. |
| 118 | + self.loop.call_soon(self._cancel_on_loop) |
| 119 | + return |
| 120 | + self._cancel_sent = True |
| 121 | + self._task.cancel() |
| 122 | + |
| 123 | + |
| 124 | +class ObservabilityLoop: |
| 125 | + """A single daemon thread running a single asyncio event loop. |
| 126 | +
|
| 127 | + Started lazily on the first :meth:`start` or :meth:`submit`, so importing |
| 128 | + the SDK does not start a thread. Fork-aware: a child process, which |
| 129 | + inherits this object but not the thread, starts its own on first use. |
| 130 | + """ |
| 131 | + |
| 132 | + def __init__(self) -> None: |
| 133 | + self._lock = threading.Lock() |
| 134 | + self._init_state() |
| 135 | + |
| 136 | + def _init_state(self) -> None: |
| 137 | + self._pid = os.getpid() |
| 138 | + self._loop: asyncio.AbstractEventLoop | None = None |
| 139 | + self._thread: threading.Thread | None = None |
| 140 | + self._tasks: set[BackgroundTask] = set() |
| 141 | + self._closed = False |
| 142 | + |
| 143 | + @property |
| 144 | + def closed(self) -> bool: |
| 145 | + """Whether :meth:`shutdown` has run. A closed loop starts nothing.""" |
| 146 | + return self._closed |
| 147 | + |
| 148 | + def in_loop_thread(self) -> bool: |
| 149 | + """Whether the caller is running on the observability loop's thread.""" |
| 150 | + thread = self._thread |
| 151 | + return thread is not None and thread.ident == threading.get_ident() |
| 152 | + |
| 153 | + def _ensure_started_locked(self) -> asyncio.AbstractEventLoop | None: |
| 154 | + """The running loop, starting it if needed; ``None`` once shut down. |
| 155 | +
|
| 156 | + Caller holds ``_lock``. |
| 157 | + """ |
| 158 | + if self._pid != os.getpid(): |
| 159 | + self._init_state() |
| 160 | + if self._closed: |
| 161 | + return None |
| 162 | + if self._loop is not None and self._thread is not None: |
| 163 | + if self._thread.is_alive(): |
| 164 | + return self._loop |
| 165 | + loop = asyncio.new_event_loop() |
| 166 | + thread = threading.Thread( |
| 167 | + target=self._run, args=(loop,), name=THREAD_NAME, daemon=True |
| 168 | + ) |
| 169 | + thread.start() |
| 170 | + self._loop, self._thread = loop, thread |
| 171 | + return loop |
| 172 | + |
| 173 | + @staticmethod |
| 174 | + def _run(loop: asyncio.AbstractEventLoop) -> None: |
| 175 | + asyncio.set_event_loop(loop) |
| 176 | + try: |
| 177 | + loop.run_forever() |
| 178 | + finally: |
| 179 | + loop.close() |
| 180 | + |
| 181 | + def submit( |
| 182 | + self, coro: Coroutine[Any, Any, Any] |
| 183 | + ) -> concurrent.futures.Future[Any] | None: |
| 184 | + """Schedule a one-off *coro* on the loop from any thread. |
| 185 | +
|
| 186 | + Returns the ``concurrent.futures.Future`` for its result, or ``None`` |
| 187 | + (and closes *coro*) when the loop has been shut down. |
| 188 | + """ |
| 189 | + with self._lock: |
| 190 | + loop = self._ensure_started_locked() |
| 191 | + if loop is None: |
| 192 | + coro.close() |
| 193 | + return None |
| 194 | + try: |
| 195 | + return asyncio.run_coroutine_threadsafe(coro, loop) |
| 196 | + except RuntimeError: |
| 197 | + # The loop closed between the lock and the schedule. |
| 198 | + coro.close() |
| 199 | + return None |
| 200 | + |
| 201 | + def start( |
| 202 | + self, coro: Coroutine[Any, Any, Any], *, name: str, stop_order: int = 0 |
| 203 | + ) -> BackgroundTask | None: |
| 204 | + """Run the long-lived *coro* on the loop until it ends or is cancelled. |
| 205 | +
|
| 206 | + *coro* must handle ``asyncio.CancelledError`` as its signal to finish |
| 207 | + up (final flush, final batch): :meth:`shutdown` cancels it once and |
| 208 | + waits for it within the shutdown bound. Tasks with a higher |
| 209 | + *stop_order* are cancelled only after every lower one has finished, |
| 210 | + so a task other tasks' final work depends on (the Segment worker and |
| 211 | + its HTTP client) outlives them. |
| 212 | +
|
| 213 | + Returns its handle, or ``None`` (and closes *coro*) when the loop has |
| 214 | + been shut down. |
| 215 | + """ |
| 216 | + with self._lock: |
| 217 | + loop = self._ensure_started_locked() |
| 218 | + if loop is None: |
| 219 | + coro.close() |
| 220 | + return None |
| 221 | + handle = BackgroundTask(name, loop, stop_order) |
| 222 | + |
| 223 | + def _spawn() -> None: |
| 224 | + handle._attach(loop.create_task(handle._run(coro), name=name)) |
| 225 | + |
| 226 | + # Scheduled under the lock: ``shutdown`` takes it to close, so its |
| 227 | + # drain is always scheduled after this spawn and finds the task. |
| 228 | + try: |
| 229 | + loop.call_soon_threadsafe(_spawn) |
| 230 | + except RuntimeError: |
| 231 | + coro.close() |
| 232 | + handle._done.set() |
| 233 | + return None |
| 234 | + self._tasks = {t for t in self._tasks if not t.done()} |
| 235 | + self._tasks.add(handle) |
| 236 | + return handle |
| 237 | + |
| 238 | + def shutdown(self, timeout: float = OBSERVABILITY_SHUTDOWN_TIMEOUT_SECONDS) -> bool: |
| 239 | + """Stop every task and the loop within *timeout* seconds. Idempotent. |
| 240 | +
|
| 241 | + Returns ``True`` if everything finished inside the bound. Past it, the |
| 242 | + loop is stopped anyway and whatever was still running is abandoned: |
| 243 | + the thread is a daemon, so it cannot hold the process open. |
| 244 | + """ |
| 245 | + with self._lock: |
| 246 | + if self._pid != os.getpid(): |
| 247 | + self._init_state() |
| 248 | + if self._closed: |
| 249 | + return True |
| 250 | + self._closed = True |
| 251 | + loop, thread = self._loop, self._thread |
| 252 | + tasks = list(self._tasks) |
| 253 | + if loop is None or thread is None or not thread.is_alive(): |
| 254 | + return True |
| 255 | + if self.in_loop_thread(): |
| 256 | + # Cannot wait on our own thread; stop without waiting. |
| 257 | + loop.call_soon(loop.stop) |
| 258 | + return False |
| 259 | + |
| 260 | + deadline = time.monotonic() + max(0.0, timeout) |
| 261 | + clean = True |
| 262 | + try: |
| 263 | + drain = asyncio.run_coroutine_threadsafe(_drain(tasks), loop) |
| 264 | + except RuntimeError: |
| 265 | + return True # Already closed. |
| 266 | + try: |
| 267 | + drain.result(timeout=max(0.0, deadline - time.monotonic())) |
| 268 | + except TimeoutError: |
| 269 | + clean = False |
| 270 | + drain.cancel() |
| 271 | + logging.warning( |
| 272 | + "Observability background loop did not finish its final flushes " |
| 273 | + "within %ss; abandoning the rest", |
| 274 | + timeout, |
| 275 | + exc_info=True, |
| 276 | + ) |
| 277 | + # conformance: ignore[E004] shutdown path; a failed drain must not block exit |
| 278 | + except Exception: |
| 279 | + clean = False |
| 280 | + logging.warning("Error stopping observability tasks", exc_info=True) |
| 281 | + try: |
| 282 | + loop.call_soon_threadsafe(loop.stop) |
| 283 | + except RuntimeError: |
| 284 | + return clean # Already closed. |
| 285 | + thread.join(max(0.0, deadline - time.monotonic())) |
| 286 | + return clean and not thread.is_alive() |
| 287 | + |
| 288 | + def _reset_after_fork(self) -> None: |
| 289 | + """Forget the parent's loop and thread in a forked child. |
| 290 | +
|
| 291 | + The lock is replaced too: another parent thread may have held it at |
| 292 | + the moment of the fork, and the child would wait on it forever. |
| 293 | + """ |
| 294 | + self._lock = threading.Lock() |
| 295 | + self._init_state() |
| 296 | + |
| 297 | + def _reset_for_testing(self, timeout: float = 5.0) -> None: |
| 298 | + """Shut down (bounded by *timeout*) and allow a fresh start. |
| 299 | +
|
| 300 | + Test-only: production shuts down once, at exit. |
| 301 | + """ |
| 302 | + self.shutdown(timeout) |
| 303 | + with self._lock: |
| 304 | + self._init_state() |
| 305 | + |
| 306 | + |
| 307 | +async def _drain(tasks: list[BackgroundTask]) -> None: |
| 308 | + """Cancel each started task once, stage by stage, then wait for the rest. |
| 309 | +
|
| 310 | + Runs on the loop. Started tasks are cancelled in ascending ``stop_order``; |
| 311 | + each stage is awaited before the next is cancelled. Tasks it did not start |
| 312 | + (a batch flush ``add_record`` began on this loop) are awaited, not |
| 313 | + cancelled. Re-checks until none are left, since a final flush can log and |
| 314 | + so start another batch flush. |
| 315 | + """ |
| 316 | + current = asyncio.current_task() |
| 317 | + for order in sorted({handle.stop_order for handle in tasks}): |
| 318 | + stage = [handle for handle in tasks if handle.stop_order == order] |
| 319 | + for handle in stage: |
| 320 | + handle._cancel_requested = True |
| 321 | + handle._cancel_on_loop() |
| 322 | + await _await_all( |
| 323 | + [h._task for h in stage if h._task is not None and h._task is not current] |
| 324 | + ) |
| 325 | + while True: |
| 326 | + pending = [t for t in asyncio.all_tasks() if t is not current] |
| 327 | + if not pending: |
| 328 | + return |
| 329 | + await _await_all(pending) |
| 330 | + |
| 331 | + |
| 332 | +async def _await_all(pending: list[asyncio.Task[Any]]) -> None: |
| 333 | + """Wait for every task in *pending*, logging any that failed.""" |
| 334 | + results = await asyncio.gather(*pending, return_exceptions=True) |
| 335 | + for result in results: |
| 336 | + if isinstance(result, Exception): |
| 337 | + logging.warning( |
| 338 | + "Observability background task failed at shutdown", |
| 339 | + exc_info=(type(result), result, result.__traceback__), |
| 340 | + ) |
| 341 | + |
| 342 | + |
| 343 | +_after_fork_resets: list[Callable[[], None]] = [] |
| 344 | +_after_fork_restarts: list[Callable[[], None]] = [] |
| 345 | + |
| 346 | + |
| 347 | +def register_after_fork( |
| 348 | + callback: Callable[[], None], *, restarts_work: bool = False |
| 349 | +) -> None: |
| 350 | + """Run *callback* in a forked child, after the observability loop's reset. |
| 351 | +
|
| 352 | + For state that belongs to the parent's loop: a sink's periodic flush, the |
| 353 | + Segment worker. Callbacks that only drop inherited state run first, in |
| 354 | + registration order; callbacks that start work on the child's loop |
| 355 | + (*restarts_work*) run after all of them, so new work never meets state |
| 356 | + still pointing at the parent's loop. Callbacks must not raise. |
| 357 | + """ |
| 358 | + (_after_fork_restarts if restarts_work else _after_fork_resets).append(callback) |
| 359 | + |
| 360 | + |
| 361 | +def _after_fork_in_child() -> None: |
| 362 | + observability_loop._reset_after_fork() |
| 363 | + for callback in [*_after_fork_resets, *_after_fork_restarts]: |
| 364 | + try: |
| 365 | + callback() |
| 366 | + # conformance: ignore[E004] fork hook; one failing callback must not stop the rest or the child |
| 367 | + except Exception: |
| 368 | + logging.warning("Observability after-fork callback failed", exc_info=True) |
| 369 | + |
| 370 | + |
| 371 | +#: The process-wide observability loop. |
| 372 | +observability_loop = ObservabilityLoop() |
| 373 | + |
| 374 | +# The background thread is a daemon, which interpreter exit stops where it |
| 375 | +# stands; give the final flushes one bounded window first. atexit runs |
| 376 | +# handlers in reverse registration order, so handlers registered after this |
| 377 | +# import (none of the sinks need one any more) run before it. |
| 378 | +atexit.register(observability_loop.shutdown) |
| 379 | + |
| 380 | +# ``os.register_at_fork`` does not exist on Windows, which has no fork. |
| 381 | +if hasattr(os, "register_at_fork"): |
| 382 | + os.register_at_fork(after_in_child=_after_fork_in_child) |
0 commit comments