diff --git a/osprey_worker/src/osprey/worker/stress/consumer.py b/osprey_worker/src/osprey/worker/stress/consumer.py new file mode 100644 index 0000000..fe8439e --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/consumer.py @@ -0,0 +1,151 @@ +"""Stress harness result consumer. + +Subscribes to the `osprey.execution_results` Kafka topic and records the +wall-clock receive time per `ActionId` so the reporter can pair every input +event with its output and compute end-to-end latency. + +Producer-agnostic: the same consumer works against synthetic load (#324), +jetstream-driven load (once #236 lands), or any other input source. The +optional `action_id_filter` decides whether we're in closed-loop matching mode +(only count IDs we know about) or open-loop throughput mode (count everything). +""" + +from __future__ import annotations + +import json +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from kafka import KafkaConsumer + + +@dataclass(frozen=True) +class ConsumerConfig: + bootstrap_servers: list[str] + topic: str + group_id: str + # If set, only record action_ids in this set. None = open-loop, count all. + action_id_filter: Optional[frozenset[int]] = None + # Max wall-clock time to keep reading after start() is called. + max_runtime_seconds: float = 60.0 + # If we've matched every id in action_id_filter, stop early. + stop_when_filter_complete: bool = True + client_id: str = 'osprey-stress-consumer' + + +class Consumer: + """Consumes execution_results on a background thread. + + `start()` -> `wait()` -> read `consumed`. `stop()` to abort early. + """ + + def __init__(self, config: ConsumerConfig, *, consumer_factory=KafkaConsumer): + self._config = config + self._consumer_factory = consumer_factory + self._consumed: dict[int, float] = {} + self._stop_event = threading.Event() + self._thread: Optional[threading.Thread] = None + self._error: Optional[BaseException] = None + + @property + def consumed(self) -> dict[int, float]: + return dict(self._consumed) + + @property + def error(self) -> Optional[BaseException]: + return self._error + + def start(self) -> None: + if self._thread is not None: + raise RuntimeError('Consumer already started') + self._thread = threading.Thread(target=self._run, name='stress-consumer', daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop_event.set() + + def wait(self, timeout: Optional[float] = None) -> None: + if self._thread is None: + return + self._thread.join(timeout=timeout) + + def _run(self) -> None: + try: + consumer = self._consumer_factory( + self._config.topic, + bootstrap_servers=self._config.bootstrap_servers, + group_id=self._config.group_id, + client_id=self._config.client_id, + # Start at 'latest' so we don't pick up old results from + # previous runs. The producer should be started after the + # consumer is assigned a partition (caller orchestrates). + auto_offset_reset='latest', + enable_auto_commit=False, + # Bounded poll so the stop_event check has a chance to fire. + consumer_timeout_ms=200, + ) + except Exception as e: + self._error = e + return + + deadline = time.monotonic() + self._config.max_runtime_seconds + try: + while not self._stop_event.is_set() and time.monotonic() < deadline: + # The kafka-python iterator yields until consumer_timeout_ms + # of inactivity, then raises StopIteration. We catch that and + # loop so we can re-check stop conditions. + try: + for message in consumer: + if self._stop_event.is_set(): + break + recorded = self._record(message.value) + if ( + recorded + and self._config.stop_when_filter_complete + and self._config.action_id_filter is not None + and len(self._consumed) >= len(self._config.action_id_filter) + ): + return + except StopIteration: + # Idle poll cycle — loop and re-check stop conditions. + continue + except Exception as e: + self._error = e + finally: + try: + consumer.close(autocommit=False) + except Exception: + # Best-effort close during shutdown: the kafka client may already + # be in a half-disconnected state, and a noisy close() error here + # would mask the consumed/error state the caller needs to read. + pass + + def _record(self, raw: bytes) -> bool: + """Parse one result message; return True if it counted toward consumed.""" + try: + payload = json.loads(raw.decode('utf-8')) + except (ValueError, UnicodeDecodeError): + # Garbage at the byte level is tolerated — Kafka can carry malformed + # bytes from broken producers and we shouldn't take down the + # measurement run for transport noise. + return False + action_id = payload.get('ActionId') + if not isinstance(action_id, int): + # Schema-level violation: a well-formed JSON result missing a valid + # ActionId means the harness was pointed at rules that didn't wire + # GetActionId() (or wired it wrong). Raise loudly so the caller sees + # the misconfiguration via consumer.error instead of getting silent + # zero matches. + raise ValueError( + f'result message has missing or non-int ActionId (got {type(action_id).__name__}); ' + f'ensure the rule set exposes ActionId via GetActionId()' + ) + if self._config.action_id_filter is not None and action_id not in self._config.action_id_filter: + return False + # First-write-wins so duplicate effect dispatches don't reset the time. + if action_id not in self._consumed: + self._consumed[action_id] = time.time() + return True + return False diff --git a/osprey_worker/src/osprey/worker/stress/tests/test_consumer.py b/osprey_worker/src/osprey/worker/stress/tests/test_consumer.py new file mode 100644 index 0000000..eb7b422 --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/tests/test_consumer.py @@ -0,0 +1,277 @@ +import json +import threading +import time +from typing import Any, Iterator, List + +from osprey.worker.stress.consumer import Consumer, ConsumerConfig + + +class FakeMessage: + def __init__(self, value: bytes) -> None: + self.value = value + + +class FakeKafkaConsumer: + """In-memory consumer mirroring kafka-python's iteration semantics. + + A real KafkaConsumer's `__next__` blocks for up to `consumer_timeout_ms` + waiting for the next message, then raises StopIteration if nothing + arrived. The Consumer under test relies on that natural backpressure to + avoid spinning. Reproduce the same shape here: when the queue is empty, + wait briefly (Event-driven so feed() can wake us early) before raising, + so the test never devolves into a busy-loop that can starve the main + thread under CI's thread contention. + """ + + _IDLE_POLL_SECONDS = 0.05 + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._lock = threading.Lock() + self._queue: List[FakeMessage] = [] + self._has_message = threading.Event() + self.kwargs = kwargs + self.close_called = False + + def feed(self, raw: bytes) -> None: + with self._lock: + self._queue.append(FakeMessage(raw)) + self._has_message.set() + + def __iter__(self) -> Iterator[FakeMessage]: + return self + + def __next__(self) -> FakeMessage: + # Drain mode: if there's already a message, return it without sleeping. + with self._lock: + if self._queue: + msg = self._queue.pop(0) + if not self._queue: + self._has_message.clear() + return msg + # Idle poll: wait briefly for a feed(), then re-check. + self._has_message.wait(timeout=self._IDLE_POLL_SECONDS) + with self._lock: + if self._queue: + msg = self._queue.pop(0) + if not self._queue: + self._has_message.clear() + return msg + raise StopIteration + + def close(self, autocommit: bool = True) -> None: + self.close_called = True + + +def make_result_payload(action_id: int) -> bytes: + return json.dumps({'ActionId': action_id, 'ActionName': 'create_post'}).encode('utf-8') + + +class TestConsumerClosedLoop: + def test_records_matching_action_ids(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='osprey.execution_results', + group_id='test', + action_id_filter=frozenset({1, 2, 3}), + max_runtime_seconds=2.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + for aid in [1, 2, 3]: + fake.feed(make_result_payload(aid)) + consumer.wait(timeout=3) + assert set(consumer.consumed.keys()) == {1, 2, 3} + + def test_ignores_non_matching_action_ids(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=frozenset({1, 2}), + max_runtime_seconds=1.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + # Feed traffic that doesn't belong to this run. + for aid in [99, 100, 101]: + fake.feed(make_result_payload(aid)) + time.sleep(0.5) + consumer.stop() + consumer.wait(timeout=2) + assert consumer.consumed == {} + + def test_stops_early_when_filter_complete(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=frozenset({1, 2, 3}), + max_runtime_seconds=60.0, # would block for a long time + stop_when_filter_complete=True, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + for aid in [1, 2, 3]: + fake.feed(make_result_payload(aid)) + start = time.monotonic() + consumer.wait(timeout=10) + elapsed = time.monotonic() - start + # Stops well under max_runtime once it has matched the whole filter. + assert elapsed < 5.0 + assert set(consumer.consumed.keys()) == {1, 2, 3} + + def test_first_write_wins_for_duplicate_action_ids(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=frozenset({1}), + max_runtime_seconds=2.0, + stop_when_filter_complete=False, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + fake.feed(make_result_payload(1)) + # Poll until the first message has actually been recorded — sleeping a + # fixed interval can race on slow CI, leaving first_ts=None and turning + # the equality check below into a meaningless assertion. + deadline = time.monotonic() + 2.0 + while consumer.consumed.get(1) is None and time.monotonic() < deadline: + time.sleep(0.05) + first_ts = consumer.consumed.get(1) + assert first_ts is not None, 'consumer never recorded the first message' + fake.feed(make_result_payload(1)) + # Give the consumer a chance to (incorrectly) overwrite if first-write-wins + # were broken. 0.5s is plenty since the poll cycle is 200ms. + time.sleep(0.5) + consumer.stop() + consumer.wait(timeout=2) + assert consumer.consumed[1] == first_ts + + +class TestConsumerOpenLoop: + def test_counts_everything_without_filter(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=None, + max_runtime_seconds=1.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + for aid in [10, 20, 30, 40, 50]: + fake.feed(make_result_payload(aid)) + consumer.wait(timeout=3) + assert set(consumer.consumed.keys()) == {10, 20, 30, 40, 50} + + +class TestConsumerMessageHygiene: + def test_skips_malformed_json(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=None, + max_runtime_seconds=1.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + fake.feed(b'not json{{{') + fake.feed(make_result_payload(42)) + consumer.wait(timeout=3) + assert consumer.consumed == {42: consumer.consumed[42]} + + def test_raises_on_missing_action_id_field(self) -> None: + # A well-formed result with no ActionId means the rule set didn't + # wire GetActionId() — that's a configuration error, not noise. The + # error should surface via consumer.error so the CLI reports it. + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=None, + max_runtime_seconds=1.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + fake.feed(json.dumps({'OtherField': 'no action id here'}).encode('utf-8')) + consumer.wait(timeout=3) + assert isinstance(consumer.error, ValueError) + assert 'ActionId' in str(consumer.error) + + def test_raises_on_non_int_action_id(self) -> None: + # A result that has ActionId but emits it as a string means a rule + # is overriding the int contract. Raise so the caller sees it instead + # of getting silent zero matches. + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=None, + max_runtime_seconds=1.0, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + fake.feed(json.dumps({'ActionId': 'not-an-int'}).encode('utf-8')) + consumer.wait(timeout=3) + assert isinstance(consumer.error, ValueError) + assert 'ActionId' in str(consumer.error) + + +class TestConsumerLifecycle: + def test_starting_twice_raises(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + max_runtime_seconds=0.5, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + consumer.start() + try: + consumer.start() + assert False, 'should have raised' + except RuntimeError: + pass + consumer.wait(timeout=2) + + def test_factory_failure_captured(self) -> None: + def boom(*_: Any, **__: Any) -> Any: + raise ConnectionError('no kafka') + + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + ) + consumer = Consumer(config, consumer_factory=boom) + consumer.start() + consumer.wait(timeout=2) + assert isinstance(consumer.error, ConnectionError) + assert consumer.consumed == {} + + def test_respects_max_runtime(self) -> None: + fake = FakeKafkaConsumer() + config = ConsumerConfig( + bootstrap_servers=['ignored'], + topic='topic', + group_id='test', + action_id_filter=frozenset({1, 2, 3}), # never fed → never completes + max_runtime_seconds=0.5, + ) + consumer = Consumer(config, consumer_factory=lambda *a, **kw: fake) + start = time.monotonic() + consumer.start() + consumer.wait(timeout=3) + elapsed = time.monotonic() - start + assert 0.4 < elapsed < 2.0