Skip to content

feat(mesh): wire TreeSnapshot into mesh sync protocol (Phase 3) #983

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
slin1237 wants to merge 7 commits into from
Closed

feat(mesh): wire TreeSnapshot into mesh sync protocol (Phase 3) #983

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
bed6cd4
feat(mesh): wire TreeSnapshot into mesh sync protocol (Phase 3)
slin1237 Mar 31, 2026
96f98f6
fix(mesh): address review feedback for tree snapshot sync
slin1237 Mar 31, 2026
dbc41f2
fix(mesh): split tree batches from policy to unblock delta delivery
slin1237 Mar 31, 2026
053c369
fix(mesh): remove full-state fallback from collector
slin1237 Mar 31, 2026
fcde807
chore: remove unused RadixTree import from incremental.rs
slin1237 Mar 31, 2026
02af584
fix(mesh): cap delta ops to prevent unbounded growth
slin1237 Mar 31, 2026
cb37fac
feat(mesh): replace raw op sync with snapshot-based tree sync
slin1237 Mar 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions crates/mesh/src/controller.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -651,17 +651,43 @@ impl MeshController {
);
}
}
} else if policy_state.policy_type
== "tree_snapshot"
{
// New compact snapshot format
match kv_index::snapshot::TreeSnapshot::from_bytes(
&policy_state.config,
)
{
Ok(snapshot) => {
sync_manager
.apply_remote_tree_snapshot(
policy_state
.model_id
.clone(),
snapshot,
policy_state.version,
actor,
);
}
Err(e) => {
log::warn!(
"Failed to deserialize tree snapshot for model {}: {e}",
policy_state.model_id
);
}
}
} else if policy_state.policy_type
== "tree_state"
{
// Full state: replace (backward compatible)
// Old format: backward compatible
match super::tree_ops::TreeState::from_bytes(
&policy_state.config,
)
{
Ok(tree_state) => {
sync_manager
.apply_remote_tree_operation(
.apply_remote_tree_state_compat(
policy_state
.model_id
.clone(),
Expand Down
164 changes: 88 additions & 76 deletions crates/mesh/src/incremental.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use tracing::{debug, trace};
use super::{
service::gossip::StateUpdate,
stores::{AppState, MembershipState, PolicyState, StateStores, StoreType, WorkerState},
tree_ops::{TreeState, TreeStateDelta},
tree_ops::TreeStateDelta,
};

/// Trait for extracting version from state types
Expand Down Expand Up @@ -252,33 +252,47 @@ impl IncrementalUpdateCollector {
let unsent_start = already_sent.min(pending.len());
let unsent_ops = &pending[unsent_start..];
if !unsent_ops.is_empty() {
// Cap ops per delta to keep serialized size
// under the gRPC message limit. Each op can
// be up to ~20 KB (long prompt text), so 200
// ops ≈ 4 MB which fits in 10 MB with overhead.
const MAX_OPS_PER_DELTA: usize = 200;
let capped_ops = if unsent_ops.len() > MAX_OPS_PER_DELTA {
&unsent_ops[..MAX_OPS_PER_DELTA]
} else {
unsent_ops
};
let capped_version = last_sent_version + capped_ops.len() as u64;

@claude claude Bot Mar 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: capped_version is computed as last_sent_version + capped_ops.len() but this assumes a 1:1 mapping between ops and version increments. If the underlying version scheme ever changes (e.g., version bumps by more than 1 per op, or version is set explicitly), this will silently produce an incorrect new_version in the delta, which the receiver uses for its version check (delta.base_version > current_version || current_version >= delta.new_version).

More importantly, when ops are capped, the remaining unsent ops will be picked up on the next collection cycle. But last_sent_version is updated to capped_version (via the StateUpdate's version field). On the next cycle, already_sent = capped_version - base_version, which correctly advances past the already-sent ops only if version increments are exactly 1-per-op. If base_version changes between cycles (e.g., due to a concurrent checkpoint that resets it), the offset calculation could skip or duplicate ops.

Consider using base_version + unsent_start + capped_ops.len() to be explicit about the absolute position, or add a debug assertion that capped_version <= current_version.


let model_id = key.strip_prefix("tree:").unwrap_or(key).to_string();
let delta = TreeStateDelta {
model_id: model_id.clone(),
operations: unsent_ops.to_vec(),
operations: capped_ops.to_vec(),
base_version: last_sent_version,
new_version: current_version,
new_version: capped_version,
};
if let Ok(delta_bytes) = delta.to_bytes() {
let delta_policy = PolicyState {
model_id,
policy_type: "tree_state_delta".to_string(),
config: delta_bytes,
version: current_version,
version: capped_version,
};
if let Ok(serialized) = bincode::serialize(&delta_policy) {
updates.push(StateUpdate {
key: key.clone(),
value: serialized,
version: current_version,
version: capped_version,
actor: self.self_name.clone(),
timestamp,
});
debug!(
"Collected tree delta: {} ({} ops, version: {})",
"Collected tree delta: {} ({}/{} ops, version: {}->{})",
key,
capped_ops.len(),
unsent_ops.len(),
current_version
last_sent_version,
capped_version,
);
sent_delta = true;
emitted_tree_keys.insert(key.clone());
Expand All @@ -288,65 +302,40 @@ impl IncrementalUpdateCollector {
}

if !sent_delta {
// Full-state fallback: build TreeState from
// tree_configs (if checkpointed) + pending ops.
let model_id = key.strip_prefix("tree:").unwrap_or(key).to_string();
// Clone bytes out of the DashMap ref so no guard
// is held while iterating tree_ops_pending.
let config_bytes = self
.stores
.tree_configs
.get(key.as_str())
.map(|r| r.clone());
let mut tree_state = match config_bytes.as_deref() {
Some(bytes) if !bytes.is_empty() => {
match TreeState::from_bytes(bytes) {
Ok(ts) => ts,
Err(_) => {
debug!("Skipping full-state fallback for {} — config corrupted", key);
continue;
}
}
}
_ => TreeState::new(model_id.clone()),
};
for op in &**pending {
tree_state.add_operation(op.clone());
}
let tree_version = tree_state.version;
let full_state = PolicyState {
model_id,
policy_type: "tree_state".to_string(),
config: tree_state.to_bytes().unwrap_or_default(),
version: tree_version,
};
if let Ok(serialized) = bincode::serialize(&full_state) {
updates.push(StateUpdate {
key: key.clone(),
value: serialized,
version: current_version,
actor: self.self_name.clone(),
timestamp,
});
debug!(
"Collected full tree state fallback: {} (version: {})",
key, current_version
);
emitted_tree_keys.insert(key.clone());
}
// No delta available (version gap). Don't send
// full-state — it can be tens of MB with long
// prompts and will exceed the gRPC message limit.
// The next checkpoint cycle will fold pending ops
// into tree_configs, resetting the delta window.
// The peer will receive the snapshot via the
// initial snapshot exchange (ping_server) or
// Phase 2 (tree_configs-only entries) instead.
debug!(
"Skipping full-state fallback for {} (version gap: last_sent={}, current={})",
key, last_sent_version, current_version
);
}
}

// Phase 2: Scan tree_configs for keys not yet emitted
// (e.g., after checkpoint + buffer drain, or remote-only entries).
// tree_configs already stores serialized TreeState bytes, so
// we use them directly as the PolicyState config instead of
// round-tripping through deserialize + re-serialize.
// Phase 2: Scan tree_configs for keys not yet emitted.
// Only emit keys that have NO pending ops — those are
// remote-only entries that can't be covered by deltas.
// Keys with pending ops are already covered by Phase 1
// (delta or full-state fallback).
for entry in &self.stores.tree_configs {
let key = entry.key();

@claude claude Bot Mar 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: Stale comment — "Keys with pending ops are already covered by Phase 1 (delta or full-state fallback)." The full-state fallback was removed in this PR (replaced by the debug!("Skipping full-state fallback...") log at line 314). Should read just "delta" or "delta delivery".

if emitted_tree_keys.contains(key.as_str()) {
continue;
}
// Skip keys that have pending ops — Phase 1 handles them.
if self
.stores
.tree_ops_pending
.get(key.as_str())
.is_some_and(|p| !p.is_empty())
{
continue;
}
let current_version = self.stores.tree_version(key);
let last_sent_version =
last_sent.policy.get(key.as_str()).copied().unwrap_or(0);
Expand All @@ -358,24 +347,18 @@ impl IncrementalUpdateCollector {
if config_bytes.is_empty() {
continue;
}
// Validate the bytes are a valid TreeState and extract
// the version, but use the raw bytes directly as the
// config payload to avoid redundant re-serialization.
let tree_version = match TreeState::from_bytes(&config_bytes) {
Ok(ts) => ts.version,
Err(_) => {
debug!(
"Skipping tree_configs full-state for {} — config corrupted",
key
);
continue;
}
};
if kv_index::snapshot::TreeSnapshot::from_bytes(&config_bytes).is_err() {
debug!(
"Skipping tree_configs full-state for {} — config corrupted",
key
);
continue;
}
let full_state = PolicyState {
model_id,
policy_type: "tree_state".to_string(),
policy_type: "tree_snapshot".to_string(),
config: config_bytes,
version: tree_version,
version: current_version,
};
if let Ok(serialized) = bincode::serialize(&full_state) {
updates.push(StateUpdate {
Expand All @@ -386,7 +369,7 @@ impl IncrementalUpdateCollector {
timestamp,
});
debug!(
"Collected tree_configs full state: {} (version: {})",
"Collected tree_configs full snapshot: {} (version: {})",
key, current_version
);
}
Expand Down Expand Up @@ -463,7 +446,10 @@ impl IncrementalUpdateCollector {
updates
}

/// Collect all incremental updates across all stores
/// Collect all incremental updates across all stores.
///
/// Tree updates are split into per-key batches so a single oversized
/// snapshot doesn't block delivery of small deltas or policy updates.
Comment on lines +333 to +336

@coderabbitai coderabbitai Bot Mar 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Splitting StoreType::Policy into multiple batches breaks retry semantics.

crates/mesh/src/controller.rs:424-520 calls collector.mark_sent(*store_type, updates) for each returned tuple. After the first successful StoreType::Policy batch, mark_sent() advances both last_scanned.policy and last_scanned.tree; any later policy/tree batch from the same collection that hits size limits or backpressure will not be re-collected until a fresh generation bump. Please defer the generation advance until all policy batches from one collection succeed, or return enough metadata for mark_sent() to know when it is handling the final policy batch.

Also applies to: 492-517

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@crates/mesh/src/incremental.rs` around lines 477 - 480, Splitting
StoreType::Policy into multiple batches causes premature advancement of
generation because collector.mark_sent(*store_type, updates) advances
last_scanned.policy and last_scanned.tree on the first successful batch; modify
the logic in the batch collection/return path (the functions that yield the
(store_type, updates) tuples consumed by controller.rs) so that
generation/last_scanned is only advanced after all policy batches for a single
collection succeed, or alternatively include explicit metadata with each
returned tuple (e.g., is_final_policy_batch or batch_index/total_batches) so
collector.mark_sent can detect the final policy batch before advancing
generations; update collector.mark_sent (and any callers) to accept and act on
that metadata and ensure the same fix is applied to the other affected location
(lines referenced 492-517) to preserve retry semantics.

pub fn collect_all_updates(&self) -> Vec<(StoreType, Vec<StateUpdate>)> {
let mut all_updates = Vec::new();

Expand All @@ -475,7 +461,33 @@ impl IncrementalUpdateCollector {
StoreType::RateLimit,
] {
let updates = self.collect_updates_for_store(store_type);
if !updates.is_empty() {
if updates.is_empty() {
continue;
}

if store_type == StoreType::Policy {
// Split tree entries from policy entries so each gets
// its own batch with independent size checking.
// A 28 MB tree snapshot must not block 1 KB deltas.
let mut policy_updates = Vec::new();
let mut tree_batches: Vec<Vec<StateUpdate>> = Vec::new();

for update in updates {
if update.key.starts_with("tree:") {
// Each tree key gets its own batch
tree_batches.push(vec![update]);
} else {
policy_updates.push(update);
}
}

if !policy_updates.is_empty() {
all_updates.push((StoreType::Policy, policy_updates));
}
for batch in tree_batches {
all_updates.push((StoreType::Policy, batch));
}
} else {
all_updates.push((store_type, updates));
}
}
Expand Down
2 changes: 2 additions & 0 deletions crates/mesh/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ mod tests;

// Re-export commonly used types
pub use crdt_kv::{CrdtOrMap, OperationLog};
// Re-export snapshot types from kv-index for use by subscribers (e.g., cache_aware routing).
pub use kv_index::snapshot::TreeSnapshot;
pub use metrics::init_mesh_metrics;
pub use mtls::{MTLSConfig, MTLSManager, OptionalMTLSManager};
pub use partition::PartitionDetector;
Expand Down
Loading
Loading