refactor: Nodes shared storage - #746
Merged
Merged
Conversation
…management - Added support for NATS JetStream KV for user synchronization and lifecycle coordination. - Introduced `NatsUserSyncStore` and `NatsNodeLifecycleCoordinator` for managing user claims and lifecycle states. - Implemented compare-and-set (CAS) semantics for user synchronization to ensure exclusive access. - Enhanced node management to utilize NATS for sharing state across multiple workers. - Added leader election mechanism for job scheduling in multi-worker setups. - Updated configuration to include new NATS KV buckets for user sync and lifecycle management. - Refactored existing node management code to integrate with new NATS functionalities. - Added unit tests for user synchronization and lifecycle lease management.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
🚧 Files skipped from review as they are similar to previous changes (6)
WalkthroughThe change adds NATS-backed CAS storage, node lifecycle coordination, scheduler leader election, multi-worker detection, lifecycle-aware node attachment, node synchronization, and all-in-one worker test configuration. ChangesMulti-worker coordination
Estimated code review effort: 4 (Complex) | ~65 minutes Sequence Diagram(s)sequenceDiagram
participant Worker1 as Worker (leader)
participant Worker2 as Worker (follower)
participant NatsKV as NATS KV
participant Scheduler
participant NodeChecker
Worker1->>NatsKV: try_become_leader()
NatsKV-->>Worker1: leadership acquired
Worker1->>Scheduler: start scheduler
Worker1->>NodeChecker: start node health checks
Worker2->>NatsKV: try_become_leader()
NatsKV-->>Worker2: leadership denied
Worker2->>NodeChecker: pause health checks
Node->>Worker1: health check
Worker1->>NatsKV: publish node action
NatsKV->>Worker2: node sync message
Worker2->>Worker2: apply local update
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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 |
|
✅ Action performedReview finished.
|
Contributor
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (8)
app/nats/leader.py (1)
17-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove
WORKER_IDout ofapp.node.nats_memory.
app/nats/leader.pyimportsWORKER_IDfromapp/node/nats_memory.py. That module importsPasarGuardNodeBridgeprotobuf types and creates NATS bridge state. A scheduler-only process therefore loads the node bridge just to read a process identifier string.Define
WORKER_IDin a neutral module such asapp/nats/__init__.pyand import it in both places.🤖 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 `@app/nats/leader.py` around lines 17 - 19, Move the WORKER_ID definition from app.node.nats_memory to a neutral app.nats module, then update both app.nats.leader and the nats_memory module to import it from that shared location. Ensure scheduler-only imports no longer load node bridge types or initialize NATS bridge state just to access the identifier.app/node/nats_memory.py (2)
338-346: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard
_nc.close()during shutdown.
shutdown_bridge_memoryruns as anon_shutdownhook. Ifclose()raises, the module globals stay populated and the remaining shutdown hooks may not run.♻️ Proposed change
async def shutdown_bridge_memory() -> None: global _nc, _user_sync_kv, _lifecycle_kv, _user_sync_store, _lifecycle_coordinator if _nc is not None: - await _nc.close() + with contextlib.suppress(Exception): + await _nc.close() _nc = None🤖 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 `@app/node/nats_memory.py` around lines 338 - 346, Update shutdown_bridge_memory so _nc.close() is guarded against exceptions, ensuring cleanup of all module globals still executes and the shutdown hook does not prevent subsequent hooks from running. Preserve the existing cleanup assignments after the guarded close.
115-142: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEvery CAS retry loop retries 32 times with no backoff or jitter. All nine loops in
NatsUserSyncStoreandNatsNodeLifecycleCoordinatorrepeatkv_get_jsonpluskv_cas_jsonimmediately after a conflict. Each attempt is two NATS round trips against one contended key per node. Concurrent workers retry in lockstep, which amplifies conflicts instead of resolving them, and the loops then raiseRuntimeErroror give up silently. Add a small jittered sleep between attempts and share one retry helper across both classes.
app/node/nats_memory.py#L115-L142: extract the retry loop ofclaim_usersinto a shared helper that applies jittered backoff, and reuse it inenqueue_users,ack_users,requeue_users, andclear.app/node/nats_memory.py#L199-L241: reuse the same helper intry_acquire,release,heartbeat, andupdate_observed.🤖 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 `@app/node/nats_memory.py` around lines 115 - 142, The nine CAS retry loops in NatsUserSyncStore and NatsNodeLifecycleCoordinator retry immediately and independently; extract a shared retry helper that performs jittered backoff between attempts, then use it for claim_users, enqueue_users, ack_users, requeue_users, clear, try_acquire, release, heartbeat, and update_observed. Apply the helper at app/node/nats_memory.py lines 115-142 and 199-241, preserving each operation’s existing success, exhaustion, and result behavior.tests/test_nats_node_memory.py (1)
88-91: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssert the exact observed state instead of a negative condition.
assert state.observed is not LifecycleStatus.BROKENpasses for any value other thanBROKEN.try_acquirewithLifecycleOperation.STOPsetsobservedtoSTOPPING, so an equality assertion pins the fencing behavior more precisely.💚 Proposed change
await coordinator.update_observed("1", LifecycleStatus.BROKEN, expected_epoch=stale.epoch) state = await coordinator.get_state("1") assert state.epoch == newer.epoch - assert state.observed is not LifecycleStatus.BROKEN + assert state.observed is LifecycleStatus.STOPPINGConsider adding a test for
heartbeat, which no test currently covers.🤖 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 `@tests/test_nats_node_memory.py` around lines 88 - 91, Update the assertions in the coordinator state test after the stale BROKEN update to require the exact expected observed state, LifecycleStatus.STOPPING, rather than merely asserting it is not BROKEN. Keep the epoch assertion unchanged; optionally add focused coverage for heartbeat if appropriate.app/operation/node.py (1)
252-269: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd type hints to
_start_or_attach_node.
core,users, andbackend_typehave no annotations. The method returns either theinfo()result of an attach or thestart()result. State the return type so callers know both paths exposenode_versionandcore_version.🤖 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 `@app/operation/node.py` around lines 252 - 269, Add precise type annotations to _start_or_attach_node for core, users, and backend_type, and annotate its return type with the shared result type exposing node_version and core_version for both _attach_if_running and pg_node.start paths. Reuse existing project type aliases or result classes rather than introducing new types.app/app_factory.py (1)
20-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
Role.is_deprecatedinstead of a second deprecation list.
role.pyalready definesis_deprecatedforBACKEND,NODE, andSCHEDULER._DEPRECATED_ROLESduplicates that list, so the two can diverge when a role is added or removed.♻️ Proposed refactor
-_DEPRECATED_ROLES = frozenset({Role.BACKEND, Role.NODE, Role.SCHEDULER})def _warn_deprecated_role(): role = runtime_settings.role - if role not in _DEPRECATED_ROLES: + if not role.is_deprecated: returnThe
Roleimport at line 20 then becomes unnecessary.Also applies to: 131-140
🤖 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 `@app/app_factory.py` around lines 20 - 24, Remove the duplicate _DEPRECATED_ROLES set and use Role.is_deprecated wherever app_factory checks role deprecation, preserving the existing behavior for BACKEND, NODE, and SCHEDULER. Update the affected logic near the later role handling accordingly, and remove the now-unused Role import if no other references remain.app/jobs/node_checker.py (2)
239-249: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAdd jitter or an initial delay to
_interval_loop.Every worker starts these loops at the same moment after startup. All workers then query the database and all nodes in lockstep. A small random initial delay spreads the load.
🤖 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 `@app/jobs/node_checker.py` around lines 239 - 249, Update _interval_loop to wait for a small randomized initial delay before entering its recurring maintenance cycle, using the existing seconds interval as the upper bound or another bounded jitter appropriate to the loop. Keep the current exception logging and periodic sleep behavior unchanged after the initial delay.
306-311: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the duplicated multi-worker shutdown guard.
Line 307 gates registration with
server_settings.workers <= 1.shutdown_nodesrepeats a similar check at line 325 withis_multi_worker() and server_settings.workers > 1. Keep one authority. The registration guard also ignores the split-role case whereis_multi_worker()is true with one worker.Also applies to: 314-319, 322-327
🤖 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 `@app/jobs/node_checker.py` around lines 306 - 311, Consolidate the multi-worker shutdown decision between the registration block and shutdown_nodes. Update the on_shutdown registration condition to exclude is_multi_worker() even when workers is 1, then remove the redundant is_multi_worker()/workers check inside shutdown_nodes so registration is the sole authority.
🤖 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 @.github/workflows/test-database-migrations.yml:
- Around line 156-159: The ROLE assignment is written to GITHUB_ENV too late to
affect make test in both workflow steps. Update the multi-worker all-in-one step
at .github/workflows/test-database-migrations.yml lines 156-159 and the
corresponding test-mariadb step at lines 320-323 to set ROLE for the step via
its env block or invoke make test with ROLE=all-in-one inline; apply the same
direct change at both sites.
In `@app/jobs/node_checker.py`:
- Around line 270-304: Update the needs_job_leader() branch so only
node_health_check is scheduled through the per-worker _interval_loop; move
check_node_limits out of that branch and register it via scheduler.add_job,
preserving its existing interval, coalescing, max_instances, and job ID settings
so it runs only on the leader scheduler.
- Around line 137-148: Update the shared-state bypass in the node-checking flow
around NodeOperation._attach_if_running so LifecycleStatus.HEALTHY is trusted
only when the owning worker’s lease or process is still active. When the
observed lifecycle state is stale or the owner is no longer active, invalidate
or ignore that state and continue to connect_single_node instead of returning.
In `@app/nats/leader.py`:
- Around line 131-150: The _heartbeat_loop function must retry failed lease
renewals for a bounded number of attempts before conceding leadership, while
still stopping immediately when the lease token is invalid or unavailable. On
final leadership loss, invoke the existing scheduler/notification shutdown
callback so running jobs are paused and the notification dispatcher stops,
rather than only setting _is_leader to False; ensure the callback is triggered
for both renewal exhaustion and token mismatch.
- Around line 59-66: Update the exception handler in _parse to use a
parenthesized tuple of exception types, preserving the existing None return
behavior for TypeError, ValueError, KeyError, and json.JSONDecodeError.
- Around line 87-97: Update the leadership acquisition flow around kv.create in
the scheduler startup logic so any create failure that may indicate LEADER_KEY
already exists proceeds to the existing read-and-steal path, rather than
permanently returning false. Preserve the specialized KeyWrongLastSequenceError
handling, and ensure generic create exceptions either detect the existing-key
condition explicitly or fall through to the same steal logic before marking
_is_leader false.
In `@app/node/nats_memory.py`:
- Around line 317-331: Update the failure branch in ensure_bridge_memory after
get_or_create_kv_bucket returns None to close the existing NATS connection _nc
before returning (None, None), using contextlib as indicated by the review.
Preserve the current warning and successful bucket initialization flow.
- Around line 96-109: The enqueue_users method currently writes the entire node
pending queue to one KV document, which can exceed the NATS payload limit and
cause all CAS retries to fail. Update enqueue_users to split users into
payload-safe batches or shard them across multiple keys, ensuring each
kv_cas_json write remains below the bucket limit while preserving per-user
deduplication and retry behavior.
- Around line 243-300: Update release, heartbeat, and update_observed so
exhausting their 32 CAS retries emits a warning-level log instead of silently
returning. Change heartbeat to return the successful CAS result, while
preserving its early returns for missing or mismatched leases, so callers can
detect renewal failure; keep release and update_observed’s existing behavior
after logging exhaustion.
In `@role.py`:
- Around line 26-33: Update the startup path before NATS handlers or queues
register to call require_nats_if_multiworker(is_multi_worker()), ensuring
multi-worker all-in-one configurations fail fast when NATS is disabled. Also
update _warn_deprecated_role() to use runtime_settings.role.is_deprecated
instead of duplicating the deprecated-role check.
---
Nitpick comments:
In `@app/app_factory.py`:
- Around line 20-24: Remove the duplicate _DEPRECATED_ROLES set and use
Role.is_deprecated wherever app_factory checks role deprecation, preserving the
existing behavior for BACKEND, NODE, and SCHEDULER. Update the affected logic
near the later role handling accordingly, and remove the now-unused Role import
if no other references remain.
In `@app/jobs/node_checker.py`:
- Around line 239-249: Update _interval_loop to wait for a small randomized
initial delay before entering its recurring maintenance cycle, using the
existing seconds interval as the upper bound or another bounded jitter
appropriate to the loop. Keep the current exception logging and periodic sleep
behavior unchanged after the initial delay.
- Around line 306-311: Consolidate the multi-worker shutdown decision between
the registration block and shutdown_nodes. Update the on_shutdown registration
condition to exclude is_multi_worker() even when workers is 1, then remove the
redundant is_multi_worker()/workers check inside shutdown_nodes so registration
is the sole authority.
In `@app/nats/leader.py`:
- Around line 17-19: Move the WORKER_ID definition from app.node.nats_memory to
a neutral app.nats module, then update both app.nats.leader and the nats_memory
module to import it from that shared location. Ensure scheduler-only imports no
longer load node bridge types or initialize NATS bridge state just to access the
identifier.
In `@app/node/nats_memory.py`:
- Around line 338-346: Update shutdown_bridge_memory so _nc.close() is guarded
against exceptions, ensuring cleanup of all module globals still executes and
the shutdown hook does not prevent subsequent hooks from running. Preserve the
existing cleanup assignments after the guarded close.
- Around line 115-142: The nine CAS retry loops in NatsUserSyncStore and
NatsNodeLifecycleCoordinator retry immediately and independently; extract a
shared retry helper that performs jittered backoff between attempts, then use it
for claim_users, enqueue_users, ack_users, requeue_users, clear, try_acquire,
release, heartbeat, and update_observed. Apply the helper at
app/node/nats_memory.py lines 115-142 and 199-241, preserving each operation’s
existing success, exhaustion, and result behavior.
In `@app/operation/node.py`:
- Around line 252-269: Add precise type annotations to _start_or_attach_node for
core, users, and backend_type, and annotate its return type with the shared
result type exposing node_version and core_version for both _attach_if_running
and pg_node.start paths. Reuse existing project type aliases or result classes
rather than introducing new types.
In `@tests/test_nats_node_memory.py`:
- Around line 88-91: Update the assertions in the coordinator state test after
the stale BROKEN update to require the exact expected observed state,
LifecycleStatus.STOPPING, rather than merely asserting it is not BROKEN. Keep
the epoch assertion unchanged; optionally add focused coverage for heartbeat if
appropriate.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f38f4f2b-0506-4eda-ad46-27ce9173c5c4
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (18)
.env.example.github/workflows/test-database-migrations.ymlapp/app_factory.pyapp/core/hosts.pyapp/core/manager.pyapp/db/crud/user.pyapp/jobs/node_checker.pyapp/jobs/record_usages.pyapp/nats/__init__.pyapp/nats/kv_cas.pyapp/nats/leader.pyapp/node/__init__.pyapp/node/nats_memory.pyapp/operation/node.pyconfig.pypyproject.tomlrole.pytests/test_nats_node_memory.py
… and corresponding tests
…nd limit checks for leader-only execution
…ging for CAS exhaustion
… and add related tests
Contributor
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (8)
app/nats/kv_cas.py (2)
118-127: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMatch JetStream filter semantics in
MemoryCasKv.keys.The in-memory implementation matches with
f in key, which is substring matching. JetStream applies subject filters. A test that passes here can still fail against a real bucket. Use prefix matching to keep the fake closer to the real backend.♻️ Proposed change
- matched = [key for key in keys if any(f in key for f in filters)] + matched = [key for key in keys if any(key.startswith(f.rstrip(">")) for f in filters)]🤖 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 `@app/nats/kv_cas.py` around lines 118 - 127, Update MemoryCasKv.keys to apply each filter as a prefix match rather than a substring match when building matched, while preserving the existing no-filter and NoKeysError behavior.
52-58: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAdd a short delay between CAS attempts.
kv_put_jsonretries 32 times with no pause. Under contention from several workers on the same key, this produces a tight round-trip loop against NATS and then raisesRuntimeError. Add a small sleep or jitter between attempts to reduce load.🤖 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 `@app/nats/kv_cas.py` around lines 52 - 58, The kv_put_json retry loop currently retries without pausing, creating a tight NATS request loop under contention. Add a small asynchronous delay or jitter between failed kv_cas_json attempts, while preserving the existing 32-attempt limit and successful return behavior.app/operation/node.py (1)
700-705: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider one bulk sync message instead of one message per node.
_connect_nodes_bulk_syncpublishes a separateconnectmessage for every node. Each sibling worker then runs a fullconnect_nodecycle per message, including a database read and a core users query. A restart of all nodes therefore producesN × (workers - 1)connect cycles. A single message that carries the node id list would let each sibling batch the work.🤖 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 `@app/operation/node.py` around lines 700 - 705, Update _connect_nodes_bulk_sync to publish one bulk “connect” synchronization message containing the IDs of all eligible nodes, instead of calling publish_node_sync once per node. Preserve the existing filtering that excludes None, disabled, and limited nodes, and use the corresponding bulk message handling path so sibling workers can process the IDs together.app/nats/router.py (1)
15-17: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
is_multi_worker()instead of duplicating the predicate.
_router_enabledrepeats the body ofis_multi_worker()inapp/nats/__init__.py(lines 14-16). Two copies can drift. Import the shared helper and keep one definition.♻️ Proposed refactor
-from config import nats_settings, runtime_settings, server_settings +from app.nats import is_multi_worker +from config import nats_settings logger = get_logger("nats-router") def _router_enabled() -> bool: - multi_worker = runtime_settings.role.requires_nats or server_settings.workers > 1 - return nats_settings.enabled and multi_worker + return nats_settings.enabled and is_multi_worker()
tests/test_node_manager_sync.pymonkeypatchesapp.nats.router.runtime_settingsandapp.nats.router.server_settings. Update those patch targets toapp.nats.*if you apply this refactor.🤖 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 `@app/nats/router.py` around lines 15 - 17, Update _router_enabled to import and call the shared is_multi_worker() helper from app.nats instead of recreating the runtime_settings/server_settings predicate, while preserving the existing nats_settings.enabled check. If the refactor changes module references, update test_node_manager_sync.py monkeypatch targets from app.nats.router.runtime_settings/server_settings to the corresponding app.nats symbols.tests/test_create_app_nats_guard.py (1)
20-39: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
pytest.warnsover patchingwarnings.warn.The test replaces the global
warnings.warnfunction through the module reference.pytest.warns(DeprecationWarning, match="deprecated")asserts the same behavior without a global patch, andwarnings.catch_warnings()withsimplefilter("error")covers the negative case.🤖 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 `@tests/test_create_app_nats_guard.py` around lines 20 - 39, Update test_warn_deprecated_role_uses_is_deprecated to use pytest.warns(DeprecationWarning, match="deprecated") for the deprecated role assertion instead of patching app_factory.warnings.warn, and use warnings.catch_warnings() with simplefilter("error") to verify the ALL_IN_ONE case emits no warning. Remove the warned list and related monkeypatch while preserving the existing logger suppression.tests/test_node_manager_sync.py (1)
60-87: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for the
connectaction.The tests cover
remove,disconnect, andupsert. Theconnectbranch inapp/node/manager_sync.pyis the most complex path. It skips disabled and limited nodes, handles anupdate_nodefailure, and callsNodeOperation.connect_node. Add cases for a skipped disabled node and for theupdate_nodefailure path. I can generate those tests if you want.🤖 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 `@tests/test_node_manager_sync.py` around lines 60 - 87, Extend the async node-message tests around handle_node_message with connect-action cases: verify disabled or limited nodes are skipped without connecting, and verify an update_node failure follows the existing failure behavior without calling NodeOperation.connect_node. Reuse the current monkeypatch patterns and assert the relevant calls and outcomes.tests/test_nats_leader_steal.py (1)
9-29: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReset the leader module globals after this test.
The test leaves
leader._is_leader = Trueandleader._tokenset. The reset fixture lives intests/test_nats_leader_heartbeat.pyand does not apply to this file. Any later test that readsis_job_leader()can then observe leader state, and the result depends on collection order. Add a fixture that restores the globals.💚 Proposed fix
from app.nats import leader +@pytest.fixture(autouse=True) +def _reset_leader_state(): + yield + leader._is_leader = False + leader._token = None + leader._kv = None + leader._on_leadership_lost = None + + `@pytest.mark.asyncio` async def test_try_become_leader_falls_through_to_steal_after_generic_create_error():🤖 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 `@tests/test_nats_leader_steal.py` around lines 9 - 29, Add test cleanup for the module globals modified by test_try_become_leader_falls_through_to_steal_after_generic_create_error. Define an autouse fixture in tests/test_nats_leader_steal.py that restores leader._is_leader and leader._token after each test, matching the reset behavior from the heartbeat tests so later tests are isolated.app/node/nats_memory.py (1)
337-345: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCorrect the
has_active_leasedocstring or exclude the local worker.The docstring says "another worker still holds an unexpired lifecycle lease". The method returns
Truefor any unexpired lease, including one this worker owns.app/jobs/node_checker.pyLine 150 uses the result to skip a reconnect. If this worker leaked its own lease, the health check keeps waiting for itself until the lease expires. Either compare the leaseworker_id, or update the docstring to state that the method ignores ownership.🤖 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 `@app/node/nats_memory.py` around lines 337 - 345, Update has_active_lease to exclude the current worker’s lease when determining whether another worker has an active lease, comparing lease_data["worker_id"] with the existing local worker identity. Preserve the false result for missing or expired leases, and keep the docstring aligned with this ownership-aware behavior.
🤖 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 `@app/app_factory.py`:
- Around line 96-105: Update `_pause_jobs_on_leadership_lost` to pause the
scheduler with `scheduler.pause()` instead of shutting it down, and add a
periodic leadership re-acquisition flow after concession that retries
`try_become_leader`; when leadership is regained, resume the scheduler and
restore the notification dispatcher as needed.
In `@app/nats/kv_cas.py`:
- Around line 61-66: Update kv_list_keys to call kv.keys with a prefix-specific
filter such as f"{prefix}>", while preserving the existing startswith(prefix)
guard and NoKeysError behavior. The sibling site app/nats/kv_cas.py:118-127
requires no direct change; it is covered by this caller-side fix.
In `@app/nats/leader.py`:
- Around line 163-195: Update _heartbeat_loop so its retry budget cannot outlast
lease_seconds, accounting for HEARTBEAT_RETRY_DELAY and the next
HEARTBEAT_INTERVAL between renewal attempts. Retry promptly without the full
interval, or constrain HEARTBEAT_MAX_RETRIES/delay so _concede_leadership occurs
before the lease expires, while preserving the existing failure logging and
concession behavior.
- Around line 39-41: Update needs_job_leader to use the same multi-process
criteria as is_multi_worker, including roles that require NATS, so split-role
scheduler and node deployments elect a single leader. Reuse the existing
is_multi_worker symbol rather than checking only server_settings.workers, while
preserving the NATS-enabled requirement.
In `@app/nats/router.py`:
- Line 79: Update the exception handler near the router gating logic to use a
parenthesized tuple for the TimeoutError and asyncio.CancelledError types,
ensuring app/nats/router.py compiles and imports successfully.
In `@app/node/manager_sync.py`:
- Around line 48-64: Restructure the action == "connect" branch so GetDB() is
used only to load db_node and the core/user mappings; exit the async with block
before calling NodeOperation.connect_node. Preserve the existing early returns
and invoke connect_node afterward using the data collected from the closed
session.
- Around line 17-21: Update publish_node_sync and the node-message handling flow
to identify the originating worker and prevent it from processing its own sync
message. Include the worker id in the published payload and have
handle_node_message ignore messages whose origin matches the current worker,
preserving processing for messages from other workers and all existing action
behavior.
In `@app/node/nats_memory.py`:
- Around line 176-185: Update the exception path in the claim flow around
kv_cas_json and self._kv.delete so a failed pending-key deletion also removes
claimed_key before continuing. Preserve the existing debug logging and retry
behavior, ensuring the temporary claim is cleaned up when the delete fails.
---
Nitpick comments:
In `@app/nats/kv_cas.py`:
- Around line 118-127: Update MemoryCasKv.keys to apply each filter as a prefix
match rather than a substring match when building matched, while preserving the
existing no-filter and NoKeysError behavior.
- Around line 52-58: The kv_put_json retry loop currently retries without
pausing, creating a tight NATS request loop under contention. Add a small
asynchronous delay or jitter between failed kv_cas_json attempts, while
preserving the existing 32-attempt limit and successful return behavior.
In `@app/nats/router.py`:
- Around line 15-17: Update _router_enabled to import and call the shared
is_multi_worker() helper from app.nats instead of recreating the
runtime_settings/server_settings predicate, while preserving the existing
nats_settings.enabled check. If the refactor changes module references, update
test_node_manager_sync.py monkeypatch targets from
app.nats.router.runtime_settings/server_settings to the corresponding app.nats
symbols.
In `@app/node/nats_memory.py`:
- Around line 337-345: Update has_active_lease to exclude the current worker’s
lease when determining whether another worker has an active lease, comparing
lease_data["worker_id"] with the existing local worker identity. Preserve the
false result for missing or expired leases, and keep the docstring aligned with
this ownership-aware behavior.
In `@app/operation/node.py`:
- Around line 700-705: Update _connect_nodes_bulk_sync to publish one bulk
“connect” synchronization message containing the IDs of all eligible nodes,
instead of calling publish_node_sync once per node. Preserve the existing
filtering that excludes None, disabled, and limited nodes, and use the
corresponding bulk message handling path so sibling workers can process the IDs
together.
In `@tests/test_create_app_nats_guard.py`:
- Around line 20-39: Update test_warn_deprecated_role_uses_is_deprecated to use
pytest.warns(DeprecationWarning, match="deprecated") for the deprecated role
assertion instead of patching app_factory.warnings.warn, and use
warnings.catch_warnings() with simplefilter("error") to verify the ALL_IN_ONE
case emits no warning. Remove the warned list and related monkeypatch while
preserving the existing logger suppression.
In `@tests/test_nats_leader_steal.py`:
- Around line 9-29: Add test cleanup for the module globals modified by
test_try_become_leader_falls_through_to_steal_after_generic_create_error. Define
an autouse fixture in tests/test_nats_leader_steal.py that restores
leader._is_leader and leader._token after each test, matching the reset behavior
from the heartbeat tests so later tests are isolated.
In `@tests/test_node_manager_sync.py`:
- Around line 60-87: Extend the async node-message tests around
handle_node_message with connect-action cases: verify disabled or limited nodes
are skipped without connecting, and verify an update_node failure follows the
existing failure behavior without calling NodeOperation.connect_node. Reuse the
current monkeypatch patterns and assert the relevant calls and outcomes.
🪄 Autofix
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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2bb6c2ba-d12e-41eb-8ba0-b2c9ef073566
📒 Files selected for processing (17)
.github/workflows/test-database-migrations.ymlapp/app_factory.pyapp/jobs/node_checker.pyapp/nats/__init__.pyapp/nats/kv_cas.pyapp/nats/leader.pyapp/nats/message.pyapp/nats/router.pyapp/node/__init__.pyapp/node/manager_sync.pyapp/node/nats_memory.pyapp/operation/node.pytests/test_create_app_nats_guard.pytests/test_nats_leader_heartbeat.pytests/test_nats_leader_steal.pytests/test_nats_node_memory.pytests/test_node_manager_sync.py
🚧 Files skipped from review as they are similar to previous changes (4)
- app/nats/init.py
- tests/test_nats_node_memory.py
- app/jobs/node_checker.py
- app/node/init.py
…s for origin handling
Contributor
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
Free-Guy-IR
pushed a commit
to Free-Guy-IR/panel
that referenced
this pull request
Sep 5, 2026
* feat(nats): Implement NATS-backed user synchronization and lifecycle management - Added support for NATS JetStream KV for user synchronization and lifecycle coordination. - Introduced `NatsUserSyncStore` and `NatsNodeLifecycleCoordinator` for managing user claims and lifecycle states. - Implemented compare-and-set (CAS) semantics for user synchronization to ensure exclusive access. - Enhanced node management to utilize NATS for sharing state across multiple workers. - Added leader election mechanism for job scheduling in multi-worker setups. - Updated configuration to include new NATS KV buckets for user sync and lifecycle management. - Refactored existing node management code to integrate with new NATS functionalities. - Added unit tests for user synchronization and lifecycle lease management. * feat(nats): Enhance logging for CAS operations in kv_cas and leader modules * refactor: Simplify test execution by setting environment variable for all-in-one role * feat(nats): Add active lease tracking to NatsNodeLifecycleCoordinator and corresponding tests * refactor(node_checker): streamline node health check initialization and limit checks for leader-only execution * fix(leader): correct exception handling in _parse function * fix(leader): improve logging for scheduler leader key creation failure * feat(leadership): implement callback for handling leadership loss and pause jobs * feat(nats): add key listing and upsert functionality to CasKv interface * fix(lifecycle): enhance heartbeat method to return status and add logging for CAS exhaustion * fix(nats_memory): handle NATS client closure on KV bucket creation failure * feat(nats): refactor NATS handling for multi-worker support and add tests * feat(node_sync): implement cross-worker node management sync via NATS and add related tests * feat(nats): enhance job leadership handling with reclaim logic and related tests * feat(heartbeat): improve retry logic and add tests for heartbeat behavior * feat(manager_sync): include origin in node sync messages and add tests for origin handling * feat(manager_sync): refactor node message handling and improve error logging
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
1 participant
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.
Type of change
Checklist
Summary by CodeRabbit
New Features
Bug Fixes
Documentation