From 079d53edf5b6a360ecf43a92083de7db508e594e Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Sun, 2 Aug 2026 15:05:33 +0800 Subject: [PATCH 01/15] perf: defer per-super-step read_channels deep clone in non-streaming mode --- crates/langgraph-core/src/graph/state.rs | 46 +++++++++++++++++++----- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index 8b53285..c293c55 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -1429,6 +1429,7 @@ impl CompiledStateGraph { let mut step: u64 = 0; let max_steps = config.get_recursion_limit().unwrap_or(self.recursion_limit); let mut last_output = JsonValue::Null; + let mut ran_super_step = false; let mut pending_writes: Vec<(String, String, JsonValue)> = Vec::new(); // Version offset: ensures new trigger writes have strictly higher @@ -1718,18 +1719,25 @@ impl CompiledStateGraph { } // ── Streaming: emit values after writes ────────────────────────── - if let Some(s) = stream { + // In streaming mode, compute `output` here (reusing the emitted + // value when stream_mode="values"). In non-streaming mode defer + // the read_channels deep clone — it clones every output channel's + // value each super-step, the dominant per-step cost at large + // histories — to the two consumers below: the interrupt_after + // return and the loop exit. + let output: JsonValue = if let Some(s) = stream { if s.has(&StreamMode::Values) { let keys = output_channel_keys(&channels); - let _ = - s.tx.send(StreamPart::values(vec![], read_channels(&channels, &keys))) - .await; + let v = read_channels(&channels, &keys); + let _ = s.tx.send(StreamPart::values(vec![], v.clone())).await; + v + } else { + let keys = output_channel_keys(&channels); + read_channels(&channels, &keys) } - } - - // Read output - let keys = output_channel_keys(&channels); - let output = read_channels(&channels, &keys); + } else { + JsonValue::Null + }; if !output.is_null() { last_output = output; } @@ -1738,11 +1746,31 @@ impl CompiledStateGraph { if !self.interrupt_after.is_empty() { let task_names: Vec = tasks.iter().map(|t| t.name.clone()).collect(); if task_names.iter().any(|n| self.interrupt_after.contains(n)) { + // Non-streaming deferred the read; materialize it now. + if stream.is_none() { + let keys = output_channel_keys(&channels); + let output = read_channels(&channels, &keys); + if !output.is_null() { + last_output = output; + } + } return Ok(last_output); } } step += 1; + ran_super_step = true; + } + + // Non-streaming: materialize the final output once at loop exit. If no + // super-step ran (empty graph / fork with no resume), keep the original + // Null result rather than reading the untouched initial state. + if stream.is_none() && ran_super_step { + let keys = output_channel_keys(&channels); + let output = read_channels(&channels, &keys); + if !output.is_null() { + last_output = output; + } } Ok(last_output) From 485d8915acdb46e6ed50165d8c080ee085c5fad5 Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 12:21:52 +0800 Subject: [PATCH 02/15] perf: wire updated_channels candidate-set through prepare_next_tasks Thread the channel set returned by apply_writes into prepare_next_tasks as the candidate set, so super-steps after the first only scan nodes reachable from bumped channels instead of all N nodes. The first super-step keeps a full scan because START-edge and input writes bypass apply_writes. Linear chain of 1000 no-op nodes: 116.8ms -> 46.9ms (~2.5x). Still super-linear because apply_writes itself sweeps all channels each super-step (step-2 max_by over channel_versions, step-6 update([]) notify sweep). Correctness: bench_sequential_chain asserts done==true at 1000 nodes (no truncation); 37 langgraph-core-rs tests pass. --- crates/langgraph-core/src/graph/state.rs | 20 +++++- tests/bench_pregel.rs | 84 ++++++++++++++++++++++-- 2 files changed, 94 insertions(+), 10 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index c293c55..d271114 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -1432,6 +1432,17 @@ impl CompiledStateGraph { let mut ran_super_step = false; let mut pending_writes: Vec<(String, String, JsonValue)> = Vec::new(); + // Candidate-set fast path for prepare_next_tasks. The first super-step + // must scan every node (input writes and START-edge trigger channels are + // written directly, not via apply_writes), so this stays `None` until the + // first super-step's apply_writes returns the channels it bumped. After + // that, a node can only trigger via a channel whose version just moved, + // and apply_writes reports exactly that set — so subsequent super-steps + // only check nodes reachable from `updated` instead of all N nodes + // (O(N) candidates -> O(updated) candidates; a linear chain goes from + // O(N²) total to O(N)). + let mut updated_channels: Option> = None; + // Version offset: ensures new trigger writes have strictly higher // versions than anything the checkpoint has already seen. let version_offset: u64 = if saved_checkpoint_exists { @@ -1511,7 +1522,7 @@ impl CompiledStateGraph { version_offset + step, &mut versions_seen, &trigger_to_nodes, - None, + updated_channels.as_ref(), &checkpoint_id, &pending_writes, &channel_versions, @@ -1684,8 +1695,10 @@ impl CompiledStateGraph { } } - // UPDATE: apply all task writes to channels - apply_writes( + // UPDATE: apply all task writes to channels. The returned set is + // exactly the channels whose version moved this super-step (and are + // available) — the next PLAN phase needs only the nodes they trigger. + let updated = apply_writes( &mut channels, &tasks, &mut versions_seen, @@ -1693,6 +1706,7 @@ impl CompiledStateGraph { &trigger_to_nodes, bump_version, ); + updated_channels = Some(updated); // ── DEBUG: 打印 apply_writes 后 messages channel 状态 ── // { diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index 635c78a..37609cd 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -58,6 +58,13 @@ fn make_message(i: usize) -> JsonValue { }) } +/// (thread_id, checkpoint_ns) -> (checkpoint_id, checkpoint_json, metadata_json, parent_cid) +type StorageEntry = (String, JsonValue, JsonValue, Option); +/// (thread_id, checkpoint_ns) -> the thread's newest checkpoint +type StorageMap = HashMap<(String, String), StorageEntry>; +/// (thread_id, checkpoint_ns, checkpoint_id) -> pending writes (interrupt-only path) +type WritesMap = HashMap<(String, String, String), Vec<(String, String, JsonValue)>>; + /// A checkpoint saver that retains only the newest checkpoint per thread. /// /// `InMemorySaver` keeps every checkpoint forever, so a benchmark that runs `N` @@ -67,13 +74,6 @@ fn make_message(i: usize) -> JsonValue { /// memory is O(latest state). The serde round-trip (to_value on `put`, /// from_value on `get_tuple`) mirrors `InMemorySaver`, so the measured per-step /// cost stays comparable. -/// (thread_id, checkpoint_ns) -> (checkpoint_id, checkpoint_json, metadata_json, parent_cid) -type StorageEntry = (String, JsonValue, JsonValue, Option); -/// (thread_id, checkpoint_ns) -> the thread's newest checkpoint -type StorageMap = HashMap<(String, String), StorageEntry>; -/// (thread_id, checkpoint_ns, checkpoint_id) -> pending writes (interrupt-only path) -type WritesMap = HashMap<(String, String, String), Vec<(String, String, JsonValue)>>; - struct LatestOnlySaver { storage: RwLock, writes: RwLock, @@ -565,6 +565,76 @@ async fn bench_parallel_fanout() { } } +/// Linear chain of `nodes` no-op nodes (START -> n0 -> n1 -> ... -> END), +/// mirroring juncture's `sequential.rs` bench 1:1 so the per-node framework +/// overhead is directly comparable (juncture reference: ~2.3µs/node at 1000, +/// measured 2026-08-02). Each node returns an empty update; the last writes a +/// `done` marker so the chain is verified to have run all `nodes` super-steps — +/// a silent truncation fails the assert, not the timing. A chain is one +/// super-step per node, so the recursion limit is raised above the 25 default. +#[tokio::test(flavor = "multi_thread")] +#[ignore] +async fn bench_sequential_chain() { + const REPS: u32 = 3; + + for nodes in [10usize, 100, 500, 1000] { + let mut channels: HashMap> = HashMap::new(); + channels.insert( + "messages".to_string(), + Box::new(BinaryOperatorAggregate::new("messages", add_messages_ref)) + as Box, + ); + channels.insert("done".to_string(), Box::new(LastValue::new("done"))); + + let mut graph = StateGraph::new(channels); + let names: Vec = (0..nodes).map(|i| format!("node_{i}")).collect(); + for (i, name) in names.iter().enumerate() { + let is_last = i + 1 == nodes; + graph + .add_node( + name.clone(), + move |_input: JsonValue, _config: RunnableConfig| { + let is_last = is_last; + async move { + Ok(if is_last { + json!({"done": true}) + } else { + json!({}) + }) + } + }, + ) + .unwrap(); + } + graph.add_edge(START, names[0].clone()).unwrap(); + for i in 0..nodes - 1 { + graph + .add_edge(names[i].clone(), names[i + 1].clone()) + .unwrap(); + } + graph.add_edge(names[nodes - 1].clone(), END).unwrap(); + let app = graph.compile().unwrap(); + + let config = RunnableConfig::new().with_recursion_limit(100_000); + + let mut best = Duration::MAX; + for _ in 0..REPS { + let t = Instant::now(); + let output = app.ainvoke(&json!({}), &config).await.unwrap(); + best = best.min(t.elapsed()); + assert_eq!( + output.get("done"), + Some(&json!(true)), + "chain was truncated: expected all {nodes} nodes to run" + ); + } + println!( + "sequential chain: {nodes:>4} no-op nodes => {best:?} ({:?}/node)", + best / nodes as u32 + ); + } +} + /// The win case for incremental writes: a large static channel (e.g. embedded /// knowledge base) written once at thread start, then untouched while /// `messages` grows every step. Without delta writes the static channel's From 255d4f8fad64039dfe15e5423ba3e9ab3849cb40 Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 12:22:23 +0800 Subject: [PATCH 03/15] chore: ignore AGENTS.md --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 13fd8e1..1aa61cd 100644 --- a/.gitignore +++ b/.gitignore @@ -110,3 +110,4 @@ crates/langgraph-tracing/frontend/node_modules/* publish.sh publish.ps1 .codebuddy/* +AGENTS.md From e6490dd9c8859fc335edfd258d970d616d0e8cf7 Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 12:39:07 +0800 Subject: [PATCH 04/15] perf: eliminate per-superstep O(#channels) scans in apply_writes Three changes on the chain hot path, together taking a 1000-node no-op chain from 46.9ms -> 4.28ms (~11x; vs the 116.8ms baseline ~27x, now within ~1.8x of juncture): 1. apply_writes takes next_version from the caller instead of scanning every channel version for the max each superstep. run_pregel_inner keeps a running version counter (initialized once; versions only ever grow). 2. The step-6 notify sweep touches only the previous super-step's 'updated' channels instead of every channel. A channel can only be available at step-6 if it was written last superstep (that sweep clears everything else), so the bounded sweep is behaviorally identical. The first superstep and get_state pass None (full sweep). 3. version_gt fast path: equal-length all-digit strings compare lexically, skipping the f64 parse and fixing precision loss past 2^53 for the 32-digit engine versions. get_state's call site keeps its exact old behavior. Correctness: all workspace tests pass except the pre-existing sqlite failure; bench_sequential_chain asserts done==true at 1000 nodes (no truncation); fanout/multi-step/sqlite benches unchanged. --- crates/langgraph-core/src/graph/state.rs | 44 +++++++---- crates/langgraph-core/src/pregel/algo.rs | 99 ++++++++++++++---------- 2 files changed, 85 insertions(+), 58 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index d271114..7b61799 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -665,19 +665,22 @@ impl CompiledStateGraph { } // Apply writes from completed tasks to get final channel state + let next_version = { + let max_version = channel_versions + .values() + .filter_map(|v| v.as_str().and_then(|s| s.parse::().ok())) + .max() + .unwrap_or(0); + JsonValue::String(format!("{:032}", max_version + 1)) + }; apply_writes( &mut channels, &tasks, &mut versions_seen, &mut channel_versions, &trigger_to_nodes, - |current| { - let num = current - .and_then(|v| v.as_str()) - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - JsonValue::String(format!("{:032}", num + 1)) - }, + &next_version, + None, ); // Read channel values @@ -1258,15 +1261,6 @@ fn output_channel_keys(channels: &HashMap>) -> Vec) -> JsonValue { - let num = current - .and_then(|v| v.as_str()) - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - JsonValue::String(format!("{:032}", num + 1)) -} - impl CompiledStateGraph { // ──────────────────────────────────────────────────────────────────────── // Public thin wrappers @@ -1509,6 +1503,16 @@ impl CompiledStateGraph { } } + // Running max channel version, used to derive next_version for apply_writes. + // Computed once here (a single O(#channels) pass after input writes) instead + // of a per-superstep max_by over every channel version. Versions only ever + // increase, so incrementing this counter per superstep stays exact. + let mut running_max_version: u64 = channel_versions + .values() + .filter_map(|v| v.as_str().and_then(|s| s.parse::().ok())) + .max() + .unwrap_or(0); + // ── Super-step loop ────────────────────────────────────────────────── while step < max_steps { @@ -1698,13 +1702,19 @@ impl CompiledStateGraph { // UPDATE: apply all task writes to channels. The returned set is // exactly the channels whose version moved this super-step (and are // available) — the next PLAN phase needs only the nodes they trigger. + running_max_version += 1; + let next_version = JsonValue::String(format!("{:032}", running_max_version)); let updated = apply_writes( &mut channels, &tasks, &mut versions_seen, &mut channel_versions, &trigger_to_nodes, - bump_version, + &next_version, + // Bounds the step-6 notify sweep to the previous super-step's + // updated channels (None on the first super-step = sweep all, + // covering directly-written input/START channels). + updated_channels.as_ref(), ); updated_channels = Some(updated); diff --git a/crates/langgraph-core/src/pregel/algo.rs b/crates/langgraph-core/src/pregel/algo.rs index c8bb690..3013607 100644 --- a/crates/langgraph-core/src/pregel/algo.rs +++ b/crates/langgraph-core/src/pregel/algo.rs @@ -22,6 +22,16 @@ fn as_f64(v: &JsonValue) -> Option { /// e.g. 10 > 9. Falls back to string lexical order for non-numeric /// version schemes (e.g. UUIDs). fn version_gt(a: &JsonValue, b: &JsonValue) -> bool { + // Fast path: engine versions are fixed-width zero-padded decimal strings + // (`format!("{:032}", n)`), for which lexical == numeric ordering. Comparing + // equal-length all-digit strings lexically is exact (avoids f64, which + // cannot represent versions past 2^53) and skips the per-compare f64 parse + // that dominated the max_by in apply_writes. + if let (Some(a_s), Some(b_s)) = (a.as_str(), b.as_str()) { + if a_s.len() == b_s.len() && a_s.bytes().all(|b| b.is_ascii_digit()) { + return a_s > b_s; + } + } if let (Some(an), Some(bn)) = (as_f64(a), as_f64(b)) { return an > bn; } @@ -201,11 +211,22 @@ fn create_scratchpad( /// /// This is the "Update" phase of the BSP cycle. It: /// 1. Updates versions_seen for each task's trigger channels -/// 2. Computes a single global next_version from the max of all channel versions -/// 3. Consumes trigger channels (flushes ephemeral values) and bumps their versions -/// 4. Groups writes by channel, applies them, and bumps versions -/// 5. Notifies un-updated channels of the new superstep (bump_step) -/// 6. Calls finish() on all channels if no trigger channels were updated +/// 2. Consumes trigger channels (flushes ephemeral values) and bumps their versions +/// 3. Groups writes by channel, applies them, and bumps versions +/// 4. Notifies un-updated channels of the new superstep (bump_step) +/// 5. Calls finish() on all channels if no trigger channels were updated +/// +/// `next_version` is supplied by the caller: the BSP loop keeps a running version +/// counter so we avoid a per-superstep max over every channel version (which was +/// O(#channels) on the hot path). All channels bumped in this superstep share it, +/// mirroring Python's behavior. +/// +/// `sweep_candidates` bounds the step-5 notify sweep: only channels that were +/// available at the end of the *previous* super-step can still hold a lingering +/// value now (that sweep clears everything else), so the BSP loop passes the +/// previous super-step's `updated` set and only those channels are touched. +/// `None` sweeps every channel (first super-step, where directly-written input +/// channels may linger, and non-loop callers). /// /// Returns the set of updated channel names. pub fn apply_writes( @@ -214,7 +235,8 @@ pub fn apply_writes( versions_seen: &mut HashMap>, channel_versions: &mut ChannelVersions, trigger_to_nodes: &TriggerToNodes, - get_next_version: impl Fn(Option<&JsonValue>) -> JsonValue, + next_version: &JsonValue, + sweep_candidates: Option<&HashSet>, ) -> HashSet { let mut updated = HashSet::new(); @@ -232,15 +254,6 @@ pub fn apply_writes( } } - // 2. Compute a single global next_version from the max of all channel versions. - // This mirrors Python's behavior: all channels updated in the same superstep - // share the same version "timestamp". - let max_version = channel_versions - .values() - .max_by(|a, b| version_gt_partial(a, b)) - .cloned(); - let next_version = get_next_version(max_version.as_ref()); - // 3. Consume trigger channels (flush ephemeral/topic values). // Filter out RESERVED channels (matching Python behavior). // If consume() returns true (state changed), bump the channel version. @@ -291,12 +304,37 @@ pub fn apply_writes( // 6. Channels that weren't updated in this step are notified of a new step. // This allows ephemeral channels to clear themselves and notify downstream. + // Only the previous super-step's `updated` channels can still be available + // here (last superstep's sweep cleared everything else), so a bounded sweep + // is behaviorally identical to the old O(#channels) scan. `None` = all. if bump_step { - for (chan, ch) in channels.iter() { - if ch.is_available() && !updated.contains(chan) && ch.update(&[]).unwrap_or(false) { - channel_versions.insert(chan.clone(), next_version.clone()); - if ch.is_available() { - updated.insert(chan.clone()); + match sweep_candidates { + Some(candidates) => { + for chan in candidates { + if updated.contains(chan) { + continue; + } + if let Some(ch) = channels.get(chan) { + if ch.is_available() && ch.update(&[]).unwrap_or(false) { + channel_versions.insert(chan.clone(), next_version.clone()); + if ch.is_available() { + updated.insert(chan.clone()); + } + } + } + } + } + None => { + for (chan, ch) in channels.iter() { + if ch.is_available() + && !updated.contains(chan) + && ch.update(&[]).unwrap_or(false) + { + channel_versions.insert(chan.clone(), next_version.clone()); + if ch.is_available() { + updated.insert(chan.clone()); + } + } } } } @@ -318,27 +356,6 @@ pub fn apply_writes( updated } -/// Helper for comparing versions in max_by. -/// -/// Tries numeric comparison first (via f64), falls back to string lexical order. -fn version_gt_partial(a: &JsonValue, b: &JsonValue) -> std::cmp::Ordering { - if let (Some(an), Some(bn)) = (as_f64(a), as_f64(b)) { - return an.partial_cmp(&bn).unwrap_or(std::cmp::Ordering::Equal); - } - // Fallback: string lexical comparison - let a_str = match a { - JsonValue::String(s) => s.as_str(), - JsonValue::Number(n) => return n.to_string().cmp(&b.to_string()), - _ => return std::cmp::Ordering::Equal, - }; - let b_str = match b { - JsonValue::String(s) => s.as_str(), - JsonValue::Number(n) => return a_str.cmp(&n.to_string()), - _ => return std::cmp::Ordering::Equal, - }; - a_str.cmp(b_str) -} - /// Check if we should interrupt before executing the given nodes. pub fn should_interrupt(interrupt_nodes: &HashSet, task_names: &[String]) -> bool { task_names.iter().any(|n| interrupt_nodes.contains(n)) From 42bb230974d74c0d76e5be76c3ec4448d1a156e1 Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 13:14:12 +0800 Subject: [PATCH 05/15] perf: cache pregel_nodes/trigger_to_nodes on CompiledStateGraph build_pregel_nodes + build_trigger_to_nodes are pure functions of the immutable CompiledStateGraph fields, yet were rebuilt on every invoke. Measured ~1.4ms at 1000 no-op nodes (~1/3 of a chain's runtime). Build both once in compile_with() and store on the struct; run_pregel_inner, get_state and get_state_history use the cached refs. PregelNode gains #[derive(Clone)] (all Arc/Vec/String fields). 1000-node no-op chain: 4.32ms -> 2.56ms (~1.7x), ~2.56us/node = juncture parity. --- crates/langgraph-core/src/graph/state.rs | 77 +++++++++++++----------- crates/langgraph-core/src/pregel/read.rs | 1 + 2 files changed, 44 insertions(+), 34 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index 7b61799..8c6f9fd 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -6,6 +6,7 @@ use crate::pregel::algo::{apply_writes, prepare_next_tasks}; use crate::pregel::io::{map_command, map_input, read_channels}; use crate::pregel::{ channels_from_checkpoint, ChannelVersions, PregelExecutableTask, PregelNode, PregelRunner, + TriggerToNodes, }; use crate::runnable::{IntoNodeFunction, Runnable, RunnableError}; use crate::stream::StreamPart; @@ -307,12 +308,27 @@ impl StateGraph { .map(|(k, c)| (k.clone(), c.clone_channel())) .collect(); + // Build the PregelNode specs and trigger index once here — they are pure + // functions of the immutable fields above, so caching them on the struct + // removes the per-invocation rebuild in run_pregel_inner (measured + // ~1.4ms at 1000 nodes, see the `pregel_nodes` field docs). + let pregel_nodes = build_pregel_nodes( + &self.nodes, + &self.edges, + &self.waiting_edges, + &self.branches, + &channels, + ); + let trigger_to_nodes = crate::pregel::build_trigger_to_nodes(&pregel_nodes); + Ok(CompiledStateGraph { nodes: self.nodes.clone(), edges: self.edges.clone(), waiting_edges: self.waiting_edges.clone(), branches: self.branches.clone(), channels, + pregel_nodes, + trigger_to_nodes, checkpointer, cache, store, @@ -439,6 +455,13 @@ pub struct CompiledStateGraph { waiting_edges: HashSet, branches: HashMap>, channels: HashMap>, + /// PregelNode specs, built once at compile time. `build_pregel_nodes` is a + /// pure function of the fields above (all immutable after `compile`), so + /// rebuilding it on every invocation was pure waste — at 1000 nodes it was + /// ~1.4ms/invoke (~1/3 of a no-op chain's runtime). See `run_pregel_inner`. + pregel_nodes: HashMap, + /// Reverse index channel -> [triggered nodes], derived from `pregel_nodes`. + trigger_to_nodes: TriggerToNodes, checkpointer: Option>, #[allow(dead_code)] cache: Option>, @@ -617,15 +640,9 @@ impl CompiledStateGraph { } } - // Build PregelNode specs and prepare next tasks - let pregel_nodes = build_pregel_nodes( - &self.nodes, - &self.edges, - &self.waiting_edges, - &self.branches, - &self.channels, - ); - let trigger_to_nodes = crate::pregel::build_trigger_to_nodes(&pregel_nodes); + // Cached PregelNode specs (built once at compile time) and next tasks + let pregel_nodes = &self.pregel_nodes; + let trigger_to_nodes = &self.trigger_to_nodes; let step = 0u64; let checkpoint_id = format!("{:032}", step); @@ -636,12 +653,12 @@ impl CompiledStateGraph { .unwrap_or_default(); let mut tasks = prepare_next_tasks( - &pregel_nodes, + pregel_nodes, &channels, config, step, &mut versions_seen, - &trigger_to_nodes, + trigger_to_nodes, None, &checkpoint_id, &pending_writes, @@ -678,7 +695,7 @@ impl CompiledStateGraph { &tasks, &mut versions_seen, &mut channel_versions, - &trigger_to_nodes, + trigger_to_nodes, &next_version, None, ); @@ -872,15 +889,9 @@ impl CompiledStateGraph { let mut snapshots = Vec::new(); - // Build PregelNode specs for task preparation - let pregel_nodes = build_pregel_nodes( - &self.nodes, - &self.edges, - &self.waiting_edges, - &self.branches, - &self.channels, - ); - let trigger_to_nodes = crate::pregel::build_trigger_to_nodes(&pregel_nodes); + // Cached PregelNode specs (built once at compile time) + let pregel_nodes = &self.pregel_nodes; + let trigger_to_nodes = &self.trigger_to_nodes; for saved in &tuples { // Reconstruct channels from checkpoint @@ -934,12 +945,12 @@ impl CompiledStateGraph { .unwrap_or_default(); let tasks = prepare_next_tasks( - &pregel_nodes, + pregel_nodes, &channels, &RunnableConfig::new(), 0, &mut versions_seen, - &trigger_to_nodes, + trigger_to_nodes, None, &checkpoint_id, &pending_writes, @@ -1004,6 +1015,8 @@ impl Clone for CompiledStateGraph { waiting_edges: self.waiting_edges.clone(), branches, channels, + pregel_nodes: self.pregel_nodes.clone(), + trigger_to_nodes: self.trigger_to_nodes.clone(), checkpointer: self.checkpointer.clone(), cache: self.cache.clone(), store: self.store.clone(), @@ -1354,14 +1367,10 @@ impl CompiledStateGraph { let mut config = config.clone(); // ── Setup ──────────────────────────────────────────────────────────── - let pregel_nodes = build_pregel_nodes( - &self.nodes, - &self.edges, - &self.waiting_edges, - &self.branches, - &self.channels, - ); - let trigger_to_nodes = crate::pregel::build_trigger_to_nodes(&pregel_nodes); + // Cached at compile time (pure function of immutable graph fields); the + // per-invocation rebuild used to cost ~1.4ms at 1000 nodes. + let pregel_nodes = &self.pregel_nodes; + let trigger_to_nodes = &self.trigger_to_nodes; // Load checkpoint (for resume support) let mut saved_checkpoint_exists = false; @@ -1520,12 +1529,12 @@ impl CompiledStateGraph { // PLAN: determine which nodes to run this step let mut tasks = prepare_next_tasks( - &pregel_nodes, + pregel_nodes, &channels, &config, version_offset + step, &mut versions_seen, - &trigger_to_nodes, + trigger_to_nodes, updated_channels.as_ref(), &checkpoint_id, &pending_writes, @@ -1709,7 +1718,7 @@ impl CompiledStateGraph { &tasks, &mut versions_seen, &mut channel_versions, - &trigger_to_nodes, + trigger_to_nodes, &next_version, // Bounds the step-6 notify sweep to the previous super-step's // updated channels (None on the first super-step = sweep all, diff --git a/crates/langgraph-core/src/pregel/read.rs b/crates/langgraph-core/src/pregel/read.rs index ca538e0..a1b77ec 100644 --- a/crates/langgraph-core/src/pregel/read.rs +++ b/crates/langgraph-core/src/pregel/read.rs @@ -5,6 +5,7 @@ use std::sync::Arc; /// /// This is NOT a Runnable itself — it's a container from which /// `PregelExecutableTask`s are built during each super-step. +#[derive(Clone)] pub struct PregelNode { /// Which channels to read as input. /// If the node reads a single channel, this is `[channel_name]`. From cbb0b6c2f2a8f7812590a9d8334f5e70ed9cebc8 Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 18:52:00 +0800 Subject: [PATCH 06/15] perf: run tasks concurrently via JoinSet in PregelRunner Fan-out super-steps now take ~max(branch time) instead of the sum: - single task executes inline (no spawn overhead on the chain path) - >=2 tasks dispatched through tokio::task::JoinSet - tasks re-sorted to index order on return so StreamMode::Updates emission and apply_writes channel ordering stay deterministic - lowest task-index error/interrupt wins regardless of JoinSet completion order; panics surface as a generic TaskFailed - run_tasks now returns (tasks, result) so the caller keeps completed writes on error/interrupt (interrupt path unchanged) Bench: 2/4/8 branches x 40ms = 54/48/46ms (was 93/185/371ms serial). Chain path unaffected (single-task inline). --- crates/langgraph-core/src/graph/state.rs | 5 +- crates/langgraph-core/src/pregel/runner.rs | 92 +++++++++++++++++++--- 2 files changed, 85 insertions(+), 12 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index 8c6f9fd..d0e715b 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -1528,7 +1528,7 @@ impl CompiledStateGraph { let checkpoint_id = format!("{:032}", version_offset + step); // PLAN: determine which nodes to run this step - let mut tasks = prepare_next_tasks( + let tasks = prepare_next_tasks( pregel_nodes, &channels, &config, @@ -1622,7 +1622,8 @@ impl CompiledStateGraph { })) }; - match runner.run_tasks(&mut tasks).await { + let (tasks, task_result) = runner.run_tasks(tasks).await; + match task_result { Ok(()) => {} Err(crate::pregel::runner::RunnerError::Interrupt { task_id, interrupt }) => { diff --git a/crates/langgraph-core/src/pregel/runner.rs b/crates/langgraph-core/src/pregel/runner.rs index 597a101..2bf91dc 100644 --- a/crates/langgraph-core/src/pregel/runner.rs +++ b/crates/langgraph-core/src/pregel/runner.rs @@ -3,6 +3,7 @@ use crate::config; use crate::runnable::RunnableError; use crate::runtime::{Runtime, StreamWriter}; use std::sync::Arc; +use tokio::task::JoinSet; /// Dispatches tasks for parallel execution using tokio. /// @@ -33,24 +34,95 @@ impl PregelRunner { /// /// Each task's runnable is invoked with its input and config. /// Writes are collected into each task's write buffer. - pub async fn run_tasks(&self, tasks: &mut [PregelExecutableTask]) -> Result<(), RunnerError> { + /// + /// A single task is executed inline to avoid spawn overhead on the common + /// sequential-chain path; multiple tasks are dispatched through a + /// `JoinSet` and run concurrently (fan-out takes ~max branch time instead + /// of the sum of branch times). + /// + /// Returns the tasks (with their write buffers populated) alongside the + /// overall result, since the caller inspects `task.writes` whether the + /// step succeeded or was interrupted. If several tasks fail or interrupt, + /// the lowest-index one wins, mirroring the order a serial runner would + /// have reported them in and keeping the outcome deterministic regardless + /// of `JoinSet` completion order. All tasks run to completion: a task that + /// would come after a failing one in serial order still executes and keeps + /// its writes (it ran in this super-step regardless). + pub async fn run_tasks( + &self, + mut tasks: Vec, + ) -> (Vec, Result<(), RunnerError>) { if tasks.is_empty() { - return Ok(()); + return (tasks, Ok(())); } if tasks.len() == 1 { let task = &mut tasks[0]; - Self::execute_single_task(task, self.runtime.as_ref(), self.stream_writer.clone()) - .await?; - return Ok(()); + if let Err(e) = + Self::execute_task(task, self.runtime.as_ref(), self.stream_writer.clone()).await + { + return (tasks, Err(e)); + } + return (tasks, Ok(())); } - for task in tasks.iter_mut() { - Self::execute_single_task(task, self.runtime.as_ref(), self.stream_writer.clone()) - .await?; + let mut set = JoinSet::new(); + for (idx, mut task) in tasks.into_iter().enumerate() { + let runtime = self.runtime.clone(); + let stream_writer = self.stream_writer.clone(); + set.spawn(async move { + let result = Self::execute_task(&mut task, runtime.as_ref(), stream_writer).await; + (idx, task, result) + }); } - Ok(()) + let mut done: Vec<(usize, PregelExecutableTask)> = Vec::with_capacity(set.len()); + let mut first_error: Option<(usize, RunnerError)> = None; + + while let Some(joined) = set.join_next().await { + match joined { + Ok((idx, task, result)) => { + if let Err(e) = result { + let replaces = match &first_error { + Some((i, _)) => idx < *i, + None => true, + }; + if replaces { + first_error = Some((idx, e)); + } + } + done.push((idx, task)); + } + Err(join_err) => { + // A task whose future panicked. The task itself is lost + // (JoinSet cannot identify it); report a generic failure. + let msg = join_err + .try_into_panic() + .ok() + .and_then(|payload| { + payload.downcast_ref::().cloned().or_else(|| { + payload.downcast_ref::<&str>().map(|s| (*s).to_string()) + }) + }) + .unwrap_or_else(|| "task panicked".to_string()); + if first_error.is_none() { + first_error = Some(( + usize::MAX, + RunnerError::TaskFailed("".to_string(), msg), + )); + } + } + } + } + + // Restore serial order so streaming update emission is deterministic. + done.sort_by_key(|(idx, _)| *idx); + let tasks: Vec = done.into_iter().map(|(_, t)| t).collect(); + + match first_error { + Some((_, e)) => (tasks, Err(e)), + None => (tasks, Ok(())), + } } /// Execute tasks synchronously (blocking). @@ -62,7 +134,7 @@ impl PregelRunner { } /// Execute a single task asynchronously. - async fn execute_single_task( + async fn execute_task( task: &mut PregelExecutableTask, runtime: Option<&Arc>, stream_writer: Option, From 5b9843e97103f284f7e85698176d55b997818e7a Mon Sep 17 00:00:00 2001 From: unknown <3058704216@qq.com> Date: Mon, 3 Aug 2026 19:30:03 +0800 Subject: [PATCH 07/15] perf: eliminate per-step checkpoint serde round-trip and sync/async bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - BaseCheckpointSaver::put takes Checkpoint by value (moves instead of deep-copying on the hot per-superstep save path) - InMemorySaver stores the Checkpoint struct directly — no to_value/ from_value round-trip per step; per-thread latest-pointer makes get/put latest lookups O(1) instead of scanning the whole history - async trait defaults only block_in_place on multi-thread runtimes (fixes current_thread panic); in-memory savers get true async overrides - run_pregel_inner/get_state/get_state_history/update_state call aput/aget_tuple/aput_writes, so the per-superstep save never bridges through block_in_place - channels_from_checkpoint consumes the checkpoint's channel values instead of cloning --- .../src/saver.rs | 6 +- .../langgraph-checkpoint-sqlite/src/saver.rs | 108 ++++--- .../src/checkpoint/base.rs | 58 +++- .../src/checkpoint/memory.rs | 159 ++++++---- .../langgraph-core/src/channels/any_value.rs | 4 +- crates/langgraph-core/src/channels/base.rs | 4 +- crates/langgraph-core/src/channels/binop.rs | 6 +- .../src/channels/ephemeral_value.rs | 4 +- .../langgraph-core/src/channels/last_value.rs | 10 +- .../src/channels/named_barrier_value.rs | 11 +- crates/langgraph-core/src/channels/topic.rs | 6 +- .../src/channels/untracked_value.rs | 2 +- crates/langgraph-core/src/graph/state.rs | 283 ++++++++++-------- crates/langgraph-core/src/pregel/mod.rs | 8 +- examples/sqlite_checkpoint.rs | 6 +- tests/bench_pregel.rs | 90 +++--- 16 files changed, 465 insertions(+), 300 deletions(-) diff --git a/crates/langgraph-checkpoint-postgres/src/saver.rs b/crates/langgraph-checkpoint-postgres/src/saver.rs index 066ada8..2a4b306 100644 --- a/crates/langgraph-checkpoint-postgres/src/saver.rs +++ b/crates/langgraph-checkpoint-postgres/src/saver.rs @@ -335,7 +335,7 @@ impl BaseCheckpointSaver for PostgresSaver { fn put( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result { @@ -437,7 +437,7 @@ impl BaseCheckpointSaver for PostgresSaver { async fn aput( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result { @@ -463,7 +463,7 @@ impl BaseCheckpointSaver for PostgresSaver { } })); - let checkpoint_json = serde_json::to_value(checkpoint) + let checkpoint_json = serde_json::to_value(&checkpoint) .map_err(|e| CheckpointError::Storage(e.to_string()))?; let metadata_json = serde_json::to_value(metadata).map_err(|e| CheckpointError::Storage(e.to_string()))?; diff --git a/crates/langgraph-checkpoint-sqlite/src/saver.rs b/crates/langgraph-checkpoint-sqlite/src/saver.rs index b39b711..73bbf8f 100644 --- a/crates/langgraph-checkpoint-sqlite/src/saver.rs +++ b/crates/langgraph-checkpoint-sqlite/src/saver.rs @@ -23,6 +23,38 @@ use crate::queries::*; /// A dumped blob row: (thread_id, checkpoint_ns, channel, type, checkpoint_id, data) type BlobDumpRow = (String, String, String, String, String, Option>); +/// Serialization view of a `Checkpoint` used for the `checkpoints` row. +/// +/// `channel_values` live in `checkpoint_blobs`, but the row's JSON must stay +/// deserializable as a full `Checkpoint`, so it is emitted as an empty object +/// placeholder — encoding the real values here only to strip them would cost +/// a full-state copy on every super-step. +#[derive(serde::Serialize)] +struct CheckpointRowView<'a> { + v: i64, + id: &'a str, + ts: &'a str, + channel_values: serde_json::Map, + channel_versions: &'a ChannelVersions, + versions_seen: &'a HashMap, + #[serde(skip_serializing_if = "Option::is_none")] + updated_channels: &'a Option>, +} + +impl<'a> From<&'a Checkpoint> for CheckpointRowView<'a> { + fn from(cp: &'a Checkpoint) -> Self { + Self { + v: cp.v, + id: &cp.id, + ts: &cp.ts, + channel_values: serde_json::Map::new(), + channel_versions: &cp.channel_versions, + versions_seen: &cp.versions_seen, + updated_channels: &cp.updated_channels, + } + } +} + /// Async SQLite checkpoint saver using sqlx. /// /// Uses a three-table schema (`checkpoints`, `checkpoint_blobs`, @@ -442,7 +474,7 @@ impl BaseCheckpointSaver for SqliteSaver { fn put( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result { @@ -529,7 +561,7 @@ impl BaseCheckpointSaver for SqliteSaver { async fn aput( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result { @@ -549,20 +581,10 @@ impl BaseCheckpointSaver for SqliteSaver { let next_config = Self::make_config(thread_id, checkpoint_ns, &checkpoint.id); - // Strip channel_values from the JSON checkpoint payload — they live in - // checkpoint_blobs. Blob rows are keyed by (channel, version) and are - // written incrementally (only channels in `new_versions`), so on load - // every channel's value resolves through the version join regardless - // of which checkpoint wrote the blob. - let mut checkpoint_value = serde_json::to_value(checkpoint) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - if let Some(obj) = checkpoint_value.as_object_mut() { - obj.insert( - "channel_values".to_string(), - JsonValue::Object(Default::default()), - ); - } - let checkpoint_text = serde_json::to_string(&checkpoint_value) + // Serialize only the checkpoint's metadata fields — channel_values live + // in checkpoint_blobs, so encoding them here (only to strip them + // afterwards) would cost a full-state copy on every super-step. + let checkpoint_text = serde_json::to_string(&CheckpointRowView::from(&checkpoint)) .map_err(|e| CheckpointError::Storage(e.to_string()))?; // Merge config-level fields (e.g. `langgraph_step`) into the // metadata before persisting, so `list(filter=...)` over those @@ -800,7 +822,10 @@ mod tests { ..Default::default() }; - let next = saver.aput(&cfg, &cp, &metadata, &versions).await.unwrap(); + let next = saver + .aput(&cfg, cp.clone(), &metadata, &versions) + .await + .unwrap(); // The returned config should reference the new checkpoint id let returned_cid = next @@ -830,7 +855,7 @@ mod tests { let (cp, versions) = make_checkpoint(vec![("a", serde_json::json!(1))]); let cfg = config_for("thread-W"); saver - .aput(&cfg, &cp, &CheckpointMetadata::default(), &versions) + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &versions) .await .unwrap(); @@ -870,7 +895,7 @@ mod tests { let (cp, versions) = make_checkpoint(vec![("x", serde_json::json!(i))]); ids.push(cp.id.clone()); saver - .aput(&cfg, &cp, &CheckpointMetadata::default(), &versions) + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &versions) .await .unwrap(); } @@ -900,7 +925,7 @@ mod tests { let (cp, versions) = make_checkpoint(vec![("x", serde_json::json!(1))]); let cfg = config_for("thread-D"); saver - .aput(&cfg, &cp, &CheckpointMetadata::default(), &versions) + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &versions) .await .unwrap(); let cfg_with_id = config_with_id("thread-D", &cp.id); @@ -936,7 +961,12 @@ mod tests { let mut versions1: ChannelVersions = HashMap::new(); versions1.insert("counter".into(), JsonValue::Number(1.into())); saver - .aput(&cfg, &cp1, &CheckpointMetadata::default(), &versions1) + .aput( + &cfg, + cp1.clone(), + &CheckpointMetadata::default(), + &versions1, + ) .await .unwrap(); @@ -951,7 +981,12 @@ mod tests { let mut versions2: ChannelVersions = HashMap::new(); versions2.insert("counter".into(), JsonValue::Number(2.into())); saver - .aput(&cfg, &cp2, &CheckpointMetadata::default(), &versions2) + .aput( + &cfg, + cp2.clone(), + &CheckpointMetadata::default(), + &versions2, + ) .await .unwrap(); @@ -989,7 +1024,12 @@ mod tests { versions1.insert("a".into(), JsonValue::String("1".into())); versions1.insert("b".into(), JsonValue::String("1".into())); let next1 = saver - .aput(&cfg, &cp1, &CheckpointMetadata::default(), &versions1) + .aput( + &cfg, + cp1.clone(), + &CheckpointMetadata::default(), + &versions1, + ) .await .unwrap(); @@ -1006,7 +1046,12 @@ mod tests { let mut versions2: ChannelVersions = HashMap::new(); versions2.insert("a".into(), JsonValue::String("2".into())); saver - .aput(&next1, &cp2, &CheckpointMetadata::default(), &versions2) + .aput( + &next1, + cp2.clone(), + &CheckpointMetadata::default(), + &versions2, + ) .await .unwrap(); @@ -1067,7 +1112,7 @@ mod tests { step: Some(step), ..Default::default() }; - saver.aput(&cfg, &cp, &meta, &vers).await.unwrap(); + saver.aput(&cfg, cp.clone(), &meta, &vers).await.unwrap(); tokio::time::sleep(std::time::Duration::from_millis(2)).await; } @@ -1142,7 +1187,7 @@ mod tests { // Metadata passed in does NOT have step set — it should be filled // from the config. saver - .aput(&cfg, &cp, &CheckpointMetadata::default(), &vers) + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &vers) .await .unwrap(); @@ -1169,12 +1214,7 @@ mod tests { let cp_clone = cp.clone(); let vers_clone = vers.clone(); let put_result = tokio::task::spawn_blocking(move || { - s2.put( - &cfg2, - &cp_clone, - &CheckpointMetadata::default(), - &vers_clone, - ) + s2.put(&cfg2, cp_clone, &CheckpointMetadata::default(), &vers_clone) }) .await .unwrap(); @@ -1196,7 +1236,7 @@ mod tests { let (cp1, vers1) = make_checkpoint(vec![("x", serde_json::json!("a"))]); let cfg = config_for("thread-P"); let next1 = saver - .aput(&cfg, &cp1, &CheckpointMetadata::default(), &vers1) + .aput(&cfg, cp1.clone(), &CheckpointMetadata::default(), &vers1) .await .unwrap(); @@ -1209,7 +1249,7 @@ mod tests { // becomes the parent_checkpoint_id of cp2. let (cp2, vers2) = make_checkpoint(vec![("x", serde_json::json!("b"))]); saver - .aput(&next1, &cp2, &CheckpointMetadata::default(), &vers2) + .aput(&next1, cp2.clone(), &CheckpointMetadata::default(), &vers2) .await .unwrap(); diff --git a/crates/langgraph-checkpoint/src/checkpoint/base.rs b/crates/langgraph-checkpoint/src/checkpoint/base.rs index 6127270..0316c42 100644 --- a/crates/langgraph-checkpoint/src/checkpoint/base.rs +++ b/crates/langgraph-checkpoint/src/checkpoint/base.rs @@ -52,10 +52,13 @@ pub trait BaseCheckpointSaver: Send + Sync { ) -> Result, CheckpointError>; /// Store a checkpoint. + /// + /// Takes ownership of the `Checkpoint` so savers can move the serialized + /// state instead of deep-copying it on the hot per-super-step path. fn put( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result; @@ -89,29 +92,42 @@ pub trait BaseCheckpointSaver: Send + Sync { } // Async mirrors with default implementations + // + // The defaults bridge to the sync methods. `block_in_place` is only valid + // on a multi-thread runtime (it panics on `current_thread`), so when no + // multi-thread runtime is present the sync method is called directly — + // the same behavior a synchronous caller gets today. async fn aget_tuple( &self, config: &RunnableConfig, ) -> Result, CheckpointError> { - let config = config.clone(); - let this = self; - // Use blocking for default impl - tokio::task::block_in_place(|| this.get_tuple(&config)) + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::CurrentThread { + let config = config.clone(); + let this = self; + return tokio::task::block_in_place(|| this.get_tuple(&config)); + } + } + self.get_tuple(config) } async fn aput( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, new_versions: &ChannelVersions, ) -> Result { - let config = config.clone(); - let checkpoint = checkpoint.clone(); - let metadata = metadata.clone(); - let new_versions = new_versions.clone(); - tokio::task::block_in_place(|| self.put(&config, &checkpoint, &metadata, &new_versions)) + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::CurrentThread { + let config = config.clone(); + return tokio::task::block_in_place(move || { + self.put(&config, checkpoint, metadata, new_versions) + }); + } + } + self.put(config, checkpoint, metadata, new_versions) } async fn aput_writes( @@ -121,13 +137,25 @@ pub trait BaseCheckpointSaver: Send + Sync { task_id: String, task_path: String, ) -> Result<(), CheckpointError> { - let config = config.clone(); - tokio::task::block_in_place(|| self.put_writes(&config, &writes, &task_id, &task_path)) + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::CurrentThread { + let config = config.clone(); + return tokio::task::block_in_place(|| { + self.put_writes(&config, &writes, &task_id, &task_path) + }); + } + } + self.put_writes(config, &writes, &task_id, &task_path) } async fn adelete_thread(&self, thread_id: String) -> Result<(), CheckpointError> { - let this = self; - tokio::task::block_in_place(|| this.delete_thread(&thread_id)) + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::CurrentThread { + let this = self; + return tokio::task::block_in_place(|| this.delete_thread(&thread_id)); + } + } + self.delete_thread(&thread_id) } } diff --git a/crates/langgraph-checkpoint/src/checkpoint/memory.rs b/crates/langgraph-checkpoint/src/checkpoint/memory.rs index fbb17da..752589f 100644 --- a/crates/langgraph-checkpoint/src/checkpoint/memory.rs +++ b/crates/langgraph-checkpoint/src/checkpoint/memory.rs @@ -10,17 +10,20 @@ use std::collections::HashMap; type StorageKey = (String, String, String); // (thread_id, checkpoint_ns, checkpoint_id) type WriteKey = (String, String, String, i64); // (thread_id, checkpoint_ns, checkpoint_id, idx) -/// (thread_id, checkpoint_ns, checkpoint_id) -> (checkpoint_json, metadata_json, parent_checkpoint_id) -type StorageValue = (JsonValue, JsonValue, Option); +/// (thread_id, checkpoint_ns, checkpoint_id) -> (checkpoint, metadata, parent_checkpoint_id) +type StorageValue = (Checkpoint, CheckpointMetadata, Option); /// (thread_id, checkpoint_ns, checkpoint_id, idx) -> (task_id, channel, value_json, task_path) type WriteValue = (String, String, JsonValue, String); /// In-memory checkpoint saver for testing and development. /// -/// Stores checkpoints, blobs, and writes in memory using DashMap -/// for concurrent access. +/// Stores the `Checkpoint` struct directly (no JSON round trip) and tracks +/// the newest checkpoint per thread so lookups are O(1) instead of scanning +/// the whole history. pub struct InMemorySaver { storage: RwLock>, + /// (thread_id, checkpoint_ns) -> newest checkpoint_id + latest: RwLock>, writes: RwLock>, } @@ -28,6 +31,7 @@ impl InMemorySaver { pub fn new() -> Self { Self { storage: RwLock::new(HashMap::new()), + latest: RwLock::new(HashMap::new()), writes: RwLock::new(HashMap::new()), } } @@ -65,34 +69,30 @@ impl BaseCheckpointSaver for InMemorySaver { config: &RunnableConfig, ) -> Result, CheckpointError> { let (thread_id, checkpoint_ns, checkpoint_id) = Self::config_to_ids(config); - let storage = self.storage.read(); - // Find the checkpoint - let key = if let Some(ref cid) = checkpoint_id { - (thread_id.clone(), checkpoint_ns.clone(), cid.clone()) - } else { - // Find the latest checkpoint for this thread/ns - let candidates: Vec<_> = storage - .keys() - .filter(|(tid, ns, _)| tid == &thread_id && ns == &checkpoint_ns) - .collect(); - - match candidates.into_iter().max_by_key(|(_, _, cid)| cid.clone()) { - Some((tid, ns, cid)) => (tid.clone(), ns.clone(), cid.clone()), + // Resolve the requested checkpoint: explicit id, else the thread's + // newest (tracked O(1) instead of scanning the whole history). + let resolved_cid = match checkpoint_id { + Some(cid) => cid, + None => match self + .latest + .read() + .get(&(thread_id.clone(), checkpoint_ns.clone())) + { + Some(cid) => cid.clone(), None => return Ok(None), - } + }, }; - let (checkpoint_json, metadata_json, parent_cid) = match storage.get(&key) { + let (checkpoint, metadata, parent_cid) = match self.storage.read().get(&( + thread_id.clone(), + checkpoint_ns.clone(), + resolved_cid.clone(), + )) { Some(v) => v.clone(), None => return Ok(None), }; - let checkpoint: Checkpoint = serde_json::from_value(checkpoint_json) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata: CheckpointMetadata = serde_json::from_value(metadata_json) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let parent_config = parent_cid.map(|pid| { let mut c = RunnableConfig::new(); c.insert( @@ -110,7 +110,9 @@ impl BaseCheckpointSaver for InMemorySaver { let writes = self.writes.read(); let pending_writes: Vec = writes .iter() - .filter(|((tid, ns, cid, _), _)| tid == &key.0 && ns == &key.1 && cid == &key.2) + .filter(|((tid, ns, cid, _), _)| { + tid == &thread_id && ns == &checkpoint_ns && cid == &resolved_cid + }) .map(|(_, (task_id, channel, value, _))| { (task_id.clone(), channel.clone(), value.clone()) }) @@ -124,7 +126,7 @@ impl BaseCheckpointSaver for InMemorySaver { serde_json::json!({ "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, - "checkpoint_id": key.2, + "checkpoint_id": resolved_cid, }), ); c @@ -182,10 +184,11 @@ impl BaseCheckpointSaver for InMemorySaver { } let mut results = Vec::new(); - for ((tid, ns, cid), (checkpoint_json, metadata_json, parent_cid)) in entries { + for ((tid, ns, cid), (checkpoint, metadata, parent_cid)) in entries { // Apply filter if let Some(filter) = filter { - let metadata_val: JsonValue = metadata_json.clone(); + let metadata_val: JsonValue = serde_json::to_value(metadata) + .map_err(|e| CheckpointError::Storage(e.to_string()))?; let mut matches = true; for (k, v) in filter { if metadata_val.get(k) != Some(v) { @@ -198,11 +201,6 @@ impl BaseCheckpointSaver for InMemorySaver { } } - let checkpoint: Checkpoint = serde_json::from_value(checkpoint_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata: CheckpointMetadata = serde_json::from_value(metadata_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let parent_config = parent_cid.as_ref().map(|pid| { let mut c = RunnableConfig::new(); c.insert( @@ -229,8 +227,8 @@ impl BaseCheckpointSaver for InMemorySaver { ); c }, - checkpoint, - metadata, + checkpoint: checkpoint.clone(), + metadata: metadata.clone(), parent_config, pending_writes: None, }); @@ -242,35 +240,37 @@ impl BaseCheckpointSaver for InMemorySaver { fn put( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, _new_versions: &ChannelVersions, ) -> Result { let (thread_id, checkpoint_ns, _) = Self::config_to_ids(config); + let cid = checkpoint.id.clone(); - let checkpoint_json = serde_json::to_value(checkpoint) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata_json = - serde_json::to_value(metadata).map_err(|e| CheckpointError::Storage(e.to_string()))?; - - // Get the current parent - let parent_id = { - let storage = self.storage.read(); - storage - .keys() - .filter(|(tid, ns, _)| tid == &thread_id && ns == &checkpoint_ns) - .max_by_key(|(_, _, cid)| cid.clone()) - .map(|(_, _, cid)| cid.clone()) - }; + // Parent = the thread's current newest checkpoint (O(1)). + let parent_id = self + .latest + .read() + .get(&(thread_id.clone(), checkpoint_ns.clone())) + .cloned(); - let key = ( - thread_id.clone(), - checkpoint_ns.clone(), - checkpoint.id.clone(), - ); + let key = (thread_id.clone(), checkpoint_ns.clone(), cid.clone()); self.storage .write() - .insert(key, (checkpoint_json, metadata_json, parent_id)); + .insert(key, (checkpoint, metadata.clone(), parent_id)); + + // Track the newest checkpoint per thread. Keep the string-max + // semantics of the previous whole-history scan: only update when the + // new id sorts higher. + let mut latest = self.latest.write(); + latest + .entry((thread_id.clone(), checkpoint_ns.clone())) + .and_modify(|existing| { + if cid > *existing { + *existing = cid.clone(); + } + }) + .or_insert_with(|| cid.clone()); let mut new_config = RunnableConfig::new(); new_config.insert( @@ -278,7 +278,7 @@ impl BaseCheckpointSaver for InMemorySaver { serde_json::json!({ "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint.id, + "checkpoint_id": cid, }), ); Ok(new_config) @@ -320,12 +320,49 @@ impl BaseCheckpointSaver for InMemorySaver { self.storage .write() .retain(|(tid, _, _), _| tid != thread_id); + self.latest.write().retain(|(tid, _), _| tid != thread_id); self.writes .write() .retain(|(tid, _, _, _), _| tid != thread_id); Ok(()) } + // Async overrides: this saver is pure in-memory, so the async mirrors are + // just the sync bodies without the trait's default block_in_place bridge + // (which requires a multi-thread runtime and adds a thread handoff per + // call on the hot checkpoint path). + + async fn aget_tuple( + &self, + config: &RunnableConfig, + ) -> Result, CheckpointError> { + self.get_tuple(config) + } + + async fn aput( + &self, + config: &RunnableConfig, + checkpoint: Checkpoint, + metadata: &CheckpointMetadata, + new_versions: &ChannelVersions, + ) -> Result { + self.put(config, checkpoint, metadata, new_versions) + } + + async fn aput_writes( + &self, + config: &RunnableConfig, + writes: Vec<(String, String, JsonValue)>, + task_id: String, + task_path: String, + ) -> Result<(), CheckpointError> { + self.put_writes(config, &writes, &task_id, &task_path) + } + + async fn adelete_thread(&self, thread_id: String) -> Result<(), CheckpointError> { + self.delete_thread(&thread_id) + } + fn get_next_version(&self, current: Option<&ChannelVersion>) -> ChannelVersion { match current { Some(JsonValue::String(s)) => { @@ -381,7 +418,7 @@ mod tests { ); let new_config = saver - .put(&config, &checkpoint, &metadata, &HashMap::new()) + .put(&config, checkpoint.clone(), &metadata, &HashMap::new()) .unwrap(); let tuple = saver.get_tuple(&new_config).unwrap(); assert!(tuple.is_some()); @@ -411,7 +448,7 @@ mod tests { ); saver - .put(&config, &checkpoint, &metadata, &HashMap::new()) + .put(&config, checkpoint, &metadata, &HashMap::new()) .unwrap(); } @@ -446,7 +483,7 @@ mod tests { ); saver - .put(&config, &checkpoint, &metadata, &HashMap::new()) + .put(&config, checkpoint, &metadata, &HashMap::new()) .unwrap(); saver.delete_thread("test-thread").unwrap(); @@ -469,7 +506,7 @@ mod tests { ); let new_config = saver - .put(&config, &checkpoint, &metadata, &HashMap::new()) + .put(&config, checkpoint, &metadata, &HashMap::new()) .unwrap(); let writes = vec![ diff --git a/crates/langgraph-core/src/channels/any_value.rs b/crates/langgraph-core/src/channels/any_value.rs index 9057b06..16a6c7e 100644 --- a/crates/langgraph-core/src/channels/any_value.rs +++ b/crates/langgraph-core/src/channels/any_value.rs @@ -24,10 +24,10 @@ impl Channel for AnyValue { self.value.read().clone() } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), - value: RwLock::new(checkpoint.cloned()), + value: RwLock::new(checkpoint), }) } diff --git a/crates/langgraph-core/src/channels/base.rs b/crates/langgraph-core/src/channels/base.rs index 8d877df..7e36424 100644 --- a/crates/langgraph-core/src/channels/base.rs +++ b/crates/langgraph-core/src/channels/base.rs @@ -12,8 +12,10 @@ pub trait Channel: Send + Sync + 'static { fn checkpoint(&self) -> Option; /// Restore channel state from a checkpoint. + /// Takes ownership of the checkpoint value so restoration moves (rather + /// than deep-copies) the serialized state out of the loaded checkpoint. #[allow(clippy::wrong_self_convention)] - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box; + fn from_checkpoint(&self, checkpoint: Option) -> Box; /// Apply a batch of updates. Returns true if the channel was modified. fn update(&self, values: &[JsonValue]) -> Result; diff --git a/crates/langgraph-core/src/channels/binop.rs b/crates/langgraph-core/src/channels/binop.rs index 3a28376..ba78103 100644 --- a/crates/langgraph-core/src/channels/binop.rs +++ b/crates/langgraph-core/src/channels/binop.rs @@ -34,10 +34,10 @@ impl Channel for BinaryOperatorAggregate { self.value.read().clone() } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), - value: RwLock::new(checkpoint.cloned()), + value: RwLock::new(checkpoint), reducer: self.reducer, }) } @@ -168,7 +168,7 @@ mod tests { ch.update(&[serde_json::json!([1, 2])]).unwrap(); let cp = ch.checkpoint(); - let restored = ch.from_checkpoint(cp.as_ref()); + let restored = ch.from_checkpoint(cp); assert_eq!(restored.get().unwrap(), serde_json::json!([1, 2])); restored.update(&[serde_json::json!([3])]).unwrap(); diff --git a/crates/langgraph-core/src/channels/ephemeral_value.rs b/crates/langgraph-core/src/channels/ephemeral_value.rs index b275ba1..8a55e19 100644 --- a/crates/langgraph-core/src/channels/ephemeral_value.rs +++ b/crates/langgraph-core/src/channels/ephemeral_value.rs @@ -29,10 +29,10 @@ impl Channel for EphemeralValue { self.value.read().clone() } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), - value: RwLock::new(checkpoint.cloned()), + value: RwLock::new(checkpoint), guard: self.guard, }) } diff --git a/crates/langgraph-core/src/channels/last_value.rs b/crates/langgraph-core/src/channels/last_value.rs index 61bc837..4b85ea4 100644 --- a/crates/langgraph-core/src/channels/last_value.rs +++ b/crates/langgraph-core/src/channels/last_value.rs @@ -25,10 +25,10 @@ impl Channel for LastValue { self.value.read().clone() } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), - value: RwLock::new(checkpoint.cloned()), + value: RwLock::new(checkpoint), }) } @@ -105,10 +105,10 @@ impl Channel for LastValueAfterFinish { } } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), - value: RwLock::new(checkpoint.cloned()), + value: RwLock::new(checkpoint), pending: RwLock::new(None), finished: RwLock::new(false), }) @@ -196,7 +196,7 @@ mod tests { let cp = ch.checkpoint(); assert_eq!(cp, Some(serde_json::json!("hello"))); - let restored = ch.from_checkpoint(cp.as_ref()); + let restored = ch.from_checkpoint(cp); assert_eq!(restored.get().unwrap(), serde_json::json!("hello")); } diff --git a/crates/langgraph-core/src/channels/named_barrier_value.rs b/crates/langgraph-core/src/channels/named_barrier_value.rs index 5cf8f23..be69c43 100644 --- a/crates/langgraph-core/src/channels/named_barrier_value.rs +++ b/crates/langgraph-core/src/channels/named_barrier_value.rs @@ -41,10 +41,10 @@ impl Channel for NamedBarrierValue { } } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { let seen = match checkpoint { Some(JsonValue::Array(arr)) => arr - .iter() + .into_iter() .filter_map(|v| v.as_str().map(|s| s.to_string())) .collect(), _ => HashSet::new(), @@ -141,13 +141,14 @@ impl Channel for NamedBarrierValueAfterFinish { } } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { + let is_some = checkpoint.is_some(); Box::new(Self { key: self.key.clone(), names: self.names.clone(), seen: RwLock::new(HashSet::new()), - value: RwLock::new(checkpoint.cloned()), - finished: RwLock::new(checkpoint.is_some()), + value: RwLock::new(checkpoint), + finished: RwLock::new(is_some), }) } diff --git a/crates/langgraph-core/src/channels/topic.rs b/crates/langgraph-core/src/channels/topic.rs index 9106e80..83289da 100644 --- a/crates/langgraph-core/src/channels/topic.rs +++ b/crates/langgraph-core/src/channels/topic.rs @@ -33,10 +33,10 @@ impl Channel for Topic { } } - fn from_checkpoint(&self, checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, checkpoint: Option) -> Box { let values = match checkpoint { - Some(JsonValue::Array(arr)) => arr.clone(), - Some(other) => vec![other.clone()], + Some(JsonValue::Array(arr)) => arr, + Some(other) => vec![other], None => Vec::new(), }; Box::new(Self { diff --git a/crates/langgraph-core/src/channels/untracked_value.rs b/crates/langgraph-core/src/channels/untracked_value.rs index 40869c0..f7e62af 100644 --- a/crates/langgraph-core/src/channels/untracked_value.rs +++ b/crates/langgraph-core/src/channels/untracked_value.rs @@ -25,7 +25,7 @@ impl Channel for UntrackedValue { None } - fn from_checkpoint(&self, _checkpoint: Option<&JsonValue>) -> Box { + fn from_checkpoint(&self, _checkpoint: Option) -> Box { Box::new(Self { key: self.key.clone(), value: RwLock::new(None), diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index d0e715b..376de57 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -510,6 +510,11 @@ impl CompiledStateGraph { /// run; only channels whose version moved since then are passed to the /// saver as `new_versions`, so persistent savers can write delta blobs /// instead of re-encoding every channel on every super-step. + /// + /// Sync variant, used by the cold-path public API (`update_state`). The + /// hot BSP loop uses [`Self::save_checkpoint_async`] to avoid the + /// sync/async bridge (block_in_place + runtime handoff) on every + /// super-step. fn save_checkpoint( &self, checkpointer: &Arc, @@ -519,36 +524,30 @@ impl CompiledStateGraph { versions_seen: &HashMap>, previous_versions: &ChannelVersions, ) -> Option { - use chrono::Utc; - use langgraph_checkpoint::checkpoint::id::uuid6; - - // Collect all channel values (including trigger channels for state history) - let channel_values: HashMap = channels - .iter() - .filter_map(|(k, v)| v.checkpoint().map(|val| (k.clone(), val))) - .collect(); - - let checkpoint = langgraph_checkpoint::Checkpoint { - v: 2, - id: uuid6(), - ts: Utc::now().to_rfc3339(), - channel_values, - channel_versions: channel_versions.clone(), - versions_seen: versions_seen.clone(), - updated_channels: None, - }; - - // Delta vs. the versions this run started from: only channels whose - // version changed since then need new blob rows. - let new_versions: ChannelVersions = channel_versions - .iter() - .filter(|(k, v)| previous_versions.get(k.as_str()) != Some(*v)) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let (checkpoint, metadata, new_versions) = + build_checkpoint(channels, channel_versions, versions_seen, previous_versions); + checkpointer + .put(config, checkpoint, &metadata, &new_versions) + .ok() + } - let metadata = CheckpointMetadata::default(); + /// Async variant of [`Self::save_checkpoint`] for the hot BSP loop: awaits + /// `aput` directly so the graph runner never bridges through + /// `block_in_place(block_on(...))` on the per-super-step save path. + async fn save_checkpoint_async( + &self, + checkpointer: &Arc, + config: &RunnableConfig, + channels: &HashMap>, + channel_versions: &ChannelVersions, + versions_seen: &HashMap>, + previous_versions: &ChannelVersions, + ) -> Option { + let (checkpoint, metadata, new_versions) = + build_checkpoint(channels, channel_versions, versions_seen, previous_versions); checkpointer - .put(config, &checkpoint, &metadata, &new_versions) + .aput(config, checkpoint, &metadata, &new_versions) + .await .ok() } @@ -618,13 +617,8 @@ impl CompiledStateGraph { }; // Reconstruct channels from checkpoint - let cp_channels: HashMap> = saved - .checkpoint - .channel_values - .iter() - .map(|(k, v)| (k.clone(), Some(v.clone()))) - .collect(); - let mut channels = channels_from_checkpoint(&self.channels, &cp_channels); + let mut channels = + channels_from_checkpoint(&self.channels, saved.checkpoint.channel_values); let mut channel_versions = saved.checkpoint.channel_versions.clone(); let mut versions_seen = saved.checkpoint.versions_seen.clone(); @@ -800,29 +794,22 @@ impl CompiledStateGraph { .map_err(|e| GraphError::Checkpoint(e.to_string()))?; // Reconstruct channels from checkpoint (or fresh if none) - let channels: HashMap> = if let Some(ref saved) = saved { - let cp_channels: HashMap> = saved - .checkpoint - .channel_values - .iter() - .map(|(k, v)| (k.clone(), Some(v.clone()))) - .collect(); - channels_from_checkpoint(&self.channels, &cp_channels) + let (channels, mut channel_versions, versions_seen) = if let Some(saved) = saved { + let checkpoint = saved.checkpoint; + let channels = channels_from_checkpoint(&self.channels, checkpoint.channel_values); + ( + channels, + checkpoint.channel_versions, + checkpoint.versions_seen, + ) } else { - self.channels + let channels = self + .channels .iter() .map(|(k, c)| (k.clone(), c.clone_channel())) - .collect() + .collect(); + (channels, HashMap::new(), HashMap::new()) }; - - let mut channel_versions = saved - .as_ref() - .map(|s| s.checkpoint.channel_versions.clone()) - .unwrap_or_default(); - let versions_seen = saved - .as_ref() - .map(|s| s.checkpoint.versions_seen.clone()) - .unwrap_or_default(); let previous_versions = channel_versions.clone(); // Apply the update values to channels @@ -893,15 +880,10 @@ impl CompiledStateGraph { let pregel_nodes = &self.pregel_nodes; let trigger_to_nodes = &self.trigger_to_nodes; - for saved in &tuples { + for saved in tuples { // Reconstruct channels from checkpoint - let cp_channels: HashMap> = saved - .checkpoint - .channel_values - .iter() - .map(|(k, v)| (k.clone(), Some(v.clone()))) - .collect(); - let channels = channels_from_checkpoint(&self.channels, &cp_channels); + let channels = + channels_from_checkpoint(&self.channels, saved.checkpoint.channel_values); let channel_versions = saved.checkpoint.channel_versions.clone(); let mut versions_seen = saved.checkpoint.versions_seen.clone(); @@ -1265,6 +1247,50 @@ fn apply_completed_writes( } } +/// Build a `Checkpoint` plus its metadata and the `new_versions` delta from the +/// current channel state. Pure CPU — both the sync and async save paths share +/// it; only the final saver call differs (`put` vs `aput`). +fn build_checkpoint( + channels: &HashMap>, + channel_versions: &ChannelVersions, + versions_seen: &HashMap>, + previous_versions: &ChannelVersions, +) -> ( + langgraph_checkpoint::Checkpoint, + CheckpointMetadata, + ChannelVersions, +) { + use chrono::Utc; + use langgraph_checkpoint::checkpoint::id::uuid6; + + // Collect all channel values (including trigger channels for state history) + let channel_values: HashMap = channels + .iter() + .filter_map(|(k, v)| v.checkpoint().map(|val| (k.clone(), val))) + .collect(); + + let checkpoint = langgraph_checkpoint::Checkpoint { + v: 2, + id: uuid6(), + ts: Utc::now().to_rfc3339(), + channel_values, + channel_versions: channel_versions.clone(), + versions_seen: versions_seen.clone(), + updated_channels: None, + }; + + // Delta vs. the versions this run started from: only channels whose + // version changed since then need new blob rows. + let new_versions: ChannelVersions = channel_versions + .iter() + .filter(|(k, v)| previous_versions.get(k.as_str()) != Some(*v)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + let metadata = CheckpointMetadata::default(); + (checkpoint, metadata, new_versions) +} + // Helper: collect output channel keys (excluding internal routing channels). fn output_channel_keys(channels: &HashMap>) -> Vec { channels @@ -1374,55 +1400,51 @@ impl CompiledStateGraph { // Load checkpoint (for resume support) let mut saved_checkpoint_exists = false; - let (mut channels, mut channel_versions, mut versions_seen) = - if let Some(ref cp) = self.checkpointer { - match cp.get_tuple(&config) { - Ok(Some(tuple)) => { - saved_checkpoint_exists = true; - let cp_channels: HashMap> = tuple - .checkpoint - .channel_values - .iter() - .map(|(k, v)| (k.clone(), Some(v.clone()))) - .collect(); - let restored = channels_from_checkpoint(&self.channels, &cp_channels); - - // Apply non-RESUME pending writes from the checkpoint - if let Some(ref pending) = tuple.pending_writes { - for (_task_id, channel, value) in pending { - if channel != RESUME { - if let Some(ch) = restored.get(channel) { - ch.update(std::slice::from_ref(value)).ok(); - } + let (mut channels, mut channel_versions, mut versions_seen) = if let Some(ref cp) = + self.checkpointer + { + match cp.aget_tuple(&config).await { + Ok(Some(tuple)) => { + saved_checkpoint_exists = true; + let restored = + channels_from_checkpoint(&self.channels, tuple.checkpoint.channel_values); + + // Apply non-RESUME pending writes from the checkpoint + if let Some(ref pending) = tuple.pending_writes { + for (_task_id, channel, value) in pending { + if channel != RESUME { + if let Some(ch) = restored.get(channel) { + ch.update(std::slice::from_ref(value)).ok(); } } } - - ( - restored, - tuple.checkpoint.channel_versions.clone(), - tuple.checkpoint.versions_seen.clone(), - ) } - _ => ( - self.channels - .iter() - .map(|(k, c)| (k.clone(), c.clone_channel())) - .collect(), - HashMap::new(), - HashMap::new(), - ), + + ( + restored, + tuple.checkpoint.channel_versions.clone(), + tuple.checkpoint.versions_seen.clone(), + ) } - } else { - ( + _ => ( self.channels .iter() .map(|(k, c)| (k.clone(), c.clone_channel())) .collect(), HashMap::new(), HashMap::new(), - ) - }; + ), + } + } else { + ( + self.channels + .iter() + .map(|(k, c)| (k.clone(), c.clone_channel())) + .collect(), + HashMap::new(), + HashMap::new(), + ) + }; // Baseline for incremental checkpoint writes: the versions as of the // start of this run. Only channels that move past these get blob rows. @@ -1567,14 +1589,17 @@ impl CompiledStateGraph { let task_names: Vec = tasks.iter().map(|t| t.name.clone()).collect(); if task_names.iter().any(|n| self.interrupt_before.contains(n)) { if let Some(ref cp) = self.checkpointer { - if let Some(new_config) = self.save_checkpoint( - cp, - &config, - &channels, - &channel_versions, - &versions_seen, - &previous_versions, - ) { + if let Some(new_config) = self + .save_checkpoint_async( + cp, + &config, + &channels, + &channel_versions, + &versions_seen, + &previous_versions, + ) + .await + { config = new_config; } } @@ -1641,14 +1666,17 @@ impl CompiledStateGraph { // Save checkpoint (now includes completed tasks' channel writes) if let Some(ref cp) = self.checkpointer { - if let Some(new_config) = self.save_checkpoint( - cp, - &config, - &channels, - &channel_versions, - &versions_seen, - &previous_versions, - ) { + if let Some(new_config) = self + .save_checkpoint_async( + cp, + &config, + &channels, + &channel_versions, + &versions_seen, + &previous_versions, + ) + .await + { config = new_config; } // Save interrupt as pending writes for get_state() @@ -1665,7 +1693,9 @@ impl CompiledStateGraph { }) .collect(); if !iw.is_empty() { - if let Err(e) = cp.put_writes(&config, &iw, &task_id, "") { + if let Err(e) = + cp.aput_writes(&config, iw, task_id, String::new()).await + { eprintln!("[CHECKPOINT] Failed to save interrupt writes: {}", e); } } @@ -1740,14 +1770,17 @@ impl CompiledStateGraph { // Save "loop" checkpoint after each completed super-step if let Some(ref cp) = self.checkpointer { - if let Some(new_config) = self.save_checkpoint( - cp, - &config, - &channels, - &channel_versions, - &versions_seen, - &previous_versions, - ) { + if let Some(new_config) = self + .save_checkpoint_async( + cp, + &config, + &channels, + &channel_versions, + &versions_seen, + &previous_versions, + ) + .await + { config = new_config; } } diff --git a/crates/langgraph-core/src/pregel/mod.rs b/crates/langgraph-core/src/pregel/mod.rs index 1ea9b69..fee8f53 100644 --- a/crates/langgraph-core/src/pregel/mod.rs +++ b/crates/langgraph-core/src/pregel/mod.rs @@ -85,13 +85,17 @@ pub fn build_trigger_to_nodes(nodes: &HashMap) -> TriggerToN } /// Reconstruct live channels from a checkpoint. +/// +/// Consumes the checkpoint's channel values so restored state is moved into +/// the channels instead of deep-copied (the checkpoint is owned by the +/// caller and not needed afterwards). pub fn channels_from_checkpoint( specs: &HashMap>, - checkpoint_channels: &HashMap>, + mut checkpoint_channels: HashMap, ) -> HashMap> { let mut channels = HashMap::new(); for (key, spec) in specs { - let cp = checkpoint_channels.get(key).and_then(|v| v.as_ref()); + let cp = checkpoint_channels.remove(key); channels.insert(key.clone(), spec.from_checkpoint(cp)); } channels diff --git a/examples/sqlite_checkpoint.rs b/examples/sqlite_checkpoint.rs index 9d5c33b..70b6631 100644 --- a/examples/sqlite_checkpoint.rs +++ b/examples/sqlite_checkpoint.rs @@ -53,7 +53,7 @@ async fn main() -> Result<(), Box> { step: Some(0), ..Default::default() }; - let next_cfg = saver.aput(&cfg, &cp1, &metadata, &vers1).await?; + let next_cfg = saver.aput(&cfg, cp1.clone(), &metadata, &vers1).await?; println!("stored checkpoint #1: id={}", cp1.id); // Second checkpoint that references the first as parent — both @@ -67,7 +67,9 @@ async fn main() -> Result<(), Box> { step: Some(1), ..Default::default() }; - saver.aput(&next_cfg, &cp2, &metadata2, &vers2).await?; + saver + .aput(&next_cfg, cp2.clone(), &metadata2, &vers2) + .await?; println!("stored checkpoint #2: id={}", cp2.id); // Fetch latest diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index 37609cd..d08c756 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -58,8 +58,8 @@ fn make_message(i: usize) -> JsonValue { }) } -/// (thread_id, checkpoint_ns) -> (checkpoint_id, checkpoint_json, metadata_json, parent_cid) -type StorageEntry = (String, JsonValue, JsonValue, Option); +/// (thread_id, checkpoint_ns) -> (checkpoint_id, checkpoint, metadata, parent_cid) +type StorageEntry = (String, Checkpoint, CheckpointMetadata, Option); /// (thread_id, checkpoint_ns) -> the thread's newest checkpoint type StorageMap = HashMap<(String, String), StorageEntry>; /// (thread_id, checkpoint_ns, checkpoint_id) -> pending writes (interrupt-only path) @@ -71,9 +71,9 @@ type WritesMap = HashMap<(String, String, String), Vec<(String, String, JsonValu /// super-steps with a growing message history retains O(N²) serialized state /// and OOMs after a few hundred steps. Production savers prune; this one does /// the same — each `put` replaces the thread's previous checkpoint, so retained -/// memory is O(latest state). The serde round-trip (to_value on `put`, -/// from_value on `get_tuple`) mirrors `InMemorySaver`, so the measured per-step -/// cost stays comparable. +/// memory is O(latest state). It stores the `Checkpoint` struct directly (no +/// JSON round trip), mirroring the current `InMemorySaver`, so the measured +/// per-step cost stays comparable. struct LatestOnlySaver { storage: RwLock, writes: RwLock, @@ -118,6 +118,7 @@ impl Default for LatestOnlySaver { } } +#[async_trait::async_trait] impl BaseCheckpointSaver for LatestOnlySaver { fn get_tuple( &self, @@ -125,7 +126,7 @@ impl BaseCheckpointSaver for LatestOnlySaver { ) -> Result, CheckpointError> { let (thread_id, checkpoint_ns, requested_id) = Self::config_to_ids(config); let storage = self.storage.read().unwrap(); - let Some((cid, checkpoint_json, metadata_json, parent_cid)) = + let Some((cid, checkpoint, metadata, parent_cid)) = storage.get(&(thread_id.clone(), checkpoint_ns.clone())) else { return Ok(None); @@ -137,11 +138,6 @@ impl BaseCheckpointSaver for LatestOnlySaver { } } - let checkpoint: Checkpoint = serde_json::from_value(checkpoint_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata: CheckpointMetadata = serde_json::from_value(metadata_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let parent_config = parent_cid.as_ref().map(|pid| { let mut c = RunnableConfig::new(); c.insert( @@ -176,8 +172,8 @@ impl BaseCheckpointSaver for LatestOnlySaver { ); c }, - checkpoint, - metadata, + checkpoint: checkpoint.clone(), + metadata: metadata.clone(), parent_config, pending_writes: if pending_writes.is_empty() { None @@ -216,11 +212,7 @@ impl BaseCheckpointSaver for LatestOnlySaver { } let mut results = Vec::new(); - for ((tid, ns), (cid, checkpoint_json, metadata_json, parent_cid)) in entries { - let checkpoint: Checkpoint = serde_json::from_value(checkpoint_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata: CheckpointMetadata = serde_json::from_value(metadata_json.clone()) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; + for ((tid, ns), (cid, checkpoint, metadata, parent_cid)) in entries { let parent_config = parent_cid.as_ref().map(|pid| { let mut c = RunnableConfig::new(); c.insert( @@ -246,8 +238,8 @@ impl BaseCheckpointSaver for LatestOnlySaver { ); c }, - checkpoint, - metadata, + checkpoint: checkpoint.clone(), + metadata: metadata.clone(), parent_config, pending_writes: None, }); @@ -258,17 +250,12 @@ impl BaseCheckpointSaver for LatestOnlySaver { fn put( &self, config: &RunnableConfig, - checkpoint: &Checkpoint, + checkpoint: Checkpoint, metadata: &CheckpointMetadata, _new_versions: &ChannelVersions, ) -> Result { let (thread_id, checkpoint_ns, _) = Self::config_to_ids(config); - let checkpoint_json = serde_json::to_value(checkpoint) - .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let metadata_json = - serde_json::to_value(metadata).map_err(|e| CheckpointError::Storage(e.to_string()))?; - // The new checkpoint's parent is the thread's current newest. let parent_id = self .storage @@ -278,19 +265,16 @@ impl BaseCheckpointSaver for LatestOnlySaver { .map(|(cid, _, _, _)| cid.clone()); // Replace the thread's checkpoint — only the newest is retained. + let cid = checkpoint.id.clone(); self.storage.write().unwrap().insert( (thread_id.clone(), checkpoint_ns.clone()), - ( - checkpoint.id.clone(), - checkpoint_json, - metadata_json, - parent_id, - ), + (cid.clone(), checkpoint, metadata.clone(), parent_id), ); // Prune pending writes down to the newest checkpoint id. - self.writes.write().unwrap().retain(|(tid, ns, cid), _| { - tid != &thread_id || ns != &checkpoint_ns || cid == &checkpoint.id - }); + self.writes + .write() + .unwrap() + .retain(|(tid, ns, cid2), _| tid != &thread_id || ns != &checkpoint_ns || cid2 == &cid); let mut new_config = RunnableConfig::new(); new_config.insert( @@ -298,7 +282,7 @@ impl BaseCheckpointSaver for LatestOnlySaver { serde_json::json!({ "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint.id, + "checkpoint_id": cid, }), ); Ok(new_config) @@ -334,6 +318,40 @@ impl BaseCheckpointSaver for LatestOnlySaver { .retain(|(tid, _, _), _| tid != thread_id); Ok(()) } + + // Async overrides mirroring InMemorySaver: pure in-memory, so skip the + // trait default's block_in_place bridge — the benches must measure the + // native in-memory path, not a thread handoff per call. + async fn aget_tuple( + &self, + config: &RunnableConfig, + ) -> Result, CheckpointError> { + self.get_tuple(config) + } + + async fn aput( + &self, + config: &RunnableConfig, + checkpoint: Checkpoint, + metadata: &CheckpointMetadata, + new_versions: &ChannelVersions, + ) -> Result { + self.put(config, checkpoint, metadata, new_versions) + } + + async fn aput_writes( + &self, + config: &RunnableConfig, + writes: Vec<(String, String, JsonValue)>, + task_id: String, + task_path: String, + ) -> Result<(), CheckpointError> { + self.put_writes(config, &writes, &task_id, &task_path) + } + + async fn adelete_thread(&self, thread_id: String) -> Result<(), CheckpointError> { + self.delete_thread(&thread_id) + } } /// Single-node graph: `messages` accumulates with the `add_messages` reducer. From 26c7909bc359522d67f3d685deb468cfdc40dcb5 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 13:00:49 +0800 Subject: [PATCH 08/15] perf: emit delta-only channel_values from build_checkpoint build_checkpoint now serializes only the channels whose version moved since the run started (new_versions delta) instead of deep-copying every channel's checkpoint() each super-step. Unchanged channels are reconstructed by the saver's version-merged reads, so nothing is lost. - InMemorySaver: store channel values as version-addressed blobs (rows + blobs + O(1) latest), reconstruct the full channel_values on get_tuple/list via reconstruct_values. Mirrors the sqlite blob model. - LatestOnlySaver (bench): merge the delta into the retained checkpoint on put instead of replacing it wholesale. - run_pregel_inner input handling: bump versions for every channel written from the input (both START-mapped and direct dict-key writes). Previously dict-key writes set the value without a version entry, so a value-bearing channel could be silently dropped from a delta-only checkpoint (regression surfaced by test_update_state). --- .../src/checkpoint/memory.rs | 230 ++++++++++++++++-- crates/langgraph-core/src/graph/state.rs | 42 +++- tests/bench_pregel.rs | 44 +++- 3 files changed, 280 insertions(+), 36 deletions(-) diff --git a/crates/langgraph-checkpoint/src/checkpoint/memory.rs b/crates/langgraph-checkpoint/src/checkpoint/memory.rs index 752589f..68e26b4 100644 --- a/crates/langgraph-checkpoint/src/checkpoint/memory.rs +++ b/crates/langgraph-checkpoint/src/checkpoint/memory.rs @@ -7,21 +7,31 @@ use parking_lot::RwLock; use serde_json::Value as JsonValue; use std::collections::HashMap; -type StorageKey = (String, String, String); // (thread_id, checkpoint_ns, checkpoint_id) +type RowKey = (String, String, String); // (thread_id, checkpoint_ns, checkpoint_id) type WriteKey = (String, String, String, i64); // (thread_id, checkpoint_ns, checkpoint_id, idx) +/// Blob key: (thread_id, checkpoint_ns, channel, version_str) +type BlobKey = (String, String, String, String); /// (thread_id, checkpoint_ns, checkpoint_id) -> (checkpoint, metadata, parent_checkpoint_id) -type StorageValue = (Checkpoint, CheckpointMetadata, Option); +type RowValue = (Checkpoint, CheckpointMetadata, Option); /// (thread_id, checkpoint_ns, checkpoint_id, idx) -> (task_id, channel, value_json, task_path) type WriteValue = (String, String, JsonValue, String); /// In-memory checkpoint saver for testing and development. /// -/// Stores the `Checkpoint` struct directly (no JSON round trip) and tracks -/// the newest checkpoint per thread so lookups are O(1) instead of scanning -/// the whole history. +/// Stores `Checkpoint`s directly (no JSON round trip). Channel values are +/// stored as version-addressed blobs — `put` receives a delta +/// (`new_versions` only, see `BaseCheckpointSaver`), and reads reconstruct +/// the full state by resolving each channel's version against the blob +/// store, mirroring the version-joined merge of the sqlite saver. `None` +/// marks a channel whose version moved but has no value (e.g. a cleared +/// ephemeral channel), matching the "empty" blob rows of the sqlite saver. +/// The newest checkpoint per thread is tracked O(1). pub struct InMemorySaver { - storage: RwLock>, + /// Checkpoint rows with `channel_values` stripped (they live in `blobs`). + rows: RwLock>, + /// (thread_id, checkpoint_ns, channel, version) -> value or empty marker. + blobs: RwLock>>, /// (thread_id, checkpoint_ns) -> newest checkpoint_id latest: RwLock>, writes: RwLock>, @@ -30,12 +40,43 @@ pub struct InMemorySaver { impl InMemorySaver { pub fn new() -> Self { Self { - storage: RwLock::new(HashMap::new()), + rows: RwLock::new(HashMap::new()), + blobs: RwLock::new(HashMap::new()), latest: RwLock::new(HashMap::new()), writes: RwLock::new(HashMap::new()), } } + /// Reconstruct a checkpoint's full `channel_values` from the + /// version-addressed blob store: every channel in `channel_versions` + /// resolves to the blob written at that version (possibly by an earlier + /// checkpoint), empty markers and missing blobs are skipped. + fn reconstruct_values( + blobs: &HashMap>, + channel_versions: &ChannelVersions, + thread_id: &str, + checkpoint_ns: &str, + ) -> HashMap { + let mut values = HashMap::new(); + for (channel, ver) in channel_versions { + let ver_str = match ver { + JsonValue::String(s) => s.clone(), + JsonValue::Number(n) => n.to_string(), + _ => continue, + }; + let key = ( + thread_id.to_string(), + checkpoint_ns.to_string(), + channel.clone(), + ver_str, + ); + if let Some(Some(val)) = blobs.get(&key) { + values.insert(channel.clone(), val.clone()); + } + } + values + } + fn config_to_ids(config: &RunnableConfig) -> (String, String, Option) { let configurable = config.get("configurable"); let thread_id = configurable @@ -84,7 +125,7 @@ impl BaseCheckpointSaver for InMemorySaver { }, }; - let (checkpoint, metadata, parent_cid) = match self.storage.read().get(&( + let (mut checkpoint, metadata, parent_cid) = match self.rows.read().get(&( thread_id.clone(), checkpoint_ns.clone(), resolved_cid.clone(), @@ -93,6 +134,19 @@ impl BaseCheckpointSaver for InMemorySaver { None => return Ok(None), }; + // Reconstruct the full state: the row holds only the delta, so merge + // version-addressed blob values over it (see struct docs). + let blob_values = { + let blobs = self.blobs.read(); + Self::reconstruct_values( + &blobs, + &checkpoint.channel_versions, + &thread_id, + &checkpoint_ns, + ) + }; + checkpoint.channel_values = blob_values; + let parent_config = parent_cid.map(|pid| { let mut c = RunnableConfig::new(); c.insert( @@ -149,7 +203,7 @@ impl BaseCheckpointSaver for InMemorySaver { before: Option<&RunnableConfig>, limit: Option, ) -> Result, CheckpointError> { - let storage = self.storage.read(); + let rows = self.rows.read(); let (thread_id, checkpoint_ns) = match config { Some(c) => { @@ -161,7 +215,7 @@ impl BaseCheckpointSaver for InMemorySaver { let before_id = before.and_then(|c| Self::config_to_ids(c).2); - let mut entries: Vec<_> = storage + let mut entries: Vec<_> = rows .iter() .filter(|((tid, ns, _), _)| { (thread_id.is_empty() || tid == &thread_id) @@ -214,6 +268,14 @@ impl BaseCheckpointSaver for InMemorySaver { c }); + // Reconstruct the full state from version-addressed blobs. + let mut checkpoint = checkpoint.clone(); + let blob_values = { + let blobs = self.blobs.read(); + Self::reconstruct_values(&blobs, &checkpoint.channel_versions, tid, ns) + }; + checkpoint.channel_values = blob_values; + results.push(CheckpointTuple { config: { let mut c = RunnableConfig::new(); @@ -227,7 +289,7 @@ impl BaseCheckpointSaver for InMemorySaver { ); c }, - checkpoint: checkpoint.clone(), + checkpoint, metadata: metadata.clone(), parent_config, pending_writes: None, @@ -240,9 +302,9 @@ impl BaseCheckpointSaver for InMemorySaver { fn put( &self, config: &RunnableConfig, - checkpoint: Checkpoint, + mut checkpoint: Checkpoint, metadata: &CheckpointMetadata, - _new_versions: &ChannelVersions, + new_versions: &ChannelVersions, ) -> Result { let (thread_id, checkpoint_ns, _) = Self::config_to_ids(config); let cid = checkpoint.id.clone(); @@ -254,8 +316,32 @@ impl BaseCheckpointSaver for InMemorySaver { .get(&(thread_id.clone(), checkpoint_ns.clone())) .cloned(); + // Store the delta as version-addressed blobs: one per moved channel, + // with the value if present or an empty marker (`None`) if the + // channel was cleared — mirroring the "empty" blob rows of the + // sqlite saver. The row keeps only the metadata fields; reads + // reconstruct the full state via `reconstruct_values`. + let mut channel_values = std::mem::take(&mut checkpoint.channel_values); + { + let mut blobs = self.blobs.write(); + for (channel, ver) in new_versions { + let ver_str = match ver { + JsonValue::String(s) => s.clone(), + JsonValue::Number(n) => n.to_string(), + _ => continue, + }; + let key = ( + thread_id.clone(), + checkpoint_ns.clone(), + channel.clone(), + ver_str, + ); + blobs.insert(key, channel_values.remove(channel)); + } + } + let key = (thread_id.clone(), checkpoint_ns.clone(), cid.clone()); - self.storage + self.rows .write() .insert(key, (checkpoint, metadata.clone(), parent_id)); @@ -317,9 +403,10 @@ impl BaseCheckpointSaver for InMemorySaver { } fn delete_thread(&self, thread_id: &str) -> Result<(), CheckpointError> { - self.storage + self.rows.write().retain(|(tid, _, _), _| tid != thread_id); + self.blobs .write() - .retain(|(tid, _, _), _| tid != thread_id); + .retain(|(tid, _, _, _), _| tid != thread_id); self.latest.write().retain(|(tid, _), _| tid != thread_id); self.writes .write() @@ -529,4 +616,115 @@ mod tests { assert!(tuple.pending_writes.is_some()); assert_eq!(tuple.pending_writes.as_ref().unwrap().len(), 2); } + + #[test] + fn test_incremental_blob_merge() { + let saver = InMemorySaver::new(); + let metadata = CheckpointMetadata::default(); + + let mut config = RunnableConfig::new(); + config.insert( + "configurable".to_string(), + serde_json::json!({ + "thread_id": "test-thread", + "checkpoint_ns": "", + }), + ); + + // cp1: a=1, b=1, both at version 1. + let mut cp1 = Checkpoint::empty(); + cp1.id = "cp-001".to_string(); + cp1.channel_versions + .insert("a".into(), serde_json::json!(1)); + cp1.channel_versions + .insert("b".into(), serde_json::json!(1)); + cp1.channel_values.insert("a".into(), serde_json::json!(1)); + cp1.channel_values.insert("b".into(), serde_json::json!(1)); + let mut new_versions = HashMap::new(); + new_versions.insert("a".into(), serde_json::json!(1)); + new_versions.insert("b".into(), serde_json::json!(1)); + saver.put(&config, cp1, &metadata, &new_versions).unwrap(); + + // cp2: only a moves to version 2; b is unchanged. + let mut cp2 = Checkpoint::empty(); + cp2.id = "cp-002".to_string(); + cp2.channel_versions + .insert("a".into(), serde_json::json!(2)); + cp2.channel_versions + .insert("b".into(), serde_json::json!(1)); + cp2.channel_values.insert("a".into(), serde_json::json!(2)); + let mut new_versions = HashMap::new(); + new_versions.insert("a".into(), serde_json::json!(2)); + saver.put(&config, cp2, &metadata, &new_versions).unwrap(); + + // Reading cp2 merges b from the earlier checkpoint. + let mut cfg2 = config.clone(); + cfg2.insert( + "configurable".to_string(), + serde_json::json!({ + "thread_id": "test-thread", + "checkpoint_ns": "", + "checkpoint_id": "cp-002", + }), + ); + let tuple = saver.get_tuple(&cfg2).unwrap().unwrap(); + assert_eq!( + tuple.checkpoint.channel_values.get("a"), + Some(&serde_json::json!(2)) + ); + assert_eq!( + tuple.checkpoint.channel_values.get("b"), + Some(&serde_json::json!(1)) + ); + + // Reading cp1 is unaffected by the later delta. + let mut cfg1 = config.clone(); + cfg1.insert( + "configurable".to_string(), + serde_json::json!({ + "thread_id": "test-thread", + "checkpoint_ns": "", + "checkpoint_id": "cp-001", + }), + ); + let tuple = saver.get_tuple(&cfg1).unwrap().unwrap(); + assert_eq!( + tuple.checkpoint.channel_values.get("a"), + Some(&serde_json::json!(1)) + ); + assert_eq!( + tuple.checkpoint.channel_values.get("b"), + Some(&serde_json::json!(1)) + ); + + // Latest (no explicit checkpoint_id) resolves to cp2 with full state. + let tuple = saver.get_tuple(&config).unwrap().unwrap(); + assert_eq!(tuple.checkpoint.id, "cp-002"); + assert_eq!( + tuple.checkpoint.channel_values.get("a"), + Some(&serde_json::json!(2)) + ); + assert_eq!( + tuple.checkpoint.channel_values.get("b"), + Some(&serde_json::json!(1)) + ); + + // cp3: a moves to version 3 with no value (cleared channel) -> the + // empty marker removes it; b stays. + let mut cp3 = Checkpoint::empty(); + cp3.id = "cp-003".to_string(); + cp3.channel_versions + .insert("a".into(), serde_json::json!(3)); + cp3.channel_versions + .insert("b".into(), serde_json::json!(1)); + let mut new_versions = HashMap::new(); + new_versions.insert("a".into(), serde_json::json!(3)); + saver.put(&config, cp3, &metadata, &new_versions).unwrap(); + let tuple = saver.get_tuple(&config).unwrap().unwrap(); + assert_eq!(tuple.checkpoint.channel_values.get("a"), None); + assert_eq!( + tuple.checkpoint.channel_values.get("b"), + Some(&serde_json::json!(1)) + ); + } } diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index 376de57..f08e5d8 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -1263,10 +1263,26 @@ fn build_checkpoint( use chrono::Utc; use langgraph_checkpoint::checkpoint::id::uuid6; - // Collect all channel values (including trigger channels for state history) - let channel_values: HashMap = channels + // Delta vs. the versions this run started from: only channels whose + // version changed since then need new blob rows. + let new_versions: ChannelVersions = channel_versions .iter() - .filter_map(|(k, v)| v.checkpoint().map(|val| (k.clone(), val))) + .filter(|(k, v)| previous_versions.get(k.as_str()) != Some(*v)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + // Delta-only: serialize only channels whose version moved this run. + // Unchanged channels are reconstructed by the saver's version-merged + // reads (see BaseCheckpointSaver::get_tuple), so `channel_values` here + // is a delta, not a snapshot. Version bumps are gated on actual channel + // state changes in apply_writes, so nothing is lost. + let channel_values: HashMap = new_versions + .keys() + .filter_map(|k| { + channels + .get(k) + .and_then(|v| v.checkpoint().map(|val| (k.clone(), val))) + }) .collect(); let checkpoint = langgraph_checkpoint::Checkpoint { @@ -1279,14 +1295,6 @@ fn build_checkpoint( updated_channels: None, }; - // Delta vs. the versions this run started from: only channels whose - // version changed since then need new blob rows. - let new_versions: ChannelVersions = channel_versions - .iter() - .filter(|(k, v)| previous_versions.get(k.as_str()) != Some(*v)) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - let metadata = CheckpointMetadata::default(); (checkpoint, metadata, new_versions) } @@ -1499,9 +1507,16 @@ impl CompiledStateGraph { // "memory confusion" / spurious LLM re-runs observed after tool denial. if !is_fork && !is_resuming { let input_writes = map_input(&[START.to_string()], input); + // Collect every channel written from the input — both the + // START-mapped writes and the direct dict-key writes — and bump + // each one's version. Delta-only checkpoints serialize only + // channels whose version moved, so a value-bearing channel with no + // version entry would be silently dropped on save. + let mut written_input: Vec = Vec::new(); for (chan, val) in &input_writes { if let Some(ch) = channels.get(chan) { ch.update(std::slice::from_ref(val)).ok(); + written_input.push(chan.clone()); } } if let Some(obj) = input.as_object() { @@ -1509,13 +1524,14 @@ impl CompiledStateGraph { if key != START && !key.starts_with("branch:") && !key.starts_with("join:") { if let Some(ch) = channels.get(key) { ch.update(std::slice::from_ref(val)).ok(); + written_input.push(key.clone()); } } } } - for (chan, _) in &input_writes { + for chan in written_input { channel_versions.insert( - chan.clone(), + chan, JsonValue::String(format!("{:032}", version_offset + step)), ); } diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index d08c756..4390f9e 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -250,9 +250,9 @@ impl BaseCheckpointSaver for LatestOnlySaver { fn put( &self, config: &RunnableConfig, - checkpoint: Checkpoint, + mut checkpoint: Checkpoint, metadata: &CheckpointMetadata, - _new_versions: &ChannelVersions, + new_versions: &ChannelVersions, ) -> Result { let (thread_id, checkpoint_ns, _) = Self::config_to_ids(config); @@ -264,12 +264,42 @@ impl BaseCheckpointSaver for LatestOnlySaver { .get(&(thread_id.clone(), checkpoint_ns.clone())) .map(|(cid, _, _, _)| cid.clone()); - // Replace the thread's checkpoint — only the newest is retained. + // Merge-on-put: the engine writes deltas (`new_versions` only), so + // move the moved channels' values over the retained state, drop + // channels cleared this step (moved with no value), and refresh the + // metadata fields. Only the newest checkpoint is retained. let cid = checkpoint.id.clone(); - self.storage.write().unwrap().insert( - (thread_id.clone(), checkpoint_ns.clone()), - (cid.clone(), checkpoint, metadata.clone(), parent_id), - ); + let mut delta_values = std::mem::take(&mut checkpoint.channel_values); + let mut storage = self.storage.write().unwrap(); + match storage.get_mut(&(thread_id.clone(), checkpoint_ns.clone())) { + Some((_, prev, prev_metadata, prev_parent)) => { + for channel in new_versions.keys() { + match delta_values.remove(channel) { + Some(val) => { + prev.channel_values.insert(channel.clone(), val); + } + None => { + prev.channel_values.remove(channel); + } + } + } + prev.v = checkpoint.v; + prev.id = checkpoint.id; + prev.ts = checkpoint.ts; + prev.channel_versions = checkpoint.channel_versions; + prev.versions_seen = checkpoint.versions_seen; + prev.updated_channels = checkpoint.updated_channels; + *prev_metadata = metadata.clone(); + *prev_parent = parent_id; + } + None => { + storage.insert( + (thread_id.clone(), checkpoint_ns.clone()), + (cid.clone(), checkpoint, metadata.clone(), parent_id), + ); + } + } + drop(storage); // Prune pending writes down to the newest checkpoint id. self.writes .write() From 41d6d2979b8679a9ef72fb3a419ea48c43dea108 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 13:01:13 +0800 Subject: [PATCH 09/15] checkpoint: add BlobCache, a version-keyed cache of deserialized blobs Cache maps (thread_id, checkpoint_ns, channel, version_str) -> JsonValue for savers that store channel values as version-addressed blobs. Repeated reads of an unchanged large channel (e.g. a static context) become O(clone) instead of O(DB fetch + parse). Correctness relies on version strings being monotonic within a (thread_id, checkpoint_ns) pair; the sole reuse point is delete_thread + thread recreation, so callers MUST call remove_thread there. Whole-cache clear on overflow bounds memory. Not yet wired into any saver (wired into the sqlite read path next). --- .../src/checkpoint/blob_cache.rs | 72 +++++++++++++++++++ .../src/checkpoint/mod.rs | 1 + 2 files changed, 73 insertions(+) create mode 100644 crates/langgraph-checkpoint/src/checkpoint/blob_cache.rs diff --git a/crates/langgraph-checkpoint/src/checkpoint/blob_cache.rs b/crates/langgraph-checkpoint/src/checkpoint/blob_cache.rs new file mode 100644 index 0000000..9d9582a --- /dev/null +++ b/crates/langgraph-checkpoint/src/checkpoint/blob_cache.rs @@ -0,0 +1,72 @@ +//! Version-keyed cache of deserialized channel values. +//! +//! Savers that store channel values as version-addressed blobs re-parse +//! every blob on every read. For large channels that rarely change (e.g. a +//! static 128KB context), that re-parse is pure overhead: the same +//! `(thread, ns, channel, version)` tuple maps to the same content for the +//! lifetime of the thread. This cache makes repeated reads of unchanged +//! channels O(clone) instead of O(parse). +//! +//! Correctness relies on version strings being monotonic and unique within a +//! `(thread_id, checkpoint_ns)` pair: the engine assigns `{:032}` counters +//! that only grow, and `update_state` derives `max + 1`. The one place a +//! version can be reused is `delete_thread` + recreation of a thread with the +//! same id, so callers MUST call [`BlobCache::remove_thread`] there. + +use parking_lot::RwLock; +use serde_json::Value as JsonValue; +use std::collections::HashMap; + +/// Default maximum number of cached values. On overflow the whole cache is +/// cleared (simple, correct; the cost is a periodic re-parse). +pub const DEFAULT_MAX_ENTRIES: usize = 1024; + +/// Key: `(thread_id, checkpoint_ns, channel, version_str)`. +pub type BlobCacheKey = (String, String, String, String); + +#[derive(Default)] +pub struct BlobCache { + inner: RwLock>, + max_entries: usize, +} + +impl BlobCache { + pub fn new() -> Self { + Self::with_capacity(DEFAULT_MAX_ENTRIES) + } + + pub fn with_capacity(max_entries: usize) -> Self { + Self { + inner: RwLock::new(HashMap::new()), + max_entries, + } + } + + /// Return a clone of the cached value, if present. + pub fn get(&self, key: &BlobCacheKey) -> Option { + self.inner.read().get(key).cloned() + } + + /// Cache a deserialized value, evicting everything when over capacity. + pub fn insert(&self, key: BlobCacheKey, value: JsonValue) { + let mut inner = self.inner.write(); + if inner.len() >= self.max_entries { + inner.clear(); + } + inner.insert(key, value); + } + + /// Drop all cached values for a thread. REQUIRED after `delete_thread`: + /// version strings restart at 1 for a recreated thread with the same id, + /// and stale entries would otherwise be served as the new values. + pub fn remove_thread(&self, thread_id: &str) { + self.inner + .write() + .retain(|(tid, _, _, _), _| tid != thread_id); + } + + #[cfg(test)] + pub fn len(&self) -> usize { + self.inner.read().len() + } +} diff --git a/crates/langgraph-checkpoint/src/checkpoint/mod.rs b/crates/langgraph-checkpoint/src/checkpoint/mod.rs index 69f7ded..55fd1cb 100644 --- a/crates/langgraph-checkpoint/src/checkpoint/mod.rs +++ b/crates/langgraph-checkpoint/src/checkpoint/mod.rs @@ -1,4 +1,5 @@ pub mod base; +pub mod blob_cache; pub mod id; pub mod memory; pub mod types; From ba629c43d31354bac149da8fdde5b299689f7097 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 13:06:03 +0800 Subject: [PATCH 10/15] checkpoint-sqlite: cache deserialized blobs on the read path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit load_blobs now drives reads by the checkpoint's channel_versions, serving cache hits (immutable per (channel, version)) as clones and fetching only the misses from the DB via a (channel, version) IN query. Repeated reads of an unchanged large channel (e.g. a static context) become O(clone) instead of O(DB fetch + parse) — the remaining P0 read-side bottleneck. - SqliteSaver gains an Arc field (shared across clones). - aget_tuple/alist pass channel_versions through; alist's rows of one thread share stable versions, so all but the first are cache hits. - adelete_thread calls BlobCache::remove_thread (required: versions restart at 1 for a recreated thread with the same id). - version_to_string helper shared by dump_blobs and load_blobs keeps write keys and read lookups in sync; SELECT_BLOBS_SQL dropped. - Cache is populated only from the read path (never warmed from aput), so cached values equal exactly what the old parse produced. - Tests: cache hit served with no blob rows in the DB; delete_thread invalidation; empty channel_versions; cleared-channel not returned. --- .../src/queries.rs | 21 -- .../langgraph-checkpoint-sqlite/src/saver.rs | 248 ++++++++++++++++-- 2 files changed, 225 insertions(+), 44 deletions(-) diff --git a/crates/langgraph-checkpoint-sqlite/src/queries.rs b/crates/langgraph-checkpoint-sqlite/src/queries.rs index d3f08bd..21ccc64 100644 --- a/crates/langgraph-checkpoint-sqlite/src/queries.rs +++ b/crates/langgraph-checkpoint-sqlite/src/queries.rs @@ -68,27 +68,6 @@ SELECT FROM checkpoints "#; -/// Fetch all blobs (channel values) for a given checkpoint by joining -/// `checkpoint.channel_versions` with the blobs table. -/// -/// SQLite stores `version` as TEXT, while `je.value` from `json_each` may -/// be a JSON number or string depending on how the checkpoint was -/// produced. The explicit `CAST(... AS TEXT)` normalizes both sides so -/// integer and string versions compare equal. -pub const SELECT_BLOBS_SQL: &str = r#" -SELECT bl.channel, bl.type, bl.blob -FROM checkpoints cp -CROSS JOIN json_each(json_extract(cp.checkpoint, '$.channel_versions')) je -INNER JOIN checkpoint_blobs bl - ON bl.thread_id = cp.thread_id - AND bl.checkpoint_ns = cp.checkpoint_ns - AND bl.channel = je.key - AND bl.version = CAST(je.value AS TEXT) -WHERE cp.thread_id = ?1 - AND cp.checkpoint_ns = ?2 - AND cp.checkpoint_id = ?3 -"#; - /// Fetch pending writes for a given checkpoint, ordered by task and idx. pub const SELECT_WRITES_SQL: &str = r#" SELECT task_id, channel, type, blob, idx, task_path diff --git a/crates/langgraph-checkpoint-sqlite/src/saver.rs b/crates/langgraph-checkpoint-sqlite/src/saver.rs index 73bbf8f..7c9d458 100644 --- a/crates/langgraph-checkpoint-sqlite/src/saver.rs +++ b/crates/langgraph-checkpoint-sqlite/src/saver.rs @@ -12,6 +12,7 @@ use sqlx::Row; use langgraph_checkpoint::checkpoint::base::{ get_checkpoint_id, get_checkpoint_metadata, writes_idx_map, BaseCheckpointSaver, }; +use langgraph_checkpoint::checkpoint::blob_cache::{BlobCache, BlobCacheKey}; use langgraph_checkpoint::checkpoint::types::*; use langgraph_checkpoint::config::RunnableConfig; use langgraph_checkpoint::error::CheckpointError; @@ -23,6 +24,23 @@ use crate::queries::*; /// A dumped blob row: (thread_id, checkpoint_ns, channel, type, checkpoint_id, data) type BlobDumpRow = (String, String, String, String, String, Option>); +/// A blob row fetched to satisfy a cache miss, indexed by (channel, version). +struct FetchedBlob { + type_tag: String, + blob: Option>, +} + +/// Normalize a channel version to the string form used in the `version` +/// column and the BlobCache key. `dump_blobs` and `load_blobs` must agree so +/// write keys and read lookups stay in sync. +fn version_to_string(ver: &JsonValue) -> Option { + match ver { + JsonValue::String(s) => Some(s.clone()), + JsonValue::Number(n) => Some(n.to_string()), + _ => None, + } +} + /// Serialization view of a `Checkpoint` used for the `checkpoints` row. /// /// `channel_values` live in `checkpoint_blobs`, but the row's JSON must stay @@ -62,6 +80,9 @@ impl<'a> From<&'a Checkpoint> for CheckpointRowView<'a> { pub struct SqliteSaver { pool: SqlitePool, serde: Arc, + /// Version-keyed cache of deserialized blob values, so repeated reads of + /// unchanged channels (e.g. a static context) skip the DB fetch + parse. + blob_cache: Arc, } impl SqliteSaver { @@ -70,12 +91,17 @@ impl SqliteSaver { Self { pool, serde: Arc::new(JsonPlusSerializer::new()), + blob_cache: Arc::new(BlobCache::new()), } } /// Create a new SqliteSaver with a custom serializer. pub fn with_serde(pool: SqlitePool, serde: Arc) -> Self { - Self { pool, serde } + Self { + pool, + serde, + blob_cache: Arc::new(BlobCache::new()), + } } /// Create a SqliteSaver from a connection string. @@ -183,31 +209,90 @@ impl SqliteSaver { &self, thread_id: &str, checkpoint_ns: &str, - checkpoint_id: &str, + channel_versions: &ChannelVersions, ) -> Result, CheckpointError> { - let rows = sqlx::query(SELECT_BLOBS_SQL) - .bind(thread_id) - .bind(checkpoint_ns) - .bind(checkpoint_id) + let mut values: HashMap = HashMap::new(); + let mut misses: Vec = Vec::new(); + + // Consult the cache first: blob values are immutable per (channel, + // version), so a hit is a clone with no DB round-trip or parse. This + // is what turns repeated reads of an unchanged large channel (e.g. a + // static context) into O(clone) instead of O(DB fetch + parse). + for (channel, ver) in channel_versions { + let Some(ver_str) = version_to_string(ver) else { + continue; + }; + let key = ( + thread_id.to_string(), + checkpoint_ns.to_string(), + channel.clone(), + ver_str, + ); + if let Some(v) = self.blob_cache.get(&key) { + values.insert(channel.clone(), v); + } else { + misses.push(key); + } + } + if misses.is_empty() { + return Ok(values); + } + + // Fetch only the missing (channel, version) pairs. Row-value IN needs + // SQLite 3.15+ (bundled with sqlx). If a checkpoint ever has more + // pairs than SQLITE_MAX_VARIABLE_NUMBER (32766 by default), the misses + // would need chunking — not needed for realistic channel counts. + let mut qb = sqlx::QueryBuilder::::new( + "SELECT channel, version, type, blob FROM checkpoint_blobs WHERE thread_id = ", + ); + qb.push_bind(thread_id) + .push(" AND checkpoint_ns = ") + .push_bind(checkpoint_ns) + .push(" AND (channel, version) IN ("); + let mut first = true; + for key in &misses { + if !first { + qb.push(", "); + } + first = false; + qb.push("(") + .push_bind(key.2.as_str()) + .push(", ") + .push_bind(key.3.as_str()) + .push(")"); + } + qb.push(")"); + let rows = qb + .build() .fetch_all(&self.pool) .await .map_err(|e| CheckpointError::Storage(e.to_string()))?; - let mut values: HashMap = HashMap::new(); + // Index fetched rows by (channel, version) so each miss maps back + // deterministically. + let mut by_key: HashMap<(String, String), FetchedBlob> = HashMap::new(); for row in rows { let channel: String = row.get("channel"); + let version: String = row.get("version"); let type_tag: String = row.get("type"); let blob: Option> = row.try_get("blob").ok(); + by_key.insert((channel, version), FetchedBlob { type_tag, blob }); + } - if type_tag == "empty" || blob.is_none() { - continue; - } - let bytes = blob.unwrap(); - let val = match self.serde.loads_typed(&type_tag, &bytes) { - Ok(any_val) => any_to_json(any_val), - Err(_) => continue, + for key in misses { + let Some(fetched) = by_key.get(&(key.2.clone(), key.3.clone())) else { + continue; // no blob row: leave out, do not cache }; - values.insert(channel, val); + if fetched.type_tag == "empty" { + continue; // empty marker: leave out, do not cache + } + if let Some(bytes) = &fetched.blob { + if let Ok(any_val) = self.serde.loads_typed(&fetched.type_tag, bytes) { + let val = any_to_json(any_val); + self.blob_cache.insert(key.clone(), val.clone()); + values.insert(key.2, val); + } + } } Ok(values) } @@ -256,10 +341,8 @@ impl SqliteSaver { ) -> Vec { let mut result = Vec::new(); for (channel, ver) in versions { - let ver_str = match ver { - JsonValue::String(s) => s.clone(), - JsonValue::Number(n) => n.to_string(), - _ => continue, + let Some(ver_str) = version_to_string(ver) else { + continue; }; if let Some(val) = values.get(channel) { if let Ok((type_tag, blob)) = self.serde.dumps_typed(val) { @@ -374,11 +457,15 @@ impl SqliteSaver { for row in rows { let mut tuple = Self::row_to_tuple(&row)?; // Merge blob values (version-joined) over the row's channel_values - // (see aget_tuple for the rationale). + // (see aget_tuple for the rationale). Rows of one thread share the + // same stable (channel, version) entries, so after the first row + // the rest are BlobCache hits. let thread_id = row.get::("thread_id"); let ns = row.get::("checkpoint_ns"); let cid = tuple.checkpoint.id.clone(); - let blob_values = self.load_blobs(&thread_id, &ns, &cid).await?; + let blob_values = self + .load_blobs(&thread_id, &ns, &tuple.checkpoint.channel_versions) + .await?; tuple.checkpoint.channel_values.extend(blob_values); tuple.pending_writes = Some(self.load_writes(&thread_id, &ns, &cid).await?); results.push(tuple); @@ -549,10 +636,12 @@ impl BaseCheckpointSaver for SqliteSaver { let mut tuple = Self::row_to_tuple(&row)?; let cid = tuple.checkpoint.id.clone(); // Blob rows are version-addressed and written incrementally, so the - // join in SELECT_BLOBS_SQL resolves every channel's value regardless + // (channel, version) lookup resolves every channel's value regardless // of which checkpoint created the blob. Merge (not replace) so values // still resolve even if a blob row is ever missing. - let blob_values = self.load_blobs(thread_id, checkpoint_ns, &cid).await?; + let blob_values = self + .load_blobs(thread_id, checkpoint_ns, &tuple.checkpoint.channel_versions) + .await?; tuple.checkpoint.channel_values.extend(blob_values); tuple.pending_writes = Some(self.load_writes(thread_id, checkpoint_ns, &cid).await?); Ok(Some(tuple)) @@ -718,6 +807,10 @@ impl BaseCheckpointSaver for SqliteSaver { } async fn adelete_thread(&self, thread_id: String) -> Result<(), CheckpointError> { + // Drop cached blob values for the thread: version strings restart at 1 + // if the thread is recreated with the same id, and stale entries would + // otherwise be served as the new values. + self.blob_cache.remove_thread(&thread_id); let mut tx = self .pool .begin() @@ -1265,4 +1358,113 @@ mod tests { .unwrap(); assert_eq!(parent_id, cp1.id); } + + #[tokio::test] + async fn test_blob_cache_serves_reads_without_db() { + let saver = fresh_saver().await; + let (cp, vers) = make_checkpoint(vec![("ctx", serde_json::json!("static-value"))]); + let cfg = config_for("thread-C"); + saver + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &vers) + .await + .unwrap(); + + // First read populates the cache. + let cfg_id = config_with_id("thread-C", &cp.id); + let t1 = saver.aget_tuple(&cfg_id).await.unwrap().unwrap(); + assert_eq!( + t1.checkpoint.channel_values.get("ctx"), + Some(&serde_json::json!("static-value")) + ); + + // Delete the blob rows behind the saver's back: a cache hit must still + // resolve the value with no blob row to fetch. + sqlx::query("DELETE FROM checkpoint_blobs") + .execute(&saver.pool) + .await + .unwrap(); + + let t2 = saver.aget_tuple(&cfg_id).await.unwrap().unwrap(); + assert_eq!( + t2.checkpoint.channel_values.get("ctx"), + Some(&serde_json::json!("static-value")), + "second read should be served from the BlobCache" + ); + } + + #[tokio::test] + async fn test_delete_thread_invalidates_blob_cache() { + let saver = fresh_saver().await; + let cfg = config_for("thread-D"); + + // First thread incarnation: value at version 1. + let (cp1, vers1) = make_checkpoint(vec![("a", serde_json::json!("first"))]); + saver + .aput(&cfg, cp1.clone(), &CheckpointMetadata::default(), &vers1) + .await + .unwrap(); + let t1 = saver + .aget_tuple(&config_with_id("thread-D", &cp1.id)) + .await + .unwrap() + .unwrap(); + assert_eq!( + t1.checkpoint.channel_values.get("a"), + Some(&serde_json::json!("first")) + ); + + // Delete the thread, then recreate it with the SAME id and a DIFFERENT + // value at the same version. The cache must not serve the stale value + // (version strings restart at 1 for a recreated thread). + saver.adelete_thread("thread-D".to_string()).await.unwrap(); + let (cp2, vers2) = make_checkpoint(vec![("a", serde_json::json!("second"))]); + saver + .aput(&cfg, cp2.clone(), &CheckpointMetadata::default(), &vers2) + .await + .unwrap(); + let t2 = saver + .aget_tuple(&config_with_id("thread-D", &cp2.id)) + .await + .unwrap() + .unwrap(); + assert_eq!( + t2.checkpoint.channel_values.get("a"), + Some(&serde_json::json!("second")), + "delete_thread must invalidate cached values for the recreated thread" + ); + } + + #[tokio::test] + async fn test_load_blobs_empty_channel_versions() { + let saver = fresh_saver().await; + let values = saver.load_blobs("t", "", &HashMap::new()).await.unwrap(); + assert!(values.is_empty()); + } + + #[tokio::test] + async fn test_cleared_channel_not_returned() { + let saver = fresh_saver().await; + let mut cp = Checkpoint::empty(); + cp.id = "cp-cleared".to_string(); + cp.channel_versions + .insert("a".into(), JsonValue::Number(2.into())); + let mut vers = HashMap::new(); + vers.insert("a".into(), JsonValue::Number(2.into())); + // No channel_values entry for `a` → an "empty" marker blob row. + let cfg = config_for("thread-E"); + saver + .aput(&cfg, cp.clone(), &CheckpointMetadata::default(), &vers) + .await + .unwrap(); + + let t = saver + .aget_tuple(&config_with_id("thread-E", &cp.id)) + .await + .unwrap() + .unwrap(); + assert!( + t.checkpoint.channel_values.get("a").is_none(), + "cleared channel must not resolve" + ); + } } From 76bfd5eef4685b8f86264821c8371c427d34383d Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 13:11:03 +0800 Subject: [PATCH 11/15] bench: assert static-context integrity and isolate BlobCache read wins - bench_sqlite_static_context asserts the 128KB static context survives every step and messages accumulate to 2*steps (each invoke appends the input message plus the node's append, matching sanity_bench_graphs_work). - bench_sqlite_static_context_reads: seed once, time one cold get_state (cache miss, parses the 128KB blob) vs N warm reads (cache hits, clone). Measured ~558us cold vs ~72us warm = ~7.7x read-side speedup from the BlobCache. --- tests/bench_pregel.rs | 106 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index 4390f9e..533a6e8 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -756,9 +756,115 @@ async fn bench_sqlite_static_context() { "sqlite static context: {steps:>4} steps, {CONTEXT_SIZE}-byte static channel => {elapsed:?} ({:?}/step)", elapsed / (steps - 1) as u32 ); + + // Correctness: the static context must survive every step (version- + // merged reads + BlobCache must reconstruct it) and messages accumulate. + let snapshot = app.get_state(&config).unwrap(); + let ctx_len = snapshot + .values + .get("context") + .and_then(|c| c.as_str()) + .map(|s| s.len()) + .unwrap_or(0); + assert_eq!(ctx_len, CONTEXT_SIZE, "static context was truncated"); + let msg_count = snapshot + .values + .get("messages") + .and_then(|m| m.as_array()) + .map(|a| a.len()) + .unwrap_or(0); + // Each invoke adds two messages (the input write plus the node's + // append), so `steps` invokes accumulate 2*steps — matching the + // sanity_bench_graphs_work invariant (2 invokes => 4 messages). + assert_eq!( + msg_count, + 2 * steps, + "message count mismatch: got {msg_count}, expected {}", + 2 * steps + ); } } +/// Read-side isolation of the BlobCache: seed a large static channel once, +/// then hammer get_state (pure reads, no writes). Each read must resolve the +/// full static context; with the BlobCache the context blob is parsed once and +/// cloned thereafter. +#[tokio::test(flavor = "multi_thread")] +#[ignore] +async fn bench_sqlite_static_context_reads() { + const CONTEXT_SIZE: usize = 128 * 1024; + const READS: usize = 200; + + let saver = SqliteSaver::from_conn_string("sqlite::memory:") + .await + .unwrap(); + saver.setup().await.unwrap(); + + let mut channels: HashMap> = HashMap::new(); + channels.insert( + "messages".to_string(), + Box::new(BinaryOperatorAggregate::new("messages", add_messages_ref)) + as Box, + ); + channels.insert( + "context".to_string(), + Box::new(LastValue::new("context")) as Box, + ); + + let mut graph = StateGraph::new(channels); + graph + .add_node( + "append", + |input: JsonValue, _config: RunnableConfig| async move { + let n = input + .get("messages") + .and_then(|m| m.as_array()) + .map(|a| a.len()) + .unwrap_or(0); + Ok(json!({"messages": [make_message(n)]})) + }, + ) + .unwrap(); + graph.add_edge(START, "append").unwrap(); + graph.add_edge("append", END).unwrap(); + let app = graph + .compile_builder() + .checkpointer(Arc::new(saver)) + .build() + .unwrap(); + + let mut config = RunnableConfig::new(); + config.insert( + "configurable".to_string(), + json!({"thread_id": "bench-ctx-reads"}), + ); + + // Seed the static channel once. + app.ainvoke( + &json!({"messages": [make_message(0)], "context": "x".repeat(CONTEXT_SIZE)}), + &config, + ) + .await + .unwrap(); + + // Cold read: the cache is empty, so the 128KB context blob must be + // fetched and parsed from the DB. + let cold_start = Instant::now(); + let _ = app.get_state(&config).unwrap(); + let cold = cold_start.elapsed(); + + // Warm reads: BlobCache hits, clone-only. + let start = Instant::now(); + for _ in 0..READS { + let _ = app.get_state(&config).unwrap(); + } + let elapsed = start.elapsed(); + println!( + "sqlite static context reads: cold {cold:?} (cache miss, parses blob) vs {READS} warm get_state reads of {CONTEXT_SIZE}-byte static channel => {elapsed:?} ({:?}/read warm)", + elapsed / READS as u32 + ); +} + /// Guards against the benchmark graphs silently becoming no-ops. #[tokio::test] async fn sanity_bench_graphs_work() { From 2fbea25a9a8c8e23427fb5b143b240399b562697 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 13:12:31 +0800 Subject: [PATCH 12/15] style: run rustfmt on bench_pregel --- tests/bench_pregel.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index 533a6e8..08b4236 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -803,8 +803,7 @@ async fn bench_sqlite_static_context_reads() { let mut channels: HashMap> = HashMap::new(); channels.insert( "messages".to_string(), - Box::new(BinaryOperatorAggregate::new("messages", add_messages_ref)) - as Box, + Box::new(BinaryOperatorAggregate::new("messages", add_messages_ref)) as Box, ); channels.insert( "context".to_string(), From 3b3cf790ef4ed14fd8a29cc555ec32fa9948ff9a Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 14:57:46 +0800 Subject: [PATCH 13/15] perf: cache a process-global Tokio runtime for sync invoke() entry points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sync invoke() entry points (CompiledStateGraph::invoke, RunnableCallable::invoke) built and dropped a full multi-thread Runtime on every call from outside a tokio context — each construction spawns num_cpus worker threads (~100us+). A single OnceLock-cached multi-thread runtime amortizes that; Handle::try_current() guard is preserved, so callers already inside a runtime keep using their own handle (no create/cache). Also unifies RunnableCallable::invoke's cold path to install the task-local config via with_config (matching ainvoke), fixing a pre-existing inconsistency where the cold path skipped it. New bench_sync_invoke_runtime_cache probe (plain #[test], outside tokio): 820.7us/invoke -> 3.3us/invoke (~248x) on the sync invoke path. --- crates/langgraph-core/src/config.rs | 31 ++++++++++++- crates/langgraph-core/src/graph/state.rs | 11 ++--- .../langgraph-core/src/runnable/callable.rs | 18 ++++---- tests/bench_pregel.rs | 44 +++++++++++++++++++ 4 files changed, 86 insertions(+), 18 deletions(-) diff --git a/crates/langgraph-core/src/config.rs b/crates/langgraph-core/src/config.rs index 71967bc..01d45b1 100644 --- a/crates/langgraph-core/src/config.rs +++ b/crates/langgraph-core/src/config.rs @@ -3,7 +3,9 @@ use langgraph_checkpoint::config::RunnableConfig; use langgraph_checkpoint::store::base::BaseStore; use serde_json::Value as JsonValue; use std::cell::RefCell; -use std::sync::Arc; +use std::future::Future; +use std::sync::{Arc, OnceLock}; +use tokio::runtime::Runtime as TokioRuntime; use tokio::sync::mpsc; // Task-local config for async contexts @@ -117,3 +119,30 @@ where }); result } + +/// Block on a future, reusing a process-global Tokio runtime when the caller +/// is outside a runtime. +/// +/// The sync `invoke()` entry points used to build and drop a full multi-thread +/// `Runtime` on every call made outside a Tokio context — each construction +/// spawns `num_cpus` worker threads (~100µs+). A single cached runtime +/// amortizes that cost while preserving the exact scheduler semantics a +/// per-call `Runtime::new()` provided (multi-thread, IO + time enabled), +/// including the `block_in_place` bridge checkpoint savers rely on (it checks +/// `runtime_flavor() != CurrentThread`). When the caller is already inside a +/// runtime, `block_on` on its handle instead — no runtime is created or cached. +pub fn block_on(fut: F) -> F::Output +where + F: Future, +{ + match tokio::runtime::Handle::try_current() { + Ok(handle) => handle.block_on(fut), + Err(_) => CACHED_RUNTIME + .get_or_init(|| { + tokio::runtime::Runtime::new().expect("failed to build the shared Tokio runtime") + }) + .block_on(fut), + } +} + +static CACHED_RUNTIME: OnceLock = OnceLock::new(); diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index f08e5d8..3dcec9e 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -1867,13 +1867,10 @@ impl Runnable for CompiledStateGraph { input: &JsonValue, config: &RunnableConfig, ) -> Result { - // Block on the async implementation - match tokio::runtime::Handle::try_current() { - Ok(handle) => handle.block_on(self.run_pregel(input, config)), - Err(_) => tokio::runtime::Runtime::new() - .unwrap() - .block_on(self.run_pregel(input, config)), - } + // Block on the async implementation, reusing the caller's runtime or a + // process-global cached one (the old per-call `Runtime::new()` spawned + // worker threads on every sync invoke from outside a tokio context). + crate::config::block_on(self.run_pregel(input, config)) } async fn ainvoke( diff --git a/crates/langgraph-core/src/runnable/callable.rs b/crates/langgraph-core/src/runnable/callable.rs index 24ab768..e84f3e9 100644 --- a/crates/langgraph-core/src/runnable/callable.rs +++ b/crates/langgraph-core/src/runnable/callable.rs @@ -71,16 +71,14 @@ impl Runnable for RunnableCallable { let input = input.clone(); let config = config.clone(); - // Try to use existing tokio runtime, otherwise create one - match tokio::runtime::Handle::try_current() { - Ok(handle) => handle.block_on(crate::config::with_config( - config.clone(), - func(input, config), - )), - Err(_) => tokio::runtime::Runtime::new() - .unwrap() - .block_on(func(input, config)), - } + // Reuse the caller's runtime, or the process-global cached one (the old + // cold path built a fresh Runtime per call). `with_config` installs the + // task-local config so `get_config()` works inside the callable, + // matching `ainvoke`, which always wraps. + crate::config::block_on(crate::config::with_config( + config.clone(), + func(input, config), + )) } async fn ainvoke( diff --git a/tests/bench_pregel.rs b/tests/bench_pregel.rs index 08b4236..aea06a2 100644 --- a/tests/bench_pregel.rs +++ b/tests/bench_pregel.rs @@ -889,3 +889,47 @@ async fn sanity_bench_graphs_work() { Some(4) ); } + +/// P2-1: sync `invoke()` from outside a tokio context used to build and drop a +/// full multi-thread Runtime per call (each construction spawns worker threads, +/// ~100µs+). A process-global cached runtime amortizes that. +/// +/// Plain `#[test]` — NOT `#[tokio::test]` — so this runs outside a tokio +/// context and actually hits the `Runtime::new()`/cached-runtime path in +/// `CompiledStateGraph::invoke`. +#[test] +#[ignore] +fn bench_sync_invoke_runtime_cache() { + let mut channels: HashMap> = HashMap::new(); + channels.insert( + "out".to_string(), + Box::new(LastValue::new("out")) as Box, + ); + let mut graph = StateGraph::new(channels); + graph + .add_node("n", |_: JsonValue, _: RunnableConfig| async move { + Ok(json!({"out": 1})) + }) + .unwrap(); + graph.add_edge(START, "n").unwrap(); + graph.add_edge("n", END).unwrap(); + let app = graph.compile().unwrap(); + let config = RunnableConfig::new(); + + // Warm up: the first invoke builds the shared runtime. + assert_eq!( + app.invoke(&json!({}), &config).unwrap().get("out"), + Some(&json!(1)) + ); + + const N: usize = 200; + let t = Instant::now(); + for _ in 0..N { + app.invoke(&json!({}), &config).unwrap(); + } + let elapsed = t.elapsed(); + println!( + "sync invoke x{N} (outside tokio) => {elapsed:?} ({:?}/invoke)", + elapsed / N as u32 + ); +} From 16fb2a1770929fb39f7401c7cd3692943e405a15 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 15:02:04 +0800 Subject: [PATCH 14/15] perf: drop per-task config clone in PregelRunner by mutating task.config in place MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit execute_task / execute_single_task_sync cloned task.config to inject CONFIG_KEY_SEND, leaving the original untouched. task.config is dead after the super-step (nothing downstream reads it), so mutate it in place: the no-store/no-stream hot path now makes 0 config clones per task, and the with_runtime branch drops from 2 clones to 1 (the remaining clone is structural — with_runtime owns a config for the task-local while the node borrows one). bench_sequential_chain @1000 nodes: 2.5909ms -> 2.3974ms (2.59 -> 2.397 us/node, ~7.5%). bench_react_agent: 50-turn ~3% better, 100/200 within noise (history serialization dominates). --- crates/langgraph-core/src/pregel/runner.rs | 30 ++++++++++++++-------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/crates/langgraph-core/src/pregel/runner.rs b/crates/langgraph-core/src/pregel/runner.rs index 2bf91dc..01e1ffe 100644 --- a/crates/langgraph-core/src/pregel/runner.rs +++ b/crates/langgraph-core/src/pregel/runner.rs @@ -139,9 +139,13 @@ impl PregelRunner { runtime: Option<&Arc>, stream_writer: Option, ) -> Result<(), RunnerError> { - let mut config = task.config.clone(); + // Inject CONFIG_KEY_SEND into task.config in place. task.config is dead + // after this super-step (nothing downstream reads it — writes, triggers + // and the checkpoint derive from other fields), so mutating it avoids a + // per-task config clone on every super-step. { - let configurable = config + let configurable = task + .config .entry("configurable".to_string()) .or_insert_with(|| serde_json::json!({})); if let Some(obj) = configurable.as_object_mut() { @@ -175,12 +179,16 @@ impl PregelRunner { }; let result = if let Some(ref rt) = effective_runtime { - config::with_runtime(config.clone(), rt.clone(), async { - task.proc.ainvoke(&task.input, &config).await + // with_runtime must own a config for the task-local while the node + // borrows one — one clone per task is unavoidable here (this branch + // is off the no-store/no-stream hot path). + let config = task.config.clone(); + config::with_runtime(config, rt.clone(), async { + task.proc.ainvoke(&task.input, &task.config).await }) .await } else { - task.proc.ainvoke(&task.input, &config).await + task.proc.ainvoke(&task.input, &task.config).await }; match result { @@ -212,9 +220,10 @@ impl PregelRunner { task: &mut PregelExecutableTask, runtime: Option<&Arc>, ) -> Result<(), RunnerError> { - let mut config = task.config.clone(); + // Same in-place CONFIG_KEY_SEND injection as the async path. { - let configurable = config + let configurable = task + .config .entry("configurable".to_string()) .or_insert_with(|| serde_json::json!({})); if let Some(obj) = configurable.as_object_mut() { @@ -226,11 +235,12 @@ impl PregelRunner { } let result = if let Some(rt) = runtime { - config::with_runtime_sync(config.clone(), rt.clone(), || { - task.proc.invoke(&task.input, &config) + let config = task.config.clone(); + config::with_runtime_sync(config, rt.clone(), || { + task.proc.invoke(&task.input, &task.config) }) } else { - task.proc.invoke(&task.input, &config) + task.proc.invoke(&task.input, &task.config) }; match result { From f1f3d64ece57f94be7b0d7646ddc4cebdaa05e62 Mon Sep 17 00:00:00 2001 From: MuFengMuXue <3058704216@qq.com> Date: Tue, 4 Aug 2026 15:09:10 +0800 Subject: [PATCH 15/15] perf: drain task write buffers in apply_writes instead of cloning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit apply_writes grouped writes by channel via `val.clone()` per write while borrowing tasks immutably. Since the streaming Updates emit (state.rs) runs over task.writes before apply_writes, and the get_state snapshot path now computes `next` (writes.is_empty()) before the call, nothing reads task.writes afterward — so take each task's write buffer by value (&mut [PregelExecutableTask] + mem::take) and move the values instead of cloning them. get_state's `next` computation is reordered before the call; the interrupt-only apply_completed_writes path is untouched (cold, mutually exclusive). Correctness: 37 core + 14 prebuilt tests green. Measurement: chain and multi-step-loop benches are within run-to-run noise on small writes (the saving is one value-clone per write, which grows with write size — worth the most for large node outputs). --- crates/langgraph-core/src/graph/state.rs | 23 +++++++++++++---------- crates/langgraph-core/src/pregel/algo.rs | 21 ++++++++++----------- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/crates/langgraph-core/src/graph/state.rs b/crates/langgraph-core/src/graph/state.rs index 3dcec9e..0db3deb 100644 --- a/crates/langgraph-core/src/graph/state.rs +++ b/crates/langgraph-core/src/graph/state.rs @@ -684,9 +684,19 @@ impl CompiledStateGraph { .unwrap_or(0); JsonValue::String(format!("{:032}", max_version + 1)) }; + // Build next BEFORE apply_writes: apply_writes drains each task's write + // buffer, so `writes.is_empty()` must be read before the writes are + // moved out. Identical result — apply_writes never modified task.writes + // before this change. + let next: Vec = tasks + .iter() + .filter(|t| t.writes.is_empty()) + .map(|t| t.name.clone()) + .collect(); + apply_writes( &mut channels, - &tasks, + &mut tasks, &mut versions_seen, &mut channel_versions, trigger_to_nodes, @@ -702,13 +712,6 @@ impl CompiledStateGraph { .collect(); let values = read_channels(&channels, &output_keys); - // Build next: names of tasks that have NOT written yet - let next: Vec = tasks - .iter() - .filter(|t| t.writes.is_empty()) - .map(|t| t.name.clone()) - .collect(); - // Extract interrupts from pending writes let interrupts: Vec = saved .pending_writes @@ -1663,7 +1666,7 @@ impl CompiledStateGraph { })) }; - let (tasks, task_result) = runner.run_tasks(tasks).await; + let (mut tasks, task_result) = runner.run_tasks(tasks).await; match task_result { Ok(()) => {} @@ -1762,7 +1765,7 @@ impl CompiledStateGraph { let next_version = JsonValue::String(format!("{:032}", running_max_version)); let updated = apply_writes( &mut channels, - &tasks, + &mut tasks, &mut versions_seen, &mut channel_versions, trigger_to_nodes, diff --git a/crates/langgraph-core/src/pregel/algo.rs b/crates/langgraph-core/src/pregel/algo.rs index 3013607..d6e6f57 100644 --- a/crates/langgraph-core/src/pregel/algo.rs +++ b/crates/langgraph-core/src/pregel/algo.rs @@ -231,7 +231,7 @@ fn create_scratchpad( /// Returns the set of updated channel names. pub fn apply_writes( channels: &mut HashMap>, - tasks: &[PregelExecutableTask], + tasks: &mut [PregelExecutableTask], versions_seen: &mut HashMap>, channel_versions: &mut ChannelVersions, trigger_to_nodes: &TriggerToNodes, @@ -245,7 +245,7 @@ pub fn apply_writes( let bump_step = tasks.iter().any(|t| !t.triggers.is_empty()); // 1. Update versions_seen for each task's trigger channels - for task in tasks { + for task in tasks.iter() { let seen = versions_seen.entry(task.name.clone()).or_default(); for trigger in &task.triggers { if let Some(ver) = channel_versions.get(trigger) { @@ -273,19 +273,18 @@ pub fn apply_writes( } } - // 4. Group writes by channel. - // Filter out all reserved keys (NO_WRITES, PUSH, RESUME, INTERRUPT, - // RETURN, ERROR, config keys, etc.) — only real channel writes proceed. + // 4. Group writes by channel, DRAINING each task's write buffer. + // Streaming Updates (state.rs) already ran over `task.writes`, and the + // get_state snapshot path computes `next` before calling apply_writes, + // so nothing reads task.writes after this point — move the values out + // instead of cloning them. let mut writes_by_channel: HashMap> = HashMap::new(); - for task in tasks { - for (chan, val) in &task.writes { + for task in tasks.iter_mut() { + for (chan, val) in std::mem::take(&mut task.writes) { if RESERVED.contains(&chan.as_str()) { continue; } - writes_by_channel - .entry(chan.clone()) - .or_default() - .push(val.clone()); + writes_by_channel.entry(chan).or_default().push(val); } }