Repository navigation
fix(mesh): prevent stale snapshot chunks from mixing across retries - #837
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves a critical snapshot corruption bug where partial data from a failed snapshot transfer could mix with chunks from a new attempt if the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughSnapshot reception buffering was restructured to key by Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 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 |
![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
The pull request effectively addresses a critical snapshot corruption bug by refining the snapshot_state tracking mechanism. By keying the snapshot_state HashMap solely by LocalStoreType and storing the expected_total chunks alongside the received_chunks, the system can now correctly identify and discard stale partial snapshot data when a new snapshot attempt for the same store type is initiated. This prevents the mixing of chunks from different attempts, which was the root cause of the corruption. The changes are well-reasoned and directly solve the problem.
| chunks.clear(); | ||
| *expected = chunk.total_chunks; | ||
| } | ||
| chunks.push(chunk.clone()); |
There was a problem hiding this comment.
The chunk.clone() operation here, and subsequently received_chunks.to_vec() on line 926, involves cloning SnapshotChunks. A SnapshotChunk contains a Vec<StateUpdate>, and each StateUpdate contains a Vec<u8> for its value. Cloning Vec<u8> results in a full memory copy, which can be inefficient for large snapshots. While necessary to ensure ownership for storage and sorting, consider the performance implications if snapshot sizes are expected to be very large. If StateUpdate.value could be bytes::Bytes (which is reference-counted), these clones would be much cheaper.
References
- Cloning
Vec<u8>results in a full memory copy, which can be inefficient for large data. Using reference-counted types likebytes::Bytescan make cloning cheaper.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 990ca17ea9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/mesh/src/ping_server.rs`:
- Around line 902-915: The snapshot assembly logic in snapshot_state (variables:
snapshot_state, store_type, chunks, expected) only clears partial state when
chunk.total_chunks changes, so if a sender restarts with the same total_chunks
stale chunks can mix; modify the block that handles inserting into
snapshot_state to also treat a received chunk with chunk.chunk_index == 0 as the
start of a fresh transfer: if chunk.chunk_index == 0 and !chunks.is_empty() then
clear chunks and set *expected = chunk.total_chunks before pushing the new
chunk, ensuring retries that reuse the same total_chunks don't corrupt the
assembly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 04b32461-d7ec-4fa3-a3cc-59a588899377
📒 Files selected for processing (1)
crates/mesh/src/ping_server.rs
990ca17 to
a4a7b7c
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/mesh/src/ping_server.rs`:
- Around line 906-917: The snapshot handling currently rewrites expected total
on every frame and treats a snapshot complete based only on buffer length;
modify the logic in the snapshot_state handling (where
snapshot_state.entry(store_type) returns (chunks, expected) and
chunk.total_chunks / chunk.chunk_index are used) so that expected is set only
when starting a new transfer (e.g., when chunks.is_empty() or when
chunk.chunk_index == 0) and not overwritten on subsequent frames, and before
applying a snapshot verify that the collected chunks contain each index
0..expected-1 exactly once (check for duplicate or missing chunk.chunk_index
values) rather than relying on Vec length; reject or restart transfers if
indexes are out of range, duplicated, or missing.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b4cf9b53-0d34-4209-9fcb-0b98c5375007
📒 Files selected for processing (1)
crates/mesh/src/ping_server.rs
a4a7b7c to
2b4e9f3
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
crates/mesh/src/ping_server.rs (1)
906-931:⚠️ Potential issue | 🟠 MajorKeep
expected_totalstable and validate chunk identity before applying.
Line 916still rewrites the tracked total on every frame, andLine 923still treatslen() == totalas completion. A retried or inconsistent stream like[0,1,1]or[0(total=3),1(total=2)]can therefore still apply a truncated snapshot. Setexpected_totalonly when starting/resetting a transfer, reject duplicate/out-of-rangechunk_indexvalues, and only apply once0..expected_total-1is present exactly once.🛠️ Suggested hardening
let (chunks, expected) = snapshot_state .entry(store_type) .or_insert_with(|| (Vec::new(), chunk.total_chunks)); if chunk.chunk_index == 0 && !chunks.is_empty() { log::info!( "New snapshot transfer for {:?}, discarding {} partial chunks", store_type, chunks.len() ); chunks.clear(); -} -*expected = chunk.total_chunks; -chunks.push(chunk.clone()); + *expected = chunk.total_chunks; +} else if *expected != chunk.total_chunks { + log::warn!( + "Snapshot total changed mid-transfer for {:?} ({} -> {}), resetting buffer", + store_type, + *expected, + chunk.total_chunks + ); + chunks.clear(); + *expected = chunk.total_chunks; +} +if chunk.chunk_index < *expected + && chunks.iter().all(|c| c.chunk_index != chunk.chunk_index) +{ + chunks.push(chunk.clone()); +} // Check if we've received all chunks if let Some((received_chunks, total)) = snapshot_state.get(&store_type) { - if received_chunks.len() as u64 == *total { + let mut sorted_chunks = received_chunks.to_vec(); + sorted_chunks.sort_by_key(|c| c.chunk_index); + let complete = sorted_chunks.len() as u64 == *total + && sorted_chunks.iter().enumerate().all(|(idx, c)| { + c.total_chunks == *total && c.chunk_index == idx as u64 + }); + if complete { // All chunks received, apply snapshot log::info!("All {} chunks received for store {:?}, applying snapshot", total, store_type); - - if let Some(ref stores) = stores { - // Sort chunks by index - let mut sorted_chunks = received_chunks.to_vec(); - sorted_chunks.sort_by_key(|c| c.chunk_index); + if let Some(ref stores) = stores {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/mesh/src/ping_server.rs` around lines 906 - 931, The current logic mutates the tracked total every frame and treats len==total as sufficient; change it so expected total is set only when starting/resetting a transfer (when inserting the entry or when chunk.chunk_index == 0 and you clear chunks), do not overwrite *expected on every incoming chunk; validate each incoming chunk: reject and ignore duplicates (same chunk_index already present) and out-of-range indices (chunk_index >= expected when expected is known), and only apply the snapshot in the completion branch (the code that currently sorts and applies) after verifying the set of received chunk_index values equals exactly 0..expected-1 with no gaps or duplicates. Use the existing snapshot_state entry tuple (Vec<Chunk>, expected_total) and the variables store_type, chunk.chunk_index and chunk.total_chunks to locate and implement these checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@crates/mesh/src/ping_server.rs`:
- Around line 906-931: The current logic mutates the tracked total every frame
and treats len==total as sufficient; change it so expected total is set only
when starting/resetting a transfer (when inserting the entry or when
chunk.chunk_index == 0 and you clear chunks), do not overwrite *expected on
every incoming chunk; validate each incoming chunk: reject and ignore duplicates
(same chunk_index already present) and out-of-range indices (chunk_index >=
expected when expected is known), and only apply the snapshot in the completion
branch (the code that currently sorts and applies) after verifying the set of
received chunk_index values equals exactly 0..expected-1 with no gaps or
duplicates. Use the existing snapshot_state entry tuple (Vec<Chunk>,
expected_total) and the variables store_type, chunk.chunk_index and
chunk.total_chunks to locate and implement these checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3b10875a-a2b1-4874-ab31-24f9fd87aa10
📒 Files selected for processing (1)
crates/mesh/src/ping_server.rs
2b4e9f3 to
69a3089
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/mesh/src/ping_server.rs`:
- Around line 919-940: The contiguous-index check on snapshot chunks (in the
snapshot_state handling around store_type / received_chunks / total) should also
validate that every chunk's declared total_chunks matches the expected total to
catch malformed or corrupted chunks; update the validation (where
sorted_chunks.iter().enumerate().all(...) is computed) to require both
c.chunk_index == i as u64 and c.total_chunks == *total, and if that combined
check fails remove snapshot_state for store_type (same behavior as the current
non-contiguous branch) so inconsistent total_chunks are treated as an invalid
snapshot.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 68399b96-d104-40ce-83bd-03d3fe32b3c6
📒 Files selected for processing (1)
crates/mesh/src/ping_server.rs
| // Check if we've received all chunks with valid indices | ||
| if let Some((received_chunks, total)) = | ||
| snapshot_state.get(&store_type) | ||
| { | ||
| if received_chunks.len() as u64 == *total { | ||
| // Verify all indices 0..total are present (no duplicates/gaps) | ||
| let mut sorted_chunks = received_chunks.to_vec(); | ||
| sorted_chunks.sort_by_key(|c| c.chunk_index); | ||
| let indices_valid = sorted_chunks.iter().enumerate().all( | ||
| |(i, c)| c.chunk_index == i as u64, | ||
| ); | ||
| if !indices_valid { | ||
| log::warn!( | ||
| "Snapshot for {:?} has {} chunks but indices are not contiguous 0..{}, discarding", | ||
| store_type, sorted_chunks.len(), total | ||
| ); | ||
| snapshot_state.remove(&store_type); | ||
| continue; | ||
| } | ||
|
|
||
| log::info!("All {} chunks received for store {:?}, applying snapshot", | ||
| chunk.total_chunks, store_type); | ||
| total, store_type); |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Good addition of contiguous index validation.
The validation that indices are contiguous 0..total correctly catches scenarios like duplicate chunks being pushed. The sort-then-enumerate approach is clean.
Consider also verifying that all chunks agree on total_chunks during the completion check to further harden against malformed/corrupted data:
let indices_valid = sorted_chunks.iter().enumerate().all(
|(i, c)| c.chunk_index == i as u64 && c.total_chunks == *total,
);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/mesh/src/ping_server.rs` around lines 919 - 940, The contiguous-index
check on snapshot chunks (in the snapshot_state handling around store_type /
received_chunks / total) should also validate that every chunk's declared
total_chunks matches the expected total to catch malformed or corrupted chunks;
update the validation (where sorted_chunks.iter().enumerate().all(...) is
computed) to require both c.chunk_index == i as u64 and c.total_chunks ==
*total, and if that combined check fails remove snapshot_state for store_type
(same behavior as the current non-contiguous branch) so inconsistent
total_chunks are treated as an invalid snapshot.
The snapshot_state HashMap keyed by (store_type, total_chunks) could mix chunks from different snapshot attempts if a peer disconnected mid-transfer and reconnected. If both attempts had the same total_chunks value, old partial chunks would mix with new ones, producing corrupted state. What changed: - crates/mesh/src/ping_server.rs: key snapshot_state by store_type only (not total_chunks). When total_chunks changes for a store (new snapshot attempt), discard the old partial chunks. This prevents stale chunk mixing across reconnections. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
69a3089 to
aaea31d
Compare
1 participant
Summary
Fixes a snapshot corruption bug where partial chunks from a failed snapshot transfer could mix with chunks from a new attempt.
Problem
The
snapshot_stateHashMap was keyed by(store_type, total_chunks). If a peer disconnected mid-snapshot and reconnected, and the new snapshot had the sametotal_chunksvalue, old partial chunks from the previous attempt would be mixed with new ones — producing corrupted state on the receiving node.What changed
snapshot_statebystore_typeonly (nottotal_chunks). Trackexpected_totalalongside the chunk vector. Whentotal_chunkschanges for a store, discard old partial chunks and start fresh.Test plan
cargo test -p smg-mesh— 152 passcargo clippy -p smg-mesh --all-targets -- -D warnings— cleanSummary by CodeRabbit