diff --git a/osprey_worker/src/osprey/worker/stress/producer.py b/osprey_worker/src/osprey/worker/stress/producer.py new file mode 100644 index 0000000..5d979f0 --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/producer.py @@ -0,0 +1,163 @@ +"""Stress harness synthetic producer. + +Emits N well-formed Osprey actions to the Kafka input topic at a configurable +rate, recording the wall-clock send time per `action_id` so the reporter can +compute end-to-end latency. + +Events match the shape expected by `example_rules/` (the rule set the +test_runner is configured with). The consumer matches each input event to its +output `ExecutionResult` via the `ActionId` extracted feature, which is +surfaced by the `GetActionId()` stdlib UDF. +""" + +from __future__ import annotations + +import json +import threading +import time +import uuid +from dataclasses import dataclass +from typing import Optional + +from kafka import KafkaProducer + +# action_ids in this range sit above what the bundled `generate_test_data.sh` +# emits but below 2**53, so they survive JSON's integer-as-float precision. +# A per-run multiplier keyed on the low bits of the run uuid keeps concurrent +# stress runs from colliding with each other. +_ACTION_ID_BASE = 1_000_000_000_000 +_RUN_BUCKET = 10_000_000 + + +def _action_id_for(run_id: str, n: int) -> int: + """Deterministic, per-run-unique integer action_id. + + The id encodes the run_id (so we can filter on the consumer side) and the + event sequence number, all without exceeding JS safe-integer precision. + """ + bucket = int(run_id[:6], 16) % _RUN_BUCKET + return _ACTION_ID_BASE + bucket * _RUN_BUCKET + n + + +@dataclass(frozen=True) +class ProducerConfig: + bootstrap_servers: list[str] + topic: str + events: int + rate_per_second: float + run_id: str + client_id: str = 'osprey-stress' + + @staticmethod + def make_run_id() -> str: + return uuid.uuid4().hex[:8] + + +def build_event(run_id: str, n: int, *, now: Optional[float] = None) -> tuple[int, bytes]: + """Return `(action_id, json_bytes)` for the n-th event in a run.""" + action_id = _action_id_for(run_id, n) + ts = now if now is not None else time.time() + send_time = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(ts)) + ip_octet = (n % 254) + 1 + payload = { + 'send_time': send_time, + 'data': { + 'action_id': action_id, + 'action_name': 'create_post', + 'data': { + 'user_id': f'stress_user_{n % 100}', + 'ip_address': f'192.168.1.{ip_octet}', + 'event_type': 'create_post', + 'post': {'text': f'hello stress {n}'}, + }, + }, + } + return action_id, json.dumps(payload).encode('utf-8') + + +class Producer: + """Produces events on a background thread. + + Use `start()` to kick off, `wait()` to block until done (or limit reached), + and `produced` to retrieve the `{action_id: timestamp}` map after stopping. + `stop()` requests early shutdown — useful for Ctrl+C. + """ + + def __init__(self, config: ProducerConfig, *, producer_factory=KafkaProducer): + self._config = config + self._producer_factory = producer_factory + self._produced: dict[int, float] = {} + self._stop_event = threading.Event() + self._thread: Optional[threading.Thread] = None + self._error: Optional[BaseException] = None + + @property + def produced(self) -> dict[int, float]: + return dict(self._produced) + + @property + def error(self) -> Optional[BaseException]: + return self._error + + def start(self) -> None: + if self._thread is not None: + raise RuntimeError('Producer already started') + self._thread = threading.Thread(target=self._run, name='stress-producer', 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: + producer = self._producer_factory( + bootstrap_servers=self._config.bootstrap_servers, + client_id=self._config.client_id, + linger_ms=0, + ) + except Exception as e: + self._error = e + return + + interval = 1.0 / self._config.rate_per_second if self._config.rate_per_second > 0 else 0.0 + next_send = time.monotonic() + + try: + for n in range(self._config.events): + if self._stop_event.is_set(): + break + action_id, value = build_event(self._config.run_id, n) + send_at = time.time() + producer.send(self._config.topic, value=value) + self._produced[action_id] = send_at + + # Sleep until the next scheduled slot. If we've drifted (Kafka + # slow / GC pause), reset the schedule to now rather than burst- + # producing to "catch up." + next_send += interval + now = time.monotonic() + if next_send > now: + time.sleep(next_send - now) + else: + next_send = now + except Exception as e: + self._error = e + finally: + # flush() and close() run in independent try blocks so a flush + # failure doesn't skip close() and leak the underlying socket. + # Either failure is captured, but only if the send loop above + # didn't already surface a more useful error. + for step in ( + lambda: producer.flush(timeout=10), + lambda: producer.close(timeout=10), + ): + try: + step() + except Exception as cleanup_error: + if self._error is None: + self._error = cleanup_error diff --git a/osprey_worker/src/osprey/worker/stress/tests/test_producer.py b/osprey_worker/src/osprey/worker/stress/tests/test_producer.py new file mode 100644 index 0000000..f414b79 --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/tests/test_producer.py @@ -0,0 +1,231 @@ +import json +import threading +import time +from typing import NoReturn + +from osprey.worker.stress.producer import ( + Producer, + ProducerConfig, + _action_id_for, + build_event, +) + + +class FakeKafkaProducer: + """Captures sends in-memory; mimics the kafka-python API surface we use.""" + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + self.sent: list[tuple[str, bytes]] = [] + self._lock = threading.Lock() + self.flush_called = False + self.close_called = False + + def send(self, topic: str, value: bytes) -> None: + with self._lock: + self.sent.append((topic, value)) + + def flush(self, timeout: float = 0) -> None: + self.flush_called = True + + def close(self, timeout: float = 0) -> None: + self.close_called = True + + +class TestActionIdGenerator: + def test_returns_int(self) -> None: + assert isinstance(_action_id_for('abcdef12', 0), int) + + def test_unique_per_event_within_run(self) -> None: + ids = {_action_id_for('abcdef12', n) for n in range(1000)} + assert len(ids) == 1000 + + def test_different_runs_dont_collide_for_same_n(self) -> None: + a = _action_id_for('11111111', 0) + b = _action_id_for('22222222', 0) + assert a != b + + def test_under_javascript_safe_integer(self) -> None: + # 2**53 = 9007199254740992. Action_ids must fit so JSON consumers (Druid, + # browsers) don't lose precision. + for run_id in ['00000000', 'ffffffff']: + for n in [0, 1_000_000, 9_999_999]: + assert _action_id_for(run_id, n) < 2**53 + + +class TestBuildEvent: + def test_action_id_is_int(self) -> None: + action_id, _ = build_event('abcdef12', 5) + assert isinstance(action_id, int) + + def test_payload_shape_matches_template(self) -> None: + action_id, raw = build_event('abcdef12', 7, now=1700000000.0) + payload = json.loads(raw) + # Must align with example_data/template.json so example_rules' SML can + # extract fields by the same JSON paths. + assert payload['data']['action_id'] == action_id + assert isinstance(payload['data']['action_id'], int) + assert payload['data']['action_name'] == 'create_post' + assert payload['data']['data']['user_id'] == 'stress_user_7' + assert payload['data']['data']['event_type'] == 'create_post' + assert 'hello' in payload['data']['data']['post']['text'] + assert payload['send_time'].endswith('Z') + + def test_user_id_is_semantically_a_user_id(self) -> None: + # Regression guard: don't overload user_id for tracking purposes. + # Tracking is done via action_id + GetActionId() UDF — user_id stays + # a plausible user identifier. + _, raw = build_event('abcdef12', 7) + payload = json.loads(raw) + assert payload['data']['data']['user_id'].startswith('stress_user_') + + def test_deterministic_for_same_inputs(self) -> None: + a_id, a_raw = build_event('abcdef12', 1, now=1700000000.0) + b_id, b_raw = build_event('abcdef12', 1, now=1700000000.0) + assert a_id == b_id + assert a_raw == b_raw + + +class TestProducer: + def test_produces_exact_count(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='osprey.actions_input', + events=20, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + producer.start() + producer.wait(timeout=5) + assert producer.error is None + assert len(fake.sent) == 20 + assert len(producer.produced) == 20 + + def test_action_ids_unique_and_ordered(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=10, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + producer.start() + producer.wait(timeout=5) + ids = list(producer.produced.keys()) + assert ids == [_action_id_for('abcdef12', n) for n in range(10)] + + def test_records_send_timestamp_per_event(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=5, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + before = time.time() + producer.start() + producer.wait(timeout=5) + after = time.time() + for ts in producer.produced.values(): + assert before <= ts <= after + + def test_rate_control_approximate(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=20, + rate_per_second=100.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + start = time.monotonic() + producer.start() + producer.wait(timeout=5) + elapsed = time.monotonic() - start + assert 0.15 < elapsed < 1.0, f'elapsed={elapsed:.3f}s' + + def test_stop_aborts_early(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=1000, + rate_per_second=100.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + producer.start() + time.sleep(0.1) + producer.stop() + producer.wait(timeout=5) + assert len(fake.sent) < 1000 + assert fake.close_called + + def test_starting_twice_raises(self) -> None: + fake = FakeKafkaProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=1, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + producer.start() + try: + producer.start() + assert False, 'should have raised' + except RuntimeError: + pass + producer.wait(timeout=2) + + def test_producer_factory_failure_captured(self) -> None: + def boom(**_: object) -> NoReturn: + raise ConnectionError('no kafka here') + + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=10, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=boom) + producer.start() + producer.wait(timeout=2) + assert isinstance(producer.error, ConnectionError) + assert producer.produced == {} + + def test_run_id_helper(self) -> None: + run_id = ProducerConfig.make_run_id() + assert len(run_id) == 8 + assert all(c in '0123456789abcdef' for c in run_id) + + def test_close_runs_even_if_flush_raises(self) -> None: + # Regression: flush() and close() were in one try, so a flush() + # exception would skip close() and leak the socket. + class FlushFailsProducer(FakeKafkaProducer): + def flush(self, timeout: float = 0) -> None: + raise RuntimeError('flush exploded') + + fake = FlushFailsProducer() + config = ProducerConfig( + bootstrap_servers=['ignored'], + topic='topic', + events=2, + rate_per_second=1000.0, + run_id='abcdef12', + ) + producer = Producer(config, producer_factory=lambda **kw: fake) + producer.start() + producer.wait(timeout=5) + assert fake.close_called + assert isinstance(producer.error, RuntimeError) + assert 'flush exploded' in str(producer.error)