diff --git a/osprey_async_worker/src/osprey/async_worker/lib/discovery/async_directory.py b/osprey_async_worker/src/osprey/async_worker/lib/discovery/async_directory.py index 715599cf..997efa09 100644 --- a/osprey_async_worker/src/osprey/async_worker/lib/discovery/async_directory.py +++ b/osprey_async_worker/src/osprey/async_worker/lib/discovery/async_directory.py @@ -12,7 +12,7 @@ import logging from random import randint, uniform from time import time -from typing import Any, Callable, ClassVar, Deque, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, ClassVar, Deque, Dict, List, Optional, Tuple from hash_ring import HashRing, HashRingNode from osprey.worker.lib import etcd @@ -133,27 +133,34 @@ def _read_overrides(raw: list) -> Dict[Any, Any]: return result async def _watch_loop(self, watcher: Any) -> None: + # Drive the SAME continue_watching() generator every iteration: recreating it + # per event resets the watcher's resume index and dedup mux, defeating its + # built-in recovery (the smite_shortlist freeze). Reopen only on + # exhaustion/error, with exponential backoff. Mirrors AsyncEtcdSourcesProvider. loop = asyncio.get_running_loop() + backoff = 1.0 + watcher_iter = watcher.continue_watching() try: while True: - event = await loop.run_in_executor(None, self._blocking_next, watcher) - if event is None: - break - self._handle_event(event) + it = watcher_iter + try: + event = await loop.run_in_executor(None, lambda: next(it)) + backoff = 1.0 + except StopIteration: + watcher_iter = watcher.continue_watching() + continue + except Exception: + logger.exception('async hash ring watch loop failed for %s', self._key) + await asyncio.sleep(backoff) + backoff = min(backoff * 2, 30.0) + watcher_iter = watcher.continue_watching() + continue + try: + self._handle_event(event) + except Exception: + logger.exception('async hash ring failed to apply event for %s', self._key) except asyncio.CancelledError: pass - except Exception: - logger.exception('hash ring watcher failed') - - @staticmethod - def _blocking_next(watcher: Any) -> Any: - """Blocking call — runs in thread pool.""" - try: - for event in watcher.continue_watching(): - return event - except Exception: - return None - return None async def stop(self) -> None: if self._watch_task: @@ -217,6 +224,9 @@ async def ensure_initialized(self) -> None: self._service_name, len(self._instances), list(self._instances.keys()), ) await self._ring.ensure_initialized() + # Events are applied on the loop thread, keeping instance mutation + # single-threaded and consistent with select(), and keeping any + # DOWN-listener create_task on a running loop. self._watch_task = asyncio.create_task(self._watch_loop(watcher)) except Exception: self._initialized = False @@ -310,31 +320,39 @@ def _remove_instance(self, instance_id: str) -> None: logger.debug('async discovery: - %s@%s', wrapper.service.name, instance_id) async def _watch_loop(self, watcher: Any) -> None: + # Drive the SAME continue_watching() generator every iteration: recreating it + # per event resets the watcher's resume index and dedup mux, defeating its + # built-in recovery (the smite_shortlist freeze). Reopen only on + # exhaustion/error, with exponential backoff. Mirrors AsyncEtcdSourcesProvider. loop = asyncio.get_running_loop() + backoff = 1.0 + watcher_iter = watcher.continue_watching() try: while True: - event = await loop.run_in_executor(None, self._blocking_next, watcher) - if event is None: - break - if isinstance(event, etcd.IncrementalSyncUpsert): - self._add_instance(Service.deserialize(event.value)) - elif isinstance(event, etcd.IncrementalSyncDelete): - self._remove_instance(Service.deserialize(event.prev_value).id) - elif isinstance(event, etcd.FullSyncRecursive): - self._handle_full_sync(event) + it = watcher_iter + try: + event = await loop.run_in_executor(None, lambda: next(it)) + backoff = 1.0 + except StopIteration: + watcher_iter = watcher.continue_watching() + continue + except Exception: + logger.exception('async service watcher failed for %s', self._service_name) + await asyncio.sleep(backoff) + backoff = min(backoff * 2, 30.0) + watcher_iter = watcher.continue_watching() + continue + try: + if isinstance(event, etcd.IncrementalSyncUpsert): + self._add_instance(Service.deserialize(event.value)) + elif isinstance(event, etcd.IncrementalSyncDelete): + self._remove_instance(Service.deserialize(event.prev_value).id) + elif isinstance(event, etcd.FullSyncRecursive): + self._handle_full_sync(event) + except Exception: + logger.exception('async service watcher failed to apply event for %s', self._service_name) except asyncio.CancelledError: pass - except Exception: - logger.exception('async service watcher failed for %s', self._service_name) - - @staticmethod - def _blocking_next(watcher: Any) -> Any: - try: - for event in watcher.continue_watching(): - return event - except Exception: - return None - return None async def stop(self) -> None: if self._watch_task: diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_async_directory.py b/osprey_async_worker/src/osprey/async_worker/tests/test_async_directory.py new file mode 100644 index 00000000..ff58d02a --- /dev/null +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_async_directory.py @@ -0,0 +1,253 @@ +"""Regression tests for the continuous-consumption etcd watcher. + +These tests cover the production incident where a worker's smite_shortlist instance +set got wiped to empty and never repopulated (select() raised "No service" forever +until the worker was restarted). The root cause was that the watch loop consumed +``continue_watching()`` one event at a time, recreating the generator per event, +which reset its resume index / dedup mux and defeated its built-in recovery. + +The fix drives a single persistent ``continue_watching()`` generator in +``_watch_loop`` (one asyncio task per watcher, pumping ``next(it)`` through the +executor), applying each event on the loop thread. + +Covered here: + * recovery-after-empty-full-sync — an empty FullSyncRecursive wipes the set (no + guard), and the SAME continuous generator repopulates it via incremental upserts; + * stop() teardown — stop() cancels the watch task (and the ring's); + * the same fix applied to _AsyncHashRing (SCALAR routing). +""" + +import asyncio +import json +import queue +from types import SimpleNamespace + +import pytest +from osprey.worker.lib.discovery.exceptions import ServiceUnavailable +from osprey.worker.lib.discovery.service import Service +from osprey.worker.lib.etcd import ( + FullSyncOne, + FullSyncOneNoKey, + FullSyncRecursive, + IncrementalSyncDelete, + IncrementalSyncUpsert, +) + +from osprey.async_worker.lib.discovery.async_directory import AsyncServiceWatcher, _AsyncHashRing + +# Sentinel pushed into a fake watcher's inbox to make continue_watching() return. +# After the watch task is cancelled, the executor thread is still parked in +# next() on inbox.get(); pushing this unblocks it so the pool thread is freed. +_UNBLOCK = object() + + +class FakeWatcher: + """A controllable stand-in for an etcd watcher. + + ``begin_watching()`` returns a preset initial sync event. ``continue_watching()`` + is an infinite generator that yields whatever the test pushes into its inbox and + only returns when ``_UNBLOCK`` is pushed (blocking on the inbox in between, like + the real watcher blocking on etcd). + """ + + def __init__(self, initial_event): + self.initial_event = initial_event + self.inbox: "queue.Queue" = queue.Queue() + self.begin_count = 0 + self.continue_count = 0 + + def begin_watching(self): + self.begin_count += 1 + return self.initial_event + + def push(self, event): + self.inbox.put(event) + + def unblock(self): + """Release a blocked continue_watching() so the parked executor thread frees.""" + self.inbox.put(_UNBLOCK) + + def continue_watching(self): + self.continue_count += 1 + while True: + item = self.inbox.get() + if item is _UNBLOCK: + return + yield item + + +class FakeEtcdClient: + """Serves preconfigured FakeWatchers, routed by the ``recursive`` flag. + + AsyncServiceWatcher requests a recursive watcher (instances); _AsyncHashRing + requests a non-recursive (scalar) watcher (ring). Extra requests beyond the + provided pools get a benign idle watcher. + """ + + def __init__(self, recursive_watchers=None, scalar_watchers=None): + self._recursive = list(recursive_watchers or []) + self._scalar = list(scalar_watchers or []) + self.calls = [] # list of (key, recursive) + self.all_watchers = [] + + def get_watcher(self, key, recursive=False, _use_mux=True): + self.calls.append((key, recursive)) + pool = self._recursive if recursive else self._scalar + if pool: + watcher = pool.pop(0) + else: + default = FullSyncRecursive(items=[]) if recursive else FullSyncOneNoKey(key=key) + watcher = FakeWatcher(initial_event=default) + self.all_watchers.append(watcher) + return watcher + + +# --- helpers ----------------------------------------------------------------- + + +def _svc(host: str, port: int) -> Service: + return Service(name='smite_shortlist', address=host, port=port) + + +def _node(value: str) -> SimpleNamespace: + return SimpleNamespace(key=f'/k/{value}', value=value) + + +def _full_sync(services) -> FullSyncRecursive: + return FullSyncRecursive(items=[_node(s.serialize()) for s in services]) + + +def _upsert(service: Service) -> IncrementalSyncUpsert: + return IncrementalSyncUpsert(key=f'/k/{service.id}', value=service.serialize()) + + +def _delete(service: Service) -> IncrementalSyncDelete: + return IncrementalSyncDelete(key=f'/k/{service.id}', prev_value=service.serialize()) + + +def _ring_full_sync(members) -> FullSyncOne: + return FullSyncOne(key='/discovery/svc/ring', value=json.dumps(list(members))) + + +async def _wait_until(predicate, timeout: float = 2.0, interval: float = 0.01) -> bool: + """Poll ``predicate`` (yielding to the loop so the watch task can run).""" + elapsed = 0.0 + while elapsed < timeout: + if predicate(): + return True + await asyncio.sleep(interval) + elapsed += interval + return predicate() + + +async def _teardown_service(watcher: AsyncServiceWatcher, client: FakeEtcdClient) -> None: + await watcher.stop() + # stop() cancels the watch tasks, but the executor threads driving next() are + # still parked on the fake inboxes — release them so the pool threads are freed. + for fw in client.all_watchers: + fw.unblock() + + +# --- AsyncServiceWatcher ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_recovers_after_empty_full_sync(): + """An empty FullSyncRecursive wipes the set; the same generator repopulates it. + + This is the incident: a key-deleted / index_cleared full sync emptied the instance + set. With continuous consumption the subsequent incremental upsert (yielded by the + SAME generator) restores the instance, so select() recovers without a restart. + """ + svc = _svc('shortlist-a', 1001) + watcher = FakeWatcher(initial_event=_full_sync([svc])) + client = FakeEtcdClient(recursive_watchers=[watcher]) + + asw = AsyncServiceWatcher(client, '/discovery', 'smite_shortlist') + await asw.ensure_initialized() + assert set(asw._instances) == {svc.id} + + # Empty full sync wipes the set (no guard — the set self-repopulates). + watcher.push(_full_sync([])) + assert await _wait_until(lambda: len(asw._instances) == 0) + with pytest.raises(ServiceUnavailable): + asw.select() + + # Incremental upsert on the same continuous generator repopulates the set. + watcher.push(_upsert(svc)) + assert await _wait_until(lambda: svc.id in asw._instances) + assert asw.select() is not None # no longer raises "No service" + + await _teardown_service(asw, client) + + +@pytest.mark.asyncio +async def test_incremental_upsert_and_delete_applied(): + """Plain upsert/delete events flow through the watch loop and mutate the set.""" + a, b = _svc('a', 1), _svc('b', 2) + watcher = FakeWatcher(initial_event=_full_sync([a])) + client = FakeEtcdClient(recursive_watchers=[watcher]) + + asw = AsyncServiceWatcher(client, '/discovery', 'smite_shortlist') + await asw.ensure_initialized() + assert set(asw._instances) == {a.id} + + watcher.push(_upsert(b)) + assert await _wait_until(lambda: set(asw._instances) == {a.id, b.id}) + + watcher.push(_delete(a)) + assert await _wait_until(lambda: set(asw._instances) == {b.id}) + + await _teardown_service(asw, client) + + +@pytest.mark.asyncio +async def test_stop_teardown(): + """stop() cancels the watch task for both the service watcher and its ring.""" + svc = _svc('a', 1) + watcher = FakeWatcher(initial_event=_full_sync([svc])) + client = FakeEtcdClient(recursive_watchers=[watcher]) + + asw = AsyncServiceWatcher(client, '/discovery', 'smite_shortlist') + await asw.ensure_initialized() + + watch_task = asw._watch_task + ring_task = asw._ring._watch_task + assert watch_task is not None and not watch_task.done() + assert ring_task is not None and not ring_task.done() + + await asw.stop() + + assert watch_task.done() + assert ring_task.done() + + # Free the executor threads still parked in next() on the fake inboxes. + for fw in client.all_watchers: + fw.unblock() + + +# --- _AsyncHashRing ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ring_recovers_after_key_deleted(): + """Same latent bug on the SCALAR ring: a key-deleted full sync wipes members, + and continuous consumption repopulates them from a later FullSyncOne.""" + key = '/discovery/svc/ring' + ring_watcher = FakeWatcher(initial_event=_ring_full_sync(['a:1', 'b:2'])) + client = FakeEtcdClient(scalar_watchers=[ring_watcher]) + + ring = _AsyncHashRing(client, key) + await ring.ensure_initialized() + assert len(ring._members) == 2 + + # Key deleted -> ring wiped. + ring_watcher.push(FullSyncOneNoKey(key=key)) + assert await _wait_until(lambda: len(ring._members) == 0) + + # Later value repopulates the ring on the same continuous generator. + ring_watcher.push(FullSyncOne(key=key, value=json.dumps(['c:3']))) + assert await _wait_until(lambda: len(ring._members) == 1) + + await ring.stop() + ring_watcher.unblock()