|
| 1 | +"""Client-side eviction notification monitor. |
| 2 | +
|
| 3 | +Backs :meth:`runloop_api_client.sdk.devbox.Devbox.on_evict`. A single account-wide |
| 4 | +SSE stream (the generated ``devboxes.watch_evictions`` endpoint) is opened lazily |
| 5 | +the moment the first devbox registers interest, and torn down as soon as the last |
| 6 | +interested devbox has been notified. |
| 7 | +
|
| 8 | +Delivery contract: |
| 9 | +
|
| 10 | +* The server replays every currently-pending eviction on connect, so a devbox that |
| 11 | + registers after its eviction was scheduled is still notified. |
| 12 | +* Notifications for devboxes not in the interest set are discarded. |
| 13 | +* A devbox is removed from the interest set *before* its callback runs, so the |
| 14 | + callback fires at most once even if the server repeats the notification. |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import logging |
| 20 | +import threading |
| 21 | +from typing import TYPE_CHECKING, Dict, Callable, Optional |
| 22 | +from weakref import WeakKeyDictionary |
| 23 | + |
| 24 | +if TYPE_CHECKING: |
| 25 | + from .._client import Runloop |
| 26 | + from .._streaming import Stream |
| 27 | + |
| 28 | + # Assumed Stainless-generated SSE event type for ``watch_evictions``; carries |
| 29 | + # ``devbox_id`` and ``eviction_deadline_ms``. Reconcile the import once the |
| 30 | + # generated code lands. |
| 31 | + from ..types import DevboxEvictionEvent |
| 32 | + |
| 33 | +EvictionCallback = Callable[["DevboxEvictionEvent"], None] |
| 34 | +"""Invoked once with the eviction event for the devbox it was registered for.""" |
| 35 | + |
| 36 | +_logger = logging.getLogger(__name__) |
| 37 | + |
| 38 | + |
| 39 | +class EvictionMonitor: |
| 40 | + """Fans account-wide eviction notifications out to per-devbox callbacks. |
| 41 | +
|
| 42 | + One monitor is shared by every :class:`~runloop_api_client.sdk.devbox.Devbox` |
| 43 | + built from the same generated client (see :func:`monitor_for`), so all |
| 44 | + registered devboxes are served by a single SSE connection. |
| 45 | + """ |
| 46 | + |
| 47 | + def __init__(self, client: "Runloop") -> None: |
| 48 | + self._client = client |
| 49 | + self._lock = threading.Lock() |
| 50 | + self._callbacks: Dict[str, EvictionCallback] = {} |
| 51 | + self._thread: Optional[threading.Thread] = None |
| 52 | + self._stream: Optional["Stream[DevboxEvictionEvent]"] = None |
| 53 | + |
| 54 | + def register(self, devbox_id: str, callback: EvictionCallback) -> None: |
| 55 | + """Add ``devbox_id`` to the interest set, starting the stream if idle.""" |
| 56 | + with self._lock: |
| 57 | + self._callbacks[devbox_id] = callback |
| 58 | + if self._thread is None or not self._thread.is_alive(): |
| 59 | + self._thread = threading.Thread( |
| 60 | + target=self._run, |
| 61 | + name="runloop-eviction-monitor", |
| 62 | + daemon=True, |
| 63 | + ) |
| 64 | + self._thread.start() |
| 65 | + |
| 66 | + def unregister(self, devbox_id: str) -> None: |
| 67 | + """Drop ``devbox_id``; close the stream if it was the last interested devbox.""" |
| 68 | + with self._lock: |
| 69 | + self._callbacks.pop(devbox_id, None) |
| 70 | + empty = not self._callbacks |
| 71 | + if empty: |
| 72 | + self._close_stream() |
| 73 | + |
| 74 | + def close(self) -> None: |
| 75 | + """Clear all interest and tear down the stream.""" |
| 76 | + with self._lock: |
| 77 | + self._callbacks.clear() |
| 78 | + self._close_stream() |
| 79 | + |
| 80 | + def _run(self) -> None: |
| 81 | + try: |
| 82 | + stream = self._client.devboxes.watch_evictions() |
| 83 | + with self._lock: |
| 84 | + self._stream = stream |
| 85 | + with stream: |
| 86 | + for event in stream: |
| 87 | + self._dispatch(event) |
| 88 | + with self._lock: |
| 89 | + if not self._callbacks: |
| 90 | + break |
| 91 | + except Exception: |
| 92 | + _logger.exception("eviction monitor stream failed") |
| 93 | + finally: |
| 94 | + with self._lock: |
| 95 | + self._stream = None |
| 96 | + self._thread = None |
| 97 | + |
| 98 | + def _dispatch(self, event: "DevboxEvictionEvent") -> None: |
| 99 | + with self._lock: |
| 100 | + callback = self._callbacks.pop(event.devbox_id, None) |
| 101 | + if callback is None: |
| 102 | + return |
| 103 | + try: |
| 104 | + callback(event) |
| 105 | + except Exception: |
| 106 | + _logger.exception("error in eviction callback for devbox %s", event.devbox_id) |
| 107 | + |
| 108 | + def _close_stream(self) -> None: |
| 109 | + with self._lock: |
| 110 | + stream = self._stream |
| 111 | + self._stream = None |
| 112 | + if stream is not None: |
| 113 | + try: |
| 114 | + stream.close() |
| 115 | + except Exception: |
| 116 | + _logger.debug("error closing eviction stream", exc_info=True) |
| 117 | + |
| 118 | + |
| 119 | +_monitors: "WeakKeyDictionary[Runloop, EvictionMonitor]" = WeakKeyDictionary() |
| 120 | +_monitors_lock = threading.Lock() |
| 121 | + |
| 122 | + |
| 123 | +def monitor_for(client: "Runloop") -> EvictionMonitor: |
| 124 | + """Return the shared :class:`EvictionMonitor` for ``client``, creating it once.""" |
| 125 | + with _monitors_lock: |
| 126 | + monitor = _monitors.get(client) |
| 127 | + if monitor is None: |
| 128 | + monitor = EvictionMonitor(client) |
| 129 | + _monitors[client] = monitor |
| 130 | + return monitor |
| 131 | + |
| 132 | + |
| 133 | +def shutdown_monitor_for(client: "Runloop") -> None: |
| 134 | + """Tear down the shared monitor for ``client`` if one exists.""" |
| 135 | + with _monitors_lock: |
| 136 | + monitor = _monitors.pop(client, None) |
| 137 | + if monitor is not None: |
| 138 | + monitor.close() |
0 commit comments