Add stress.producer — synthetic Kafka producer for the stress harness - #329
Merged
Conversation
julietshen
requested review from
a team,
EXBreder,
ayubun,
haileyok and
vinaysrao1
as code owners
June 1, 2026 13:35
Contributor
|
Warning Review limit reached
More reviews will be available in 16 minutes and 59 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughAdds a new ChangesStress Synthetic Producer
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
3 tasks
Contributor
|
@coderabbitai review |
Contributor
✅ Action performedReview finished.
|
reitblatt
approved these changes
Jun 16, 2026
reitblatt
left a comment
reitblatt
left a comment
Contributor
There was a problem hiding this comment.
LGTM, just take the github-code-quality suggestion on the empty exception block.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@osprey_worker/src/osprey/worker/stress/producer.py`:
- Line 22: Add `kafka-python` as a dependency to the dependencies list in the
osprey_worker package's pyproject.toml file. This ensures that the kafka imports
used throughout the osprey_worker modules (including the KafkaProducer import in
the stress/producer.py module) are properly declared as package requirements,
allowing the osprey_worker package to function correctly when installed outside
the workspace context.
- Around line 42-53: The ProducerConfig class currently uses Python's dataclass
decorator instead of Pydantic's BaseModel, which is inconsistent with the
codebase guidelines and other config classes like FeatureFlagsConfig,
LabelsConfig, and AnalyticsConfig. Convert ProducerConfig to inherit from
pydantic.BaseModel instead of using the `@dataclass`(frozen=True) decorator.
Remove the dataclass decorator, import BaseModel from pydantic, and update the
class definition accordingly. The field definitions and the make_run_id() static
method should remain unchanged, as Pydantic will handle the field validation
automatically.
In `@osprey_worker/src/osprey/worker/stress/tests/test_producer.py`:
- Line 4: Replace the overly broad `Any` type annotations with more specific
types to comply with the coding guideline against using `Any`. Update the import
statement to import `NoReturn` instead of `Any`. In the
`FakeKafkaProducer.__init__` method, change the `Any` annotation for the kwargs
parameter to `object` since it accepts arbitrary keyword arguments for API
compatibility. In the `boom()` function, change the return type annotation from
`Any` to `NoReturn` since it always raises immediately, and change any parameter
type annotations from `Any` to `object` for the ignored parameters.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 0d245382-607e-417e-8900-2adebbc67349
📒 Files selected for processing (4)
osprey_worker/src/osprey/worker/stress/__init__.pyosprey_worker/src/osprey/worker/stress/producer.pyosprey_worker/src/osprey/worker/stress/tests/__init__.pyosprey_worker/src/osprey/worker/stress/tests/test_producer.py
Comment on lines
+42
to
+53
| @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] |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, check if pyproject.toml exists and look for Pydantic
find . -name "pyproject.toml" -type fRepository: roostorg/osprey
Length of output: 168
🏁 Script executed:
# Check the producer.py file to verify the code
find . -path "*/osprey/worker/stress/producer.py" -type fRepository: roostorg/osprey
Length of output: 112
🏁 Script executed:
# Search for Pydantic imports across the codebase to understand the pattern
rg "from pydantic|import pydantic" --type py | head -20Repository: roostorg/osprey
Length of output: 2058
🏁 Script executed:
# Search for dataclass usage to see the pattern
rg "`@dataclass`" --type py | head -20Repository: roostorg/osprey
Length of output: 1622
🏁 Script executed:
# Check osprey_worker/pyproject.toml for Pydantic version
cat ./osprey_worker/pyproject.toml | grep -A 5 -B 5 "pydantic"Repository: roostorg/osprey
Length of output: 41
🏁 Script executed:
# Look at the actual producer.py file
cat -n ./osprey_worker/src/osprey/worker/stress/producer.py | head -100Repository: roostorg/osprey
Length of output: 4080
🏁 Script executed:
# Search for other config classes to understand the pattern
rg "class.*Config" --type py -A 3 | grep -E "(class|BaseModel|`@dataclass`)" | head -30Repository: roostorg/osprey
Length of output: 3073
🏁 Script executed:
# Check if there are any exceptions or guidelines about dataclass usage
rg "dataclass|Pydantic" ./coding_guidelines* 2>/dev/null || echo "No coding guidelines file found"Repository: roostorg/osprey
Length of output: 91
🏁 Script executed:
# Check osprey_worker dependencies more broadly
head -100 ./osprey_worker/pyproject.tomlRepository: roostorg/osprey
Length of output: 612
🏁 Script executed:
# Look for how other config classes in the codebase are structured with Pydantic
cat -n ./osprey_worker/src/osprey/engine/stdlib/configs/feature_flags_config.py | head -30Repository: roostorg/osprey
Length of output: 1030
🏁 Script executed:
# Check if ProducerConfig is used elsewhere that might affect conversion
rg "ProducerConfig" --type py -B 2 -A 2Repository: roostorg/osprey
Length of output: 5267
🏁 Script executed:
# Check the root pyproject.toml for pydantic dependency
cat ./pyproject.toml | grep -A 20 "dependencies\|pydantic"Repository: roostorg/osprey
Length of output: 2710
🏁 Script executed:
# Look at the test file to understand ProducerConfig usage
cat -n ./osprey_worker/src/osprey/worker/stress/tests/test_producer.py | head -60Repository: roostorg/osprey
Length of output: 2251
🏁 Script executed:
# Check how ProducerConfig is instantiated in tests - see if frozen behavior matters
cat -n ./osprey_worker/src/osprey/worker/stress/tests/test_producer.py | grep -A 8 "ProducerConfig("Repository: roostorg/osprey
Length of output: 2630
🏁 Script executed:
# Check if there are any other usages of ProducerConfig in the codebase
rg "ProducerConfig" --type py -B 1 -A 3 | tail -50Repository: roostorg/osprey
Length of output: 4014
Convert ProducerConfig to Pydantic BaseModel.
The coding guidelines require using Pydantic for data models throughout the Python codebase. Other config classes in the codebase (e.g., FeatureFlagsConfig, LabelsConfig, AnalyticsConfig) follow this pattern. Converting to Pydantic will add automatic validation at construction time and maintain consistency with the codebase.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@osprey_worker/src/osprey/worker/stress/producer.py` around lines 42 - 53, The
ProducerConfig class currently uses Python's dataclass decorator instead of
Pydantic's BaseModel, which is inconsistent with the codebase guidelines and
other config classes like FeatureFlagsConfig, LabelsConfig, and AnalyticsConfig.
Convert ProducerConfig to inherit from pydantic.BaseModel instead of using the
`@dataclass`(frozen=True) decorator. Remove the dataclass decorator, import
BaseModel from pydantic, and update the class definition accordingly. The field
definitions and the make_run_id() static method should remain unchanged, as
Pydantic will handle the field validation automatically.
Source: Coding guidelines
julietshen
added a commit
that referenced
this pull request
Jun 16, 2026
Addresses two review notes on #329: * github-code-quality + @reitblatt: replace the empty `except: pass` around `flush()`/`close()` in `Producer._run`. Cleanup failures are now recorded in `self._error`, but only if no earlier send-loop error was already captured — preserves "don't raise from cleanup" without silently dropping the failure. * CodeRabbit: replace `Any` with `object` (kwargs we accept but don't touch) and `NoReturn` (always-raises `boom()`) in test_producer.py. Stricter than `Any` without over-constraining.
julietshen
force-pushed
the
add-stress-producer
branch
from
June 16, 2026 20:09
a494e2d to
505c423
Compare
Member
Author
|
Rebased on latest main and addressed the comments in Accepted (both in this commit):
Declining with reasoning:
16/16 producer tests still pass, ruff + mypy clean. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@osprey_worker/src/osprey/worker/stress/producer.py`:
- Around line 151-154: The producer cleanup code in the try block has `flush()`
and `close()` together, which means if `flush()` raises an exception, `close()`
will never be executed and resources will remain open. To fix this, ensure
`close()` always runs by wrapping the `flush()` call in a separate try-except
block or moving `close()` to a finally block. This way, the producer is properly
closed even when `flush()` fails.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: e6ac6f21-ef5d-4ef3-9ec5-dc8156ae6ef8
📒 Files selected for processing (4)
osprey_worker/src/osprey/worker/stress/__init__.pyosprey_worker/src/osprey/worker/stress/producer.pyosprey_worker/src/osprey/worker/stress/tests/__init__.pyosprey_worker/src/osprey/worker/stress/tests/test_producer.py
🚧 Files skipped from review as they are similar to previous changes (1)
- osprey_worker/src/osprey/worker/stress/tests/test_producer.py
Threaded Kafka producer that emits N well-formed Osprey actions to a configurable topic at a configurable rate. Returns the wall-clock send time per action_id so the reporter can later compute end-to-end latency. Key design calls: * Deterministic integer action_ids of the form `(base + run_bucket * 10M + n)`, encoding the run_id and sequence so concurrent runs don't collide and the consumer can filter on the run. * All action_ids stay below 2**53 so JSON consumers (Druid, browsers) don't lose precision. * Rate control via "sleep until next slot" with drift correction: a long stall resets the schedule to now rather than triggering a burst-catchup. * Factory-injected KafkaProducer so unit tests run without live Kafka. Part of the stress harness work for #324, split into smaller PRs per AGENTS.md. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Addresses two review notes on #329: * github-code-quality + @reitblatt: replace the empty `except: pass` around `flush()`/`close()` in `Producer._run`. Cleanup failures are now recorded in `self._error`, but only if no earlier send-loop error was already captured — preserves "don't raise from cleanup" without silently dropping the failure. * CodeRabbit: replace `Any` with `object` (kwargs we accept but don't touch) and `NoReturn` (always-raises `boom()`) in test_producer.py. Stricter than `Any` without over-constraining.
julietshen
force-pushed
the
add-stress-producer
branch
from
June 16, 2026 20:12
505c423 to
49295df
Compare
Addresses CodeRabbit follow-up on #329. The previous fix put flush() and close() in the same try block, so a flush() exception would skip close() and leak the underlying socket / kafka client. Split into a two-step loop: each step's failure is captured into self._error (but only if the send loop above didn't already record a more useful one), and either step's failure does not prevent the other from running. Adds a regression test (test_close_runs_even_if_flush_raises) that pins the new behavior with a subclassed FakeKafkaProducer whose flush() raises.
Member
Author
|
Good catch from coderabbit, fixed in Added a regression test ( |
cassidyjames
approved these changes
Jun 16, 2026
5 tasks
julietshen
added a commit
that referenced
this pull request
Jun 17, 2026
Orchestrates the stress reporter (#328), producer (#329), and consumer (#330) into a CLI that produces synthetic events at a configurable rate, observes their ExecutionResults, and reports drop rate + p50/p95/p99 latency. Exits non-zero on threshold breach so it can gate CI on pipeline health. Subcommands: * `run` — closed-loop synthetic testing. Blocked on #330; the body lazy-imports the consumer and short-circuits with a clear pointer message until that PR merges, then activates. * `measure` — open-loop measurement against an external source. Stub; will activate once #236 (jetstream input stream plugin) lands. Argparse + threshold gates + report dispatch + entry-point registration are all live today, so the CLI structure (~240 lines) is reviewable independently. When #330 merges, the only change needed here is dropping the `try/except ImportError` wrapper around the consumer import — the orchestration body below it already references Consumer/ConsumerConfig. Adds: * `osprey-stress` console script in `osprey_worker/pyproject.toml` * CHANGELOG entry citing the full stack of PRs the harness landed across * 7 unit tests covering arg parsing, the measure stub, and the #330-blocked-on-import path so a regression in either is loud
4 tasks
haileyok
added a commit
that referenced
this pull request
Jun 17, 2026
OspreyEngine allocates a gevent ThreadPool per instance and never releases it, so the test suite (which builds many engines) accumulates idle worker threads. The integration-tests job runs pytest under gevent monkey-patching (python -m gevent.monkey --module pytest); those leftover native worker threads block interpreter shutdown, so the process never exits after pytest reports '1096 passed' and the job hangs until the 30-min timeout. It is timing/environment-sensitive (deadlocks on the CI runner, exits cleanly on faster boxes); #330 (stress.consumer) perturbed scheduling enough to make it deterministic in CI, while #328/#329 happened to win the shutdown race. Kill any surviving gevent ThreadPools in pytest_sessionfinish so the worker threads exit and the process can finalize. Locally this drops the live OS-thread count at exit from ~69 to 2.
haileyok
added a commit
that referenced
this pull request
Jun 17, 2026
OspreyEngine held a gevent ThreadPool(maxsize=1) for its whole lifetime. The test suite builds many engines, so their idle threadpool worker threads accumulated (~67 by the end of a run). Under the gevent-monkey-patched test runner (python -m gevent.monkey --module pytest) those leftover native worker threads block interpreter shutdown, so the process never exits after pytest reports '1096 passed' and the integration-tests job hangs until the 30-min timeout. It is timing/environment-sensitive (deadlocks on the CI runner, exits cleanly on faster boxes); #330 perturbed scheduling enough to make it deterministic while #328/#329 happened to win the shutdown race. Use a short-lived ThreadPool per compilation and kill() it afterwards so an engine never keeps a worker thread alive. Verified in an isolated full-stack run: the live OS-thread count at exit drops from ~69 to 3 with all 1096 tests passing.
julietshen
added a commit
that referenced
this pull request
Jun 17, 2026
Orchestrates the stress reporter (#328), producer (#329), and consumer (#330) into a CLI that produces synthetic events at a configurable rate, observes their ExecutionResults, and reports drop rate + p50/p95/p99 latency. Exits non-zero on threshold breach so it can gate CI on pipeline health. Subcommands: * `run` — closed-loop synthetic testing. Blocked on #330; the body lazy-imports the consumer and short-circuits with a clear pointer message until that PR merges, then activates. * `measure` — open-loop measurement against an external source. Stub; will activate once #236 (jetstream input stream plugin) lands. Argparse + threshold gates + report dispatch + entry-point registration are all live today, so the CLI structure (~240 lines) is reviewable independently. When #330 merges, the only change needed here is dropping the `try/except ImportError` wrapper around the consumer import — the orchestration body below it already references Consumer/ConsumerConfig. Adds: * `osprey-stress` console script in `osprey_worker/pyproject.toml` * CHANGELOG entry citing the full stack of PRs the harness landed across * 7 unit tests covering arg parsing, the measure stub, and the #330-blocked-on-import path so a regression in either is loud
julietshen
added a commit
that referenced
this pull request
Jun 30, 2026
Orchestrates the stress reporter (#328), producer (#329), and consumer (#330) into a CLI that produces synthetic events at a configurable rate, observes their ExecutionResults, and reports drop rate + p50/p95/p99 latency. Exits non-zero on threshold breach so it can gate CI on pipeline health. Subcommands: * `run` — closed-loop synthetic testing. Blocked on #330; the body lazy-imports the consumer and short-circuits with a clear pointer message until that PR merges, then activates. * `measure` — open-loop measurement against an external source. Stub; will activate once #236 (jetstream input stream plugin) lands. Argparse + threshold gates + report dispatch + entry-point registration are all live today, so the CLI structure (~240 lines) is reviewable independently. When #330 merges, the only change needed here is dropping the `try/except ImportError` wrapper around the consumer import — the orchestration body below it already references Consumer/ConsumerConfig. Adds: * `osprey-stress` console script in `osprey_worker/pyproject.toml` * CHANGELOG entry citing the full stack of PRs the harness landed across * 7 unit tests covering arg parsing, the measure stub, and the #330-blocked-on-import path so a regression in either is loud
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Threaded Kafka producer that emits N well-formed Osprey actions to a configurable topic at a configurable rate. Returns the wall-clock send time per
action_idso the reporter (#328) can later compute end-to-end latency.Key design calls:
(base + run_bucket * 10M + n), encoding the run_id and sequence so concurrent runs don't collide and the consumer can filter on the run.KafkaProducerso unit tests run without live Kafka.Part 3 of 5 in the stress-harness split for #324, per AGENTS.md. Sibling PRs:
GetActionIdUDFIndependent of all siblings.
Test plan
FakeKafkaProducer: produces exact count, action_ids unique within a run, action_ids unique across runs, timestamps recorded per event, rate control approximately honored,stop()aborts early, factory failure captured, run_id helper formatuv run ruff check— cleanuv run mypy— clean🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Tests