feat: velo::sync — PendingMap and CloseSignal session primitives (0.4.2) - #45
Open
ryanolson wants to merge 3 commits into
Open
feat: velo::sync — PendingMap and CloseSignal session primitives (0.4.2)#45ryanolson wants to merge 3 commits into
ryanolson wants to merge 3 commits into
Conversation
PendingMap<K, V>: caller-keyed pending-operation map with atomic
drain-on-close. One mutex over enum { Open(map) | Closed(reason) } makes
the insert-after-drain race structurally unrepresentable: register and
close serialize on a single critical section, and insertion exists only
in the Open arm. Sends happen after the lock is released (waker
re-entrancy discipline). close() is sync and callable from non-tokio
threads.
CloseSignal: reason-carrying fire-once close signal. Exactly-once gate
is the reason OnceLock (is_closed() flips at reason-set, before the
drain); close order is reason -> on_close subscribers -> token cancel,
so subscribers (e.g. a PendingMap drain) resolve waiters before any
select! arm on the token wins.
Extracted from the kvbm v2 session-teardown hardening: replaces the
hand-rolled pending_pulls DashMap + fail_pending_pulls + closed flag +
post-insert re-check pattern that was wrong twice under review.
…nseManager docs Additive-only release (new velo::sync module + re-exports); velo-ext untouched at =0.2.0. ResponseManager stays pub(crate) — boundary rationale lives in the velo::sync module doc.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThis PR adds a new session-scoped sync subsystem exposing CloseSignal and PendingMap, registers integration tests and re-exports the new types at the crate root; it also bumps the workspace version to 0.4.2. ChangesSession-scoped synchronization primitives
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
resolve() returns whether a pending entry was found and removed, not whether the value was delivered — a dropped Waiter (cancelled future) still yields true, since the completion matched a real registration. The previous doc claimed false on dropped Waiter, which the code never did and which would make stray-completion detection misreport that race. Pinned with a regression test.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 `@lib/velo/src/sync/close_signal.rs`:
- Around line 140-147: In CloseSignal::close, wrap the subscriber invocation
loop in std::panic::catch_unwind so you can always call
self.inner.token.cancel() afterwards; if catch_unwind returns Err, call
std::panic::resume_unwind to rethrow the panic after cancelling. Concretely:
replace the direct for f in &subscribers { f(&reason); } with a catch_unwind(||
{ for f in &subscribers { f(&reason); } }), then unconditionally call
self.inner.token.cancel(); finally, if catch_unwind produced Err,
resume_unwind(err). This keeps the existing symbols (CloseSignal::close,
subscribers loop, self.inner.token.cancel()) and ensures cancellation even on
subscriber panic.
In `@lib/velo/src/sync/pending_map.rs`:
- Around line 171-186: The resolve method currently returns true unconditionally
after removing the waiter from the map; change it to reflect the documented
contract by returning the result of the send operation so it returns false if
the waiter was dropped—i.e., in the resolve function (the block that obtains tx
via map.remove and then calls tx.send(Ok(value))), replace the unconditional
"true" return with returning tx.send(...).is_ok() (or equivalent) so resolve
returns false when send fails.
- Around line 152-157: In PendingMap::register, replace the contains_key/insert
pattern on the local HashMap (variable map) with a HashMap::entry() check in the
State::Open branch: use map.entry(key) and match Vacant to insert the oneshot
sender (tx) and return Ok(Waiter { rx }), or match Occupied to return
Err(RegisterError::Occupied); update logic in the State::Open arm so no separate
contains_key call remains and clippy’s map_entry lint is satisfied.
🪄 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: Enterprise
Run ID: fbdd30d0-8e3a-486c-9e65-cd21b598b9c0
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (9)
Cargo.tomllib/velo/Cargo.tomllib/velo/src/lib.rslib/velo/src/messenger/common/responses.rslib/velo/src/sync.rslib/velo/src/sync/close_signal.rslib/velo/src/sync/pending_map.rslib/velo/tests/sync/close_signal.rslib/velo/tests/sync/pending_map.rs
Comment on lines
+140
to
+147
| for f in &subscribers { | ||
| f(&reason); | ||
| } | ||
|
|
||
| // ── Step 3: wake async awaiters ─────────────────────────────────── | ||
| // Token is cancelled *last* so that any task woken by the cancel | ||
| // observes subscriber side-effects that were set in step 2. | ||
| self.inner.token.cancel(); |
There was a problem hiding this comment.
Ensure CloseSignal::close always cancels even if a subscriber panics.
A panic in the subscriber loop can skip self.inner.token.cancel(), leaving cancelled().await potentially pending. Make the cancellation unconditional and rethrow the panic after cancelling.
Suggested fix
+use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, OnceLock};
@@
- for f in &subscribers {
- f(&reason);
- }
-
- // ── Step 3: wake async awaiters ───────────────────────────────────
- // Token is cancelled *last* so that any task woken by the cancel
- // observes subscriber side-effects that were set in step 2.
- self.inner.token.cancel();
+ let callback_result = panic::catch_unwind(AssertUnwindSafe(|| {
+ for f in &subscribers {
+ f(&reason);
+ }
+ }));
+
+ // ── Step 3: wake async awaiters ───────────────────────────────────
+ // Always cancel, even if a callback panicked.
+ self.inner.token.cancel();
+ if let Err(payload) = callback_result {
+ panic::resume_unwind(payload);
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| for f in &subscribers { | |
| f(&reason); | |
| } | |
| // ── Step 3: wake async awaiters ─────────────────────────────────── | |
| // Token is cancelled *last* so that any task woken by the cancel | |
| // observes subscriber side-effects that were set in step 2. | |
| self.inner.token.cancel(); | |
| let callback_result = panic::catch_unwind(AssertUnwindSafe(|| { | |
| for f in &subscribers { | |
| f(&reason); | |
| } | |
| })); | |
| // ── Step 3: wake async awaiters ─────────────────────────────────── | |
| // Always cancel, even if a callback panicked. | |
| self.inner.token.cancel(); | |
| if let Err(payload) = callback_result { | |
| panic::resume_unwind(payload); | |
| } |
🤖 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 `@lib/velo/src/sync/close_signal.rs` around lines 140 - 147, In
CloseSignal::close, wrap the subscriber invocation loop in
std::panic::catch_unwind so you can always call self.inner.token.cancel()
afterwards; if catch_unwind returns Err, call std::panic::resume_unwind to
rethrow the panic after cancelling. Concretely: replace the direct for f in
&subscribers { f(&reason); } with a catch_unwind(|| { for f in &subscribers {
f(&reason); } }), then unconditionally call self.inner.token.cancel(); finally,
if catch_unwind produced Err, resume_unwind(err). This keeps the existing
symbols (CloseSignal::close, subscribers loop, self.inner.token.cancel()) and
ensures cancellation even on subscriber panic.
Comment on lines
+152
to
+157
| if map.contains_key(&key) { | ||
| return Err(RegisterError::Occupied); | ||
| } | ||
| let (tx, rx) = oneshot::channel(); | ||
| map.insert(key, tx); | ||
| Ok(Waiter { rx }) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="lib/velo/src/sync/pending_map.rs"
echo "== File head =="
sed -n '1,120p' "$FILE" | nl -ba | head -n 120
echo
echo "== Around register (lines ~120-220) =="
sed -n '120,240p' "$FILE" | nl -ba
echo
echo "== Search for contains_key/insert patterns =="
rg -n "contains_key|entry\(|insert\(" "$FILE"
echo
echo "== Search for resolve docs/implementation =="
rg -n "resolve|Closed|Waiter|RegisterError|dropped|false|contract" "$FILE"Repository: ai-dynamo/velo
Length of output: 135
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="lib/velo/src/sync/pending_map.rs"
if [ ! -f "$FILE" ]; then
echo "File not found: $FILE" >&2
exit 1
fi
echo "== File head =="
sed -n '1,120p' "$FILE" | awk '{printf "%4d: %s\n", NR, $0}'
echo
echo "== Around register (lines ~120-240) =="
sed -n '120,240p' "$FILE" | awk '{printf "%4d: %s\n", NR+119, $0}'
echo
echo "== Search for contains_key/entry patterns =="
rg -n "contains_key|entry\(|insert\(" "$FILE" || true
echo
echo "== Search for resolve docs/implementation markers =="
rg -n "resolve|Closed|Waiter|RegisterError|dropped|false|contract" "$FILE" || trueRepository: ai-dynamo/velo
Length of output: 17664
Fix clippy map_entry lint in PendingMap::register
lib/velo/src/sync/pending_map.rs lines 152-157 use map.contains_key(&key) followed by map.insert(key, tx); switch to HashMap::entry (Occupied/Vacant) in the State::Open branch to avoid cargo clippy ... -D warnings failures.
🤖 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 `@lib/velo/src/sync/pending_map.rs` around lines 152 - 157, In
PendingMap::register, replace the contains_key/insert pattern on the local
HashMap (variable map) with a HashMap::entry() check in the State::Open branch:
use map.entry(key) and match Vacant to insert the oneshot sender (tx) and return
Ok(Waiter { rx }), or match Occupied to return Err(RegisterError::Occupied);
update logic in the State::Open arm so no separate contains_key call remains and
clippy’s map_entry lint is satisfied.
Source: Coding guidelines
Comment on lines
+171
to
+186
| pub fn resolve(&self, key: &K, value: V) -> bool { | ||
| let tx = { | ||
| let mut guard = self.inner.lock(); | ||
| match &mut *guard { | ||
| State::Closed(_) => return false, | ||
| State::Open(map) => match map.remove(key) { | ||
| Some(tx) => tx, | ||
| None => return false, | ||
| }, | ||
| } | ||
| }; | ||
| // Send outside the lock so the receiver's waker cannot re-enter | ||
| // self.inner — mirroring the drop(guard)-before-send discipline in | ||
| // close(). Ignore send error: the Waiter may have been dropped. | ||
| let _ = tx.send(Ok(value)); | ||
| true |
There was a problem hiding this comment.
Align resolve return value with its documented contract.
Line 164-166 says resolve should return false when the waiter was dropped, but Line 185-186 always returns true after map.remove. Return tx.send(...).is_ok() (or update docs if this is intentional).
Suggested fix
- let _ = tx.send(Ok(value));
- true
+ tx.send(Ok(value)).is_ok()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| pub fn resolve(&self, key: &K, value: V) -> bool { | |
| let tx = { | |
| let mut guard = self.inner.lock(); | |
| match &mut *guard { | |
| State::Closed(_) => return false, | |
| State::Open(map) => match map.remove(key) { | |
| Some(tx) => tx, | |
| None => return false, | |
| }, | |
| } | |
| }; | |
| // Send outside the lock so the receiver's waker cannot re-enter | |
| // self.inner — mirroring the drop(guard)-before-send discipline in | |
| // close(). Ignore send error: the Waiter may have been dropped. | |
| let _ = tx.send(Ok(value)); | |
| true | |
| pub fn resolve(&self, key: &K, value: V) -> bool { | |
| let tx = { | |
| let mut guard = self.inner.lock(); | |
| match &mut *guard { | |
| State::Closed(_) => return false, | |
| State::Open(map) => match map.remove(key) { | |
| Some(tx) => tx, | |
| None => return false, | |
| }, | |
| } | |
| }; | |
| // Send outside the lock so the receiver's waker cannot re-enter | |
| // self.inner — mirroring the drop(guard)-before-send discipline in | |
| // close(). Ignore send error: the Waiter may have been dropped. | |
| tx.send(Ok(value)).is_ok() |
🤖 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 `@lib/velo/src/sync/pending_map.rs` around lines 171 - 186, The resolve method
currently returns true unconditionally after removing the waiter from the map;
change it to reflect the documented contract by returning the result of the send
operation so it returns false if the waiter was dropped—i.e., in the resolve
function (the block that obtains tx via map.remove and then calls
tx.send(Ok(value))), replace the unconditional "true" return with returning
tx.send(...).is_ok() (or equivalent) so resolve returns false when send fails.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
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.
What
New top-level
velo::syncmodule with two session-scoped primitives, extracted from the kvbm v2 session-teardown hardening in dynamo:PendingMap<K, V = ()>— caller-keyed pending-operation map with atomic drain-on-close. Oneparking_lot::Mutexoverenum State { Open(HashMap<K, oneshot::Sender<…>>) | Closed(reason) }: register and close serialize on a single critical section and insertion exists only in theOpenarm, so the insert-after-drain race is unrepresentable, not re-checked. All sends happen after the lock is released.close(reason)is sync, idempotent (first reason wins), and callable from non-tokio threads (PyO3/vLLM close paths).CloseSignal— cloneable, reason-carrying, fire-once close signal. Exactly-once gate is the reasonOnceLock(is_closed()flips at reason-set, before the drain); close order is reason →on_closesubscribers → token cancel, so subscribers (e.g. aPendingMapdrain viasignal.on_close(move |r| p.close(r.clone()))) resolve their waiters before anyselect!arm parked on the token wins.Plus a doc-only cross-reference in
responses.rsand the 0.4.1 → 0.4.2 bump.velo-extuntouched at=0.2.0.Why
Dynamo's
kvbm-enginesession layer (p2p/session/velo.rs) hand-rolls this aspending_pulls: DashMap<u64, oneshot>+ afail_pending_pullsdrain + aclosed: Mutex<bool>flag + a post-insert re-check — a two-lock pattern that was wrong twice under adversarial review. An adversarial design pass confirmed two velo-side changes out of six proposals; this PR is both.Why not generalize
ResponseManager: its registration path spans a tokio Semaphore → arena free-list mutex → per-slot mutexes, so it structurally cannot offer the single-critical-section register/close; it's per-worker shared state where a drain-all would fail every in-flight unary worker-wide; and its manager-mintedResponseIdkeys /Bytespayloads don't fit consumer-chosen keys. Full boundary rationale lives in thevelo::syncmodule doc;ResponseManagerstayspub(crate).Adoption evidence (dynamo follow-up PR, after publish)
Against current
kvbm-refactor:pending_pullsDashMap field +closed: Mutex<bool>→ onePendingMap<u64>field;fail_pending_pulls(~16 lines) deleted entirelypull()'s insert-then-recheck →pending_pulls.register(pull_id).map_err(…)?(atomic check-and-insert)pull()'s three-outcome match →Waiter's two-outcomeResult(RecvErrorfolds intoErr(Closed))*closed.lock() = true; fail_pending_pulls(…)writers (monitor exit +close()) → oneclose_signal.close(reason); theon_closehook drains the map synchronously, still safe on the non-tokio shutdown threadFrame::PullCompletearm →pending_pulls.resolve(&pull_id, ())Verification
--all-features --all-targets -D warnings: clean; fmt clean;cargo machetecleanCloseSignal→PendingMapdrain-before-token-cancel ordering testscripts/check-semver.sh: no breaking changes (additive-only; patch bump sufficient)cargo tree -p velo-ext | grep -c prometheus= 0 (boundary intact)resolve()/cancel()now send after unlock (same discipline asclose());cancel()delivers a distinct"operation cancelled"reason; theon_closeexample doctest moved to a public item so it stays compile-checkeddiscovery_etcd::test_cluster_isolationfails locally from stale etcd keys; unrelatedDeferred (tracked here, no tickets per discussion)
closed_token()exposure — the underlyingAnchorEntry::cancel_tokenmisses four permanent-close paths today (TransportError/deser arms, clean channel close, EOF); complete it firstResponseManager::fail_all_pending(reason)for the messenger's own Gate→Drain→Teardown — needs its own pass on the per-slot generation TOCTOUAnchorAttachRequeststays byte-identicalVeloWorkerClient/ConnectorWorkerClientpending-unary hang class) — orthogonalkvbm-refactor— blocked on publishing 0.4.2🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Tests
Chores