Skip to content

Commit cdfbbdb

Browse files
authored
refactor(observability): run every periodic flush on one background loop with one shutdown bound [FND-4648] (#4819)
1 parent 7af6581 commit cdfbbdb

19 files changed

Lines changed: 1295 additions & 621 deletions

‎application_sdk/constants.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -991,6 +991,14 @@ def _load_service_version() -> str:
991991
SIZING_FILE_NAME = os.environ.get("ATLAN_SIZING_FILE_NAME", "sizing.parquet")
992992
TRACES_FILE_NAME = os.environ.get("ATLAN_TRACES_FILE_NAME", "traces.parquet")
993993

994+
#: Upper bound, in seconds, on stopping the shared observability background
995+
#: loop at process exit: the final flush of every periodic sink (logs,
996+
#: metrics, traces, sizing) and the Segment worker's last batch, together.
997+
#: See ``application_sdk.observability.background_loop``.
998+
OBSERVABILITY_SHUTDOWN_TIMEOUT_SECONDS = float(
999+
os.environ.get("ATLAN_OBSERVABILITY_SHUTDOWN_TIMEOUT_SECONDS", "10")
1000+
)
1001+
9941002
# Segment Configuration
9951003
#: Segment API URL for sending events. Defaults to https://api.segment.io/v1/batch
9961004
SEGMENT_API_URL = os.getenv("ATLAN_SEGMENT_API_URL", "https://api.segment.io/v1/batch")
Lines changed: 382 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,382 @@
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

Comments
 (0)