Repository navigation
refactor(worker): replace healthy AtomicBool with status AtomicU8 - #1101
Conversation
📝 WalkthroughWalkthroughReplaced boolean worker health with a WorkerStatus-backed AtomicU8 across worker implementations and FFI; added Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean mechanical refactor. The AtomicBool→AtomicU8 swap, new trait API (status()/set_status()), and backward-compatible shims are all correct. The set_healthy(false) guard against Pending→NotReady is a sound forward-looking choice for PR 6b. Memory orderings preserved, metrics updates correct, all callers compatible. No issues found.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors worker health management by replacing the boolean healthy flag with a more granular WorkerStatus enum stored as an AtomicU8. The Worker trait and its implementations, including GrpcWorker and BasicWorker, have been updated to include status() and set_status() methods, while maintaining compatibility through shims for is_healthy() and set_healthy(). Feedback was provided regarding the memory ordering used in the GrpcWorker implementation within the Go bindings, suggesting a shift from Relaxed to Acquire/Release semantics to ensure proper synchronization and consistency with the core gateway.
| fn is_healthy(&self) -> bool { | ||
| self.healthy.load(Ordering::Relaxed) | ||
| fn status(&self) -> WorkerStatus { | ||
| WorkerStatus::from_u8(self.status.load(Ordering::Relaxed)) |
There was a problem hiding this comment.
The status() method currently uses Ordering::Relaxed for loading the worker status. While this might be acceptable for a simple health flag, using Ordering::Acquire is more consistent with the BasicWorker implementation in the core gateway and ensures that any state changes synchronized via the status update are visible to the thread performing the load. This is particularly important when the status is used by load balancing policies to make routing decisions.
| WorkerStatus::from_u8(self.status.load(Ordering::Relaxed)) | |
| WorkerStatus::from_u8(self.status.load(Ordering::Acquire)) |
References
- To maintain consistency, changes to a feature on one execution path should align with its behavior on related paths or core implementations.
| fn set_healthy(&self, healthy: bool) { | ||
| self.healthy.store(healthy, Ordering::Relaxed); | ||
| fn set_status(&self, status: WorkerStatus) { | ||
| self.status.store(status as u8, Ordering::Relaxed); |
There was a problem hiding this comment.
The set_status() method currently uses Ordering::Relaxed for storing the worker status. To ensure proper synchronization with threads performing an Acquire load (as suggested for the status() method), Ordering::Release should be used. This ensures that all previous memory operations in the current thread are visible to other threads that subsequently load the status with Acquire ordering.
| self.status.store(status as u8, Ordering::Relaxed); | |
| self.status.store(status as u8, Ordering::Release); |
References
- To maintain consistency, changes to a feature on one execution path should align with its behavior on related paths or core implementations.
What changed: - worker.rs: replaced `healthy: Arc<AtomicBool>` with `status: Arc<AtomicU8>` in BasicWorker, storing WorkerStatus enum as u8 - worker.rs: added `status()` and `set_status()` to Worker trait as required methods, with `is_healthy()` and `set_healthy()` becoming default impls - worker.rs: is_healthy() returns status() == Ready (routing predicate) - worker.rs: set_healthy(false) only transitions Ready→NotReady, not Pending→NotReady (a Pending worker hasn't proven itself yet) - worker.rs: set_status() updates metrics via Metrics::set_worker_health() - worker.rs: Debug impl shows WorkerStatus instead of bool - builder.rs: initializes status as WorkerStatus::Ready (same behavior — workers start routable, PR 6b will change to Pending) - bindings/golang/src/policy.rs: same AtomicBool→AtomicU8 swap for GrpcWorker FFI binding, uses is_healthy()/set_healthy() trait defaults Why: The single bool conflates "starting up", "can serve traffic", and "broken". AtomicU8 stores the four-state WorkerStatus enum (Pending/Ready/NotReady/ Failed) from PR 2, enabling the state machine transitions in PR 6b. How: Pure type swap with no behavioral change. Workers still start Ready. is_healthy() still returns true only for Ready workers. set_healthy() maps to set_status() as a compatibility shim. All existing call sites (~40 for is_healthy, ~15 for set_healthy) continue working unchanged through the trait default implementations. Refs: worker-module-deep-refactor plan (PR 6a) Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
13fbfba to
102eeef
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/worker/worker.rs (1)
992-1010: 🧹 Nitpick | 🔵 Trivial
worker_to_info()derives status fromis_healthy()rather thanworker.status().This implementation maps the boolean
is_healthy()back toWorkerStatus, which means it will only ever reportReadyorNotReady— neverPendingorFailed.If the API should expose the full four-state lifecycle (for observability or debugging), consider calling
worker.status()directly:WorkerInfo { ... is_healthy, - status: Some(if is_healthy { - WorkerStatus::Ready - } else { - WorkerStatus::NotReady - }), + status: Some(worker.status()), ... }If this is intentional for backward compatibility (API consumers only expect Ready/NotReady), consider adding a comment explaining the choice.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/worker/worker.rs` around lines 992 - 1010, The current worker_to_info() maps worker.is_healthy() to WorkerStatus, losing Pending/Failed states; change it to use the worker.status() accessor directly for the status field (i.e., set status: Some(worker.status())) so the full lifecycle is exposed, and remove the boolean-to-enum mapping; if the Ready/NotReady behavior was intentional for compatibility, instead add a clear comment above worker_to_info() explaining that mapping to only Ready/NotReady is deliberate and why, and keep the rest of the function (references: worker_to_info, worker.is_healthy, worker.status, WorkerStatus).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@bindings/golang/src/policy.rs`:
- Around line 183-189: The health-check methods grpc_health_check and
http_health_check now return the cached state from is_healthy() which can be
confusing since callers expect an active check; update both functions to include
a short clarifying comment like "FFI workers don't perform active health checks;
report current state" above the Ok(self.is_healthy()) return (or alternatively
rename the methods to reflect "is_healthy" semantics), and keep the existing
behavior if that was intended so readers understand this is a deliberate design
choice.
In `@model_gateway/src/worker/worker.rs`:
- Around line 162-172: The current set_healthy(…) does a non-atomic
read-then-write (status() then set_status()) which has a TOCTOU window; replace
this with an atomic compare-and-swap on the worker status so the transition to
NotReady only occurs if the current status is still Ready. Implement or use a
method like compare_and_set_status(expected: WorkerStatus, new: WorkerStatus)
(or add an atomic field and CAS helper) and call
compare_and_set_status(WorkerStatus::Ready, WorkerStatus::NotReady) inside
set_healthy(false) instead of calling status() then set_status(); keep
WorkerStatus enum values (Ready, NotReady, Pending) for comparisons.
---
Outside diff comments:
In `@model_gateway/src/worker/worker.rs`:
- Around line 992-1010: The current worker_to_info() maps worker.is_healthy() to
WorkerStatus, losing Pending/Failed states; change it to use the worker.status()
accessor directly for the status field (i.e., set status: Some(worker.status()))
so the full lifecycle is exposed, and remove the boolean-to-enum mapping; if the
Ready/NotReady behavior was intentional for compatibility, instead add a clear
comment above worker_to_info() explaining that mapping to only Ready/NotReady is
deliberate and why, and keep the rest of the function (references:
worker_to_info, worker.is_healthy, worker.status, WorkerStatus).
🪄 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: ASSERTIVE
Plan: Pro
Run ID: ba5400bb-e1ae-4856-9bcd-6ecec1f941fe
📒 Files selected for processing (3)
bindings/golang/src/policy.rsmodel_gateway/src/worker/builder.rsmodel_gateway/src/worker/worker.rs
| async fn grpc_health_check(&self) -> WorkerResult<bool> { | ||
| Ok(self.healthy.load(Ordering::Relaxed)) | ||
| Ok(self.is_healthy()) | ||
| } | ||
|
|
||
| async fn http_health_check(&self) -> WorkerResult<bool> { | ||
| Ok(self.healthy.load(Ordering::Relaxed)) | ||
| Ok(self.is_healthy()) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Health check stubs return current health state — verify intent.
Both grpc_health_check and http_health_check now return self.is_healthy() instead of always returning true. This changes the semantics: previously, these would indicate "check succeeded", now they report current health state.
Given the comment "FFI workers don't do their own health checks", the intent seems correct — but returning the current health state from a "health check" method is a bit unusual. The caller expects this to perform a check, not reflect cached state.
Consider clarifying with a comment or renaming the semantics in a follow-up:
async fn grpc_health_check(&self) -> WorkerResult<bool> {
// FFI workers don't perform active health checks; report current state
Ok(self.is_healthy())
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bindings/golang/src/policy.rs` around lines 183 - 189, The health-check
methods grpc_health_check and http_health_check now return the cached state from
is_healthy() which can be confusing since callers expect an active check; update
both functions to include a short clarifying comment like "FFI workers don't
perform active health checks; report current state" above the
Ok(self.is_healthy()) return (or alternatively rename the methods to reflect
"is_healthy" semantics), and keep the existing behavior if that was intended so
readers understand this is a deliberate design choice.
| fn set_healthy(&self, healthy: bool) { | ||
| if healthy { | ||
| self.set_status(WorkerStatus::Ready); | ||
| } else { | ||
| // Only transition to NotReady if currently Ready. | ||
| // Don't transition Pending→NotReady (hasn't proven itself). | ||
| if self.status() == WorkerStatus::Ready { | ||
| self.set_status(WorkerStatus::NotReady); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Minor TOCTOU window in set_healthy(false) — acceptable for current usage.
The read-then-write pattern between self.status() and self.set_status() has a theoretical race window. However:
- Health checks run sequentially per worker
- The guard only prevents
Pending → NotReady - Worst case: a concurrent
set_healthy(true)wins, which is benign
If future usage requires stricter atomicity, consider a compare-and-swap pattern. For now, this is fine.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/worker/worker.rs` around lines 162 - 172, The current
set_healthy(…) does a non-atomic read-then-write (status() then set_status())
which has a TOCTOU window; replace this with an atomic compare-and-swap on the
worker status so the transition to NotReady only occurs if the current status is
still Ready. Implement or use a method like compare_and_set_status(expected:
WorkerStatus, new: WorkerStatus) (or add an atomic field and CAS helper) and
call compare_and_set_status(WorkerStatus::Ready, WorkerStatus::NotReady) inside
set_healthy(false) instead of calling status() then set_status(); keep
WorkerStatus enum values (Ready, NotReady, Pending) for comparisons.
1 participant
Description
Problem
Worker health is stored as a single
Arc<AtomicBool>— can only represent healthy/unhealthy. TheWorkerStatusenum (PR #1093) defines four states (Pending, Ready, NotReady, Failed) but has no storage backing.Solution
Replace
Arc<AtomicBool>withArc<AtomicU8>storingWorkerStatusas u8. Addstatus()/set_status()to the Worker trait. Makeis_healthy()/set_healthy()default implementations that delegate to the new methods.Changes
worker.rs:healthy: Arc<AtomicBool>→status: Arc<AtomicU8>. Addedstatus()andset_status()as required trait methods.is_healthy()andset_healthy()become default impls.set_healthy(false)guards Pending→NotReady (only Ready→NotReady is valid).builder.rs: Initializes withWorkerStatus::Ready(preserves current behavior).bindings/golang/src/policy.rs: Same swap forGrpcWorkerFFI binding.Design notes
is_healthy()=status() == Ready— it's a routing predicate, not a health assessment. A Pending worker isn't unhealthy, just unverified.set_healthy(false)only transitionsReady→NotReady, neverPending→NotReady. A worker that hasn't proven itself shouldn't be marked as "was ready, now failing".Ready— PR 6b will change this toPendingfor health-checked workers.is_healthy()and ~15set_healthy()call sites continue working unchanged through trait defaults.Test Plan
cargo check --package smg --package smg-golang— cleancargo test --package smg --lib worker::— 133 tests passcargo test --package smg --test api_tests— 100 integration tests passRefs: worker-module-deep-refactor plan (PR 6a)
Checklist
Summary by CodeRabbit