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 all 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
178 changes: 37 additions & 141 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,6 @@ use tracing::{debug, trace};
use super::{
service::gossip::StateUpdate,
stores::{AppState, MembershipState, PolicyState, StateStores, StoreType, WorkerState},
tree_ops::{TreeState, TreeStateDelta},
};

/// Trait for extracting version from state types
Expand Down Expand Up @@ -222,131 +221,12 @@ impl IncrementalUpdateCollector {
}

// --- Tree keys (driven by atomic tree_versions, not CRDT) ---
// Snapshot-based sync: tree_configs stores compact TreeSnapshot
// bytes written by the checkpoint cycle. No raw operations are
// buffered or sent — only the tree structure (with prefix sharing).
if tree_changed {
let mut emitted_tree_keys = std::collections::HashSet::new();

// Phase 1: Scan tree_ops_pending for keys with pending ops.
for entry in &self.stores.tree_ops_pending {
let key = entry.key();
let pending = entry.value();
if pending.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);

if current_version <= last_sent_version {
continue;
}

// Try to send a delta with only the unsent pending ops
let mut sent_delta = false;
let total_pending = pending.len() as u64;
let base_version = current_version.saturating_sub(total_pending);

if last_sent_version >= base_version {
let already_sent =
last_sent_version.saturating_sub(base_version) as usize;
let unsent_start = already_sent.min(pending.len());
let unsent_ops = &pending[unsent_start..];
if !unsent_ops.is_empty() {
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(),
base_version: last_sent_version,
new_version: current_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,
};
if let Ok(serialized) = bincode::serialize(&delta_policy) {
updates.push(StateUpdate {
key: key.clone(),
value: serialized,
version: current_version,
actor: self.self_name.clone(),
timestamp,
});
debug!(
"Collected tree delta: {} ({} ops, version: {})",
key,
unsent_ops.len(),
current_version
);
sent_delta = true;
emitted_tree_keys.insert(key.clone());
}
}
}
}

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());
}
}
}

// 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.
for entry in &self.stores.tree_configs {
let key = entry.key();
if emitted_tree_keys.contains(key.as_str()) {
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 +238,11 @@ 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;
}
};
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 +253,7 @@ impl IncrementalUpdateCollector {
timestamp,
});
debug!(
"Collected tree_configs full state: {} (version: {})",
"Collected tree snapshot: {} (version: {})",
key, current_version
);
}
Expand Down Expand Up @@ -463,7 +330,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 +345,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
5 changes: 4 additions & 1 deletion 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 All @@ -39,6 +41,7 @@ pub use stores::{
GLOBAL_RATE_LIMIT_COUNTER_KEY, GLOBAL_RATE_LIMIT_KEY,
};
pub use sync::{
MeshSyncManager, OptionalMeshSyncManager, TreeStateSubscriber, WorkerStateSubscriber,
MeshSyncManager, OptionalMeshSyncManager, TreeSnapshotProvider, TreeStateSubscriber,
WorkerStateSubscriber,
};
pub use tree_ops::{TreeInsertOp, TreeKey, TreeOperation, TreeRemoveOp, TreeState};
Loading
Loading