diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index 0bfd4ac8b..fcdcc26ad 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -77,11 +77,17 @@ type CompiledGraph[StateT any] struct { } type Context struct { - engine *RustEngine + engine *RustEngine + resumeEvent *WaitEvent } func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) { - return c.engine.WaitAnyOf(cond) + if c.resumeEvent != nil { + event := *c.resumeEvent + c.resumeEvent = nil + return event, nil + } + return WaitEvent{}, ErrWaitRequested{Condition: cond} } func (c *Context) PublishToChannel(channel string, value any) error { @@ -134,7 +140,8 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*H if fallbackState == nil { return Command{}, fmt.Errorf("node `%s` expected map state argument", node) } - return fn(&Context{engine: engine}, nodeInput, fallbackState) + resolvedInput, resumeEvent := unwrapResumeInput(nodeInput) + return fn(&Context{engine: engine, resumeEvent: resumeEvent}, resolvedInput, fallbackState) }, ) if err != nil { @@ -304,3 +311,24 @@ func mustTypeOf[T any]() reflect.Type { return reflect.TypeOf((*T)(nil)).Elem() } +func unwrapResumeInput(input any) (any, *WaitEvent) { + wrapper, ok := input.(map[string]any) + if !ok { + return input, nil + } + rawArg, hasArg := wrapper["__lg_resume_arg__"] + rawEvent, hasEvent := wrapper["__lg_resume_event__"] + if !hasArg || !hasEvent { + return input, nil + } + eventPayload, err := json.Marshal(rawEvent) + if err != nil { + return rawArg, nil + } + var event WaitEvent + if err := json.Unmarshal(eventPayload, &event); err != nil { + return rawArg, nil + } + return rawArg, &event +} + diff --git a/langgraph-go/advancedgraph/rust_engine.go b/langgraph-go/advancedgraph/rust_engine.go index a7df2f71b..163c0471d 100644 --- a/langgraph-go/advancedgraph/rust_engine.go +++ b/langgraph-go/advancedgraph/rust_engine.go @@ -51,6 +51,9 @@ func goNodeCallback(userData C.ulong, node *C.char, argJSON *C.char, stateJSON * cmd, err := ctx.exec(nodeName, nodeInput, state) if err != nil { + if waitReq, ok := AsErrWaitRequested(err); ok { + return cCallbackEnvelopeSuspend(waitReq.Condition) + } return cCallbackEnvelopeError(err.Error()) } @@ -201,6 +204,17 @@ func cCallbackEnvelopeError(message string) *C.char { return C.CString(string(raw)) } +func cCallbackEnvelopeSuspend(cond AnyOfCondition) *C.char { + raw, _ := json.Marshal(map[string]any{ + "ok": true, + "suspend": map[string]any{ + "kind": "any_of", + "any_of": cond, + }, + }) + return C.CString(string(raw)) +} + func coerceJSONValue(v any) any { switch t := v.(type) { case map[string]any: diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go index 8960344a3..8be5be54a 100644 --- a/langgraph-go/advancedgraph/types.go +++ b/langgraph-go/advancedgraph/types.go @@ -1,6 +1,9 @@ package advancedgraph -import "encoding/json" +import ( + "encoding/json" + "errors" +) type WaitCondition interface { toAny() map[string]any @@ -63,6 +66,22 @@ type Command struct { Goto []Send } +type ErrWaitRequested struct { + Condition AnyOfCondition +} + +func (e ErrWaitRequested) Error() string { + return "wait requested" +} + +func AsErrWaitRequested(err error) (ErrWaitRequested, bool) { + var target ErrWaitRequested + if !errors.As(err, &target) { + return target, false + } + return target, true +} + func DecodeString(raw json.RawMessage) string { var s string _ = json.Unmarshal(raw, &s) diff --git a/libs/langgraph/langgraph/advanced_graph/state.py b/libs/langgraph/langgraph/advanced_graph/state.py index 7929c335a..efb6f4701 100644 --- a/libs/langgraph/langgraph/advanced_graph/state.py +++ b/libs/langgraph/langgraph/advanced_graph/state.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import inspect +import threading from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import timedelta @@ -38,6 +39,12 @@ class AnyOfCondition: WaitCondition = ChannelCondition | TimerCondition +class WaitRequested(Exception): + def __init__(self, payload: dict[str, Any]) -> None: + super().__init__("wait requested") + self.payload = payload + + class AdvancedStateGraph(Generic[StateT]): """Experimental in-memory graph engine with async channels.""" @@ -74,20 +81,14 @@ class AdvancedStateGraph(Generic[StateT]): raise ValueError(f"Channel `{name}` already exists") self._async_channels[name] = _ChannelSpec(typ=typ) - def set_entry_point(self, name_or_node: str | Callable[..., Any]) -> None: - self._entry_point = self._resolve_node_name(name_or_node) - - def set_finish_point(self, name_or_node: str | Callable[..., Any]) -> None: - self._finish_point = self._resolve_node_name(name_or_node) - def add_entry_node(self, node: Callable[..., Any]) -> str: node_name = self.add_node(node) - self.set_entry_point(node_name) + self._entry_point = self._resolve_node_name(node_name) return node_name def add_finish_node(self, node: Callable[..., Any]) -> str: node_name = self.add_node(node) - self.set_finish_point(node_name) + self._finish_point = self._resolve_node_name(node_name) return node_name def _resolve_node_name(self, name_or_node: str | Callable[..., Any]) -> str: @@ -153,7 +154,10 @@ class Context: self._run = run async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any: - return await self._run.wait_for(target) + resumed = self._run._consume_resume_event(target) + if resumed is not None: + return resumed + raise WaitRequested(_target_to_suspend_payload(target)) def publish_to_channel(self, channel: str, value: Any) -> None: self._run.publish_nowait(channel, value) @@ -199,6 +203,7 @@ class _GraphEngineRun: self._tasks: set[asyncio.Task[list[Send]]] = set() self._finished = False self._state: Any = None + self._local = threading.local() self.context = Context(self) async def run(self, initial_state: StateT) -> StateT: @@ -252,12 +257,19 @@ class _GraphEngineRun: def _execute_node_for_rust( self, node_name: str, node_input: Any, state: Any ) -> dict[str, Any]: + node_input, resume_event = _unwrap_resume_input(node_input) + self._set_resume_event(resume_event) if node_name not in self._nodes: raise ValueError(f"Unknown node `{node_name}`") node = self._nodes[node_name] - result = _invoke_node(node, self.context, node_input, state) - if inspect.isawaitable(result): - result = asyncio.run(cast(Coroutine[Any, Any, Any], result)) + try: + result = _invoke_node(node, self.context, node_input, state) + if inspect.isawaitable(result): + result = asyncio.run(cast(Coroutine[Any, Any, Any], result)) + except WaitRequested as suspend: + return {"suspend": suspend.payload} + finally: + self._set_resume_event(None) if isinstance(result, Command): update = result.update @@ -274,6 +286,16 @@ class _GraphEngineRun: ], } + def _set_resume_event(self, event: dict[str, Any] | None) -> None: + self._local.resume_event = event + + def _consume_resume_event(self, target: WaitCondition | AnyOfCondition) -> Any | None: + event = cast(dict[str, Any] | None, getattr(self._local, "resume_event", None)) + if event is None: + return None + self._local.resume_event = None + return event + def _normalize_result_to_sends(result: Any, *, default_input: Any) -> list[Send]: if result is None: return [] @@ -365,6 +387,29 @@ def _condition_to_rust(condition: WaitCondition) -> dict[str, Any]: raise TypeError(f"Unsupported condition type: {type(condition)!r}") +def _target_to_suspend_payload(target: WaitCondition | AnyOfCondition) -> dict[str, Any]: + if isinstance(target, AnyOfCondition): + return { + "kind": "any_of", + "any_of": { + "conditions": [_condition_to_rust(cond) for cond in target.conditions] + }, + } + return {"kind": "condition", "condition": _condition_to_rust(target)} + + +def _unwrap_resume_input(node_input: Any) -> tuple[Any, dict[str, Any] | None]: + if not isinstance(node_input, dict): + return node_input, None + if "__lg_resume_arg__" not in node_input or "__lg_resume_event__" not in node_input: + return node_input, None + resume_arg = node_input["__lg_resume_arg__"] + resume_event = node_input["__lg_resume_event__"] + if isinstance(resume_event, dict): + return resume_arg, resume_event + return resume_arg, None + + def _infer_node_name(node: Callable[..., Any]) -> str: node_name = getattr(node, "__name__", "") if not node_name or node_name == "": diff --git a/libs/langgraph/tests/advanced-graph/test_run_pool_size.py b/libs/langgraph/tests/advanced-graph/test_run_pool_size.py new file mode 100644 index 000000000..db79c0046 --- /dev/null +++ b/libs/langgraph/tests/advanced-graph/test_run_pool_size.py @@ -0,0 +1,55 @@ +import os +import subprocess +import sys + + +def test_run_pool_size_one_still_allows_parallel_runs() -> None: + script = r""" +import asyncio +import time +from typing_extensions import TypedDict +from langgraph.advanced_graph import AdvancedStateGraph, Context, timer_condition +from langgraph.types import Command, Send + + +class RunState(TypedDict): + done: bool + + +async def wait_node(ctx: Context, _: object, state: RunState) -> Command: + await ctx.wait_for(timer_condition(seconds=0.2)) + return Command(goto=Send("finish_node", None), update=state) + + +async def finish_node(_: object, state: RunState) -> dict[str, bool]: + return {"done": True} + + +async def main() -> None: + graph = AdvancedStateGraph(RunState) + graph.add_entry_node(wait_node) + graph.add_finish_node(finish_node) + compiled = graph.compile() + started = time.perf_counter() + await asyncio.gather( + compiled.ainvoke({"done": False}), + compiled.ainvoke({"done": False}), + ) + elapsed = time.perf_counter() - started + print(f"{elapsed:.6f}") + + +asyncio.run(main()) +""" + env = os.environ.copy() + env["LANGGRAPH_RUN_POOL_SIZE"] = "1" + env.setdefault("LANGGRAPH_NODE_POOL_SIZE", "2") + completed = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=True, + ) + elapsed = float(completed.stdout.strip().splitlines()[-1]) + assert elapsed < 0.35, completed.stdout diff --git a/rust-core/Cargo.lock b/rust-core/Cargo.lock index 10132774b..082228bd7 100644 --- a/rust-core/Cargo.lock +++ b/rust-core/Cargo.lock @@ -50,6 +50,7 @@ dependencies = [ "pyo3", "serde", "serde_json", + "tokio", ] [[package]] @@ -111,6 +112,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "portable-atomic" version = "1.13.1" @@ -285,6 +292,27 @@ version = "0.12.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +[[package]] +name = "tokio" +version = "1.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +dependencies = [ + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/rust-core/Cargo.toml b/rust-core/Cargo.toml index 79c508e66..70fcbd3ac 100644 --- a/rust-core/Cargo.toml +++ b/rust-core/Cargo.toml @@ -17,4 +17,5 @@ serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" parking_lot = "0.12" libc = "0.2" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] } diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index 6658df5af..c9f82180f 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -1,14 +1,16 @@ -use parking_lot::{Condvar, Mutex}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::{HashMap, VecDeque}; use std::env; +use std::future::Future; use std::sync::mpsc; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; use std::thread; use std::time::{Duration, Instant}; +use tokio::runtime::Runtime; +use tokio::sync::Notify; #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(tag = "kind")] @@ -36,6 +38,15 @@ pub enum WaitEvent { Timer { seconds: f64 }, } +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(tag = "kind")] +pub enum WaitRequest { + #[serde(rename = "condition")] + Condition { condition: WaitCondition }, + #[serde(rename = "any_of")] + AnyOf { any_of: AnyOfCondition }, +} + pub struct SendPayload { pub node: String, pub arg: A, @@ -46,6 +57,11 @@ pub struct NodeExecResult { pub sends: Vec>, } +pub enum NodeOutcome { + Completed(NodeExecResult), + Suspended { wait: WaitRequest }, +} + type Task = Box; fn debug_enabled() -> bool { @@ -70,6 +86,13 @@ fn debug_log(message: &str) { } } +fn pool_size_from_env(var_name: &str, default: usize, min: usize) -> usize { + let parsed = env::var(var_name) + .ok() + .and_then(|raw| raw.trim().parse::().ok()); + parsed.unwrap_or(default).max(min) +} + struct ThreadPool { tx: mpsc::Sender, _workers: Vec>, @@ -122,9 +145,10 @@ where debug_log("node_pool_execute() called"); static NODE_POOL: OnceLock = OnceLock::new(); let pool = NODE_POOL.get_or_init(|| { - let size = thread::available_parallelism() + let default_size = thread::available_parallelism() .map(|n| n.get().max(2)) .unwrap_or(4); + let size = pool_size_from_env("LANGGRAPH_NODE_POOL_SIZE", default_size, 1); ThreadPool::new(size, "langgraph-node") }); pool.execute(task) @@ -135,9 +159,39 @@ where F: FnOnce() + Send + 'static, { debug_log("run_loop_pool_execute() called"); - static RUN_LOOP_POOL: OnceLock = OnceLock::new(); - let pool = RUN_LOOP_POOL.get_or_init(|| ThreadPool::new(2, "langgraph-runloop")); - pool.execute(task) + run_runtime().spawn_blocking(task); + Ok(()) +} + +pub fn run_loop_spawn(future: F) -> Result<(), String> +where + F: Future + Send + 'static, +{ + run_runtime().spawn(future); + Ok(()) +} + +pub fn run_loop_block_on(future: F) -> F::Output +where + F: Future, +{ + run_runtime().block_on(future) +} + +fn run_runtime() -> &'static Runtime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + let default_size = thread::available_parallelism() + .map(|n| n.get().max(2)) + .unwrap_or(2); + let worker_threads = pool_size_from_env("LANGGRAPH_RUN_POOL_SIZE", default_size, 1); + tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .thread_name("langgraph-runloop") + .enable_all() + .build() + .expect("failed to build tokio runtime") + }) } pub fn run_scheduler_loop( @@ -227,8 +281,8 @@ pub fn merge_json_update(state: &mut Value, update: Option) { #[derive(Clone, Default)] pub struct Engine { - channels: Arc>>>, - channel_notify: Arc, + channels: Arc>>>, + channel_notify: Arc, } impl Engine { @@ -239,78 +293,55 @@ impl Engine { pub fn add_async_channel(&self, name: &str) { debug_log(&format!("Engine::add_async_channel(name={name})")); - let mut channels = self.channels.lock(); + let mut channels = self.channels.lock().expect("channels mutex poisoned"); channels.entry(name.to_owned()).or_default(); } pub fn publish_json(&self, channel: &str, value: serde_json::Value) -> Result<(), String> { debug_log(&format!("Engine::publish_json(channel={channel})")); - let mut channels = self.channels.lock(); + let mut channels = self.channels.lock().expect("channels mutex poisoned"); let queue = channels .get_mut(channel) .ok_or_else(|| format!("Unknown channel `{channel}`"))?; queue.push_back(value); - // Wake up waiters blocked on channel conditions/any_of. - self.channel_notify.notify_all(); + self.channel_notify.notify_waiters(); Ok(()) } - pub fn wait_for(&self, cond: &WaitCondition) -> Result { - debug_log(&format!("Engine::wait_for(cond={cond:?})")); + pub async fn wait_request_async(&self, wait: &WaitRequest) -> Result { + match wait { + WaitRequest::Condition { condition } => self.wait_for_async(condition).await, + WaitRequest::AnyOf { any_of } => self.wait_for_any_of_async(any_of).await, + } + } + + pub async fn wait_for_async(&self, cond: &WaitCondition) -> Result { + debug_log(&format!("Engine::wait_for_async(cond={cond:?})")); match cond { WaitCondition::Channel { channel, n } => { if *n < 1 { return Err("channel condition n must be >= 1".to_string()); } - let mut channels = self.channels.lock(); loop { - let queue = channels - .get_mut(channel) - .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - if queue.len() >= *n { - debug_log(&format!( - "Engine::wait_for channel ready (channel={channel}, n={n})" - )); - if *n == 1 { - if let Some(value) = queue.pop_front() { - return Ok(WaitEvent::Channel { - channel: channel.clone(), - value, - }); - } - } else { - let mut values = Vec::with_capacity(*n); - for _ in 0..*n { - if let Some(v) = queue.pop_front() { - values.push(v); - } - } - return Ok(WaitEvent::Channel { - channel: channel.clone(), - value: serde_json::Value::Array(values), - }); - } + if let Some(event) = self.try_take_channel_event(channel, *n)? { + return Ok(event); } - debug_log(&format!( - "Engine::wait_for waiting on channel condvar (channel={channel}, n={n})" - )); - self.channel_notify.wait(&mut channels); + self.channel_notify.notified().await; } } WaitCondition::Timer { seconds } => { if *seconds <= 0.0 { return Err("timer condition must be > 0".to_string()); } - std::thread::sleep(Duration::from_secs_f64(*seconds)); - debug_log(&format!("Engine::wait_for timer fired (seconds={seconds})")); + tokio::time::sleep(Duration::from_secs_f64(*seconds)).await; Ok(WaitEvent::Timer { seconds: *seconds }) } } } - pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result { + pub async fn wait_for_any_of_async(&self, any_of: &AnyOfCondition) -> Result { debug_log(&format!( - "Engine::wait_for_any_of(conditions={})", + "Engine::wait_for_any_of_async(conditions={})", any_of.conditions.len() )); if any_of.conditions.is_empty() { @@ -328,39 +359,14 @@ impl Engine { } } - let mut channels = self.channels.lock(); loop { for cond in &any_of.conditions { if let WaitCondition::Channel { channel, n } = cond { if *n < 1 { return Err("channel condition n must be >= 1".to_string()); } - let queue = channels - .get_mut(channel) - .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - if queue.len() >= *n { - debug_log(&format!( - "Engine::wait_for_any_of channel hit (channel={channel}, n={n})" - )); - if *n == 1 { - if let Some(value) = queue.pop_front() { - return Ok(WaitEvent::Channel { - channel: channel.clone(), - value, - }); - } - } else { - let mut values = Vec::with_capacity(*n); - for _ in 0..*n { - if let Some(v) = queue.pop_front() { - values.push(v); - } - } - return Ok(WaitEvent::Channel { - channel: channel.clone(), - value: serde_json::Value::Array(values), - }); - } + if let Some(event) = self.try_take_channel_event(channel, *n)? { + return Ok(event); } } } @@ -369,21 +375,59 @@ impl Engine { let timeout = Duration::from_secs_f64(seconds); let elapsed = started.elapsed(); if elapsed >= timeout { - debug_log(&format!( - "Engine::wait_for_any_of timer hit (seconds={seconds})" - )); return Ok(WaitEvent::Timer { seconds }); } let remaining = timeout.saturating_sub(elapsed); - debug_log(&format!( - "Engine::wait_for_any_of waiting on condvar with timeout {:?}", - remaining - )); - self.channel_notify.wait_for(&mut channels, remaining); + tokio::select! { + _ = self.channel_notify.notified() => {} + _ = tokio::time::sleep(remaining) => { + return Ok(WaitEvent::Timer { seconds }); + } + } } else { - debug_log("Engine::wait_for_any_of waiting on condvar without timeout"); - self.channel_notify.wait(&mut channels); + self.channel_notify.notified().await; } } } + + pub fn wait_for(&self, cond: &WaitCondition) -> Result { + run_loop_block_on(self.wait_for_async(cond)) + } + + pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result { + run_loop_block_on(self.wait_for_any_of_async(any_of)) + } + + fn try_take_channel_event( + &self, + channel: &str, + n: usize, + ) -> Result, String> { + let mut channels = self.channels.lock().expect("channels mutex poisoned"); + let queue = channels + .get_mut(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + if queue.len() < n { + return Ok(None); + } + if n == 1 { + if let Some(value) = queue.pop_front() { + return Ok(Some(WaitEvent::Channel { + channel: channel.to_string(), + value, + })); + } + return Ok(None); + } + let mut values = Vec::with_capacity(n); + for _ in 0..n { + if let Some(v) = queue.pop_front() { + values.push(v); + } + } + Ok(Some(WaitEvent::Channel { + channel: channel.to_string(), + value: serde_json::Value::Array(values), + })) + } } diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index c4da9cfdc..1760ff75a 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -1,6 +1,6 @@ use crate::engine::{ - merge_json_update, node_pool_execute, run_loop_pool_execute, run_scheduler_loop, - AnyOfCondition, Engine, NodeExecResult, SendPayload, + merge_json_update, node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition, + Engine, NodeExecResult, NodeOutcome, SendPayload, WaitEvent, WaitRequest, }; use serde::Deserialize; use serde_json::Value; @@ -8,6 +8,7 @@ use std::ffi::{CStr, CString}; use std::os::raw::c_char; use std::sync::mpsc; use std::sync::{Arc, Mutex}; +use tokio::sync::mpsc as tokio_mpsc; #[derive(Debug, Deserialize)] struct SendPayloadJson { @@ -29,6 +30,8 @@ struct CallbackEnvelopeIn { #[serde(default)] payload: Option, #[serde(default)] + suspend: Option, + #[serde(default)] error: Option, } @@ -63,7 +66,7 @@ fn into_c_ptr(s: String) -> *mut c_char { fn parse_c_callback_result( raw: String, node_name: &str, -) -> Result, String> { +) -> Result, String> { let parsed: CallbackEnvelopeIn = serde_json::from_str(&raw) .map_err(|e| format!("decode callback envelope for `{node_name}` failed: {e}"))?; if !parsed.ok { @@ -71,6 +74,9 @@ fn parse_c_callback_result( .error .unwrap_or_else(|| format!("callback reported error for `{node_name}`"))); } + if let Some(wait) = parsed.suspend { + return Ok(NodeOutcome::Suspended { wait }); + } let payload = parsed .payload .ok_or_else(|| format!("callback payload missing for `{node_name}`"))?; @@ -82,22 +88,40 @@ fn parse_c_callback_result( arg: s.arg, }) .collect(); - Ok(NodeExecResult { + Ok(NodeOutcome::Completed(NodeExecResult { update: payload.update, sends, - }) + })) +} + +enum SchedulerEventJson { + Node(Result), + Resume { + node: String, + arg: Value, + event: WaitEvent, + }, + WaitError(String), +} + +struct NodeExecutionJson { + node: String, + arg: Value, + outcome: NodeOutcome, } fn spawn_json_node_task( node: String, arg: Value, state_snapshot: Value, - tx: mpsc::Sender), String>>, + tx: tokio_mpsc::UnboundedSender, user_data_bits: libc::c_ulong, callback: CNodeCallback, ) -> Result<(), String> { node_pool_execute(move || { - let result = (|| -> Result<(String, NodeExecResult), String> { + let node_for_result = node.clone(); + let arg_for_result = arg.clone(); + let result = (|| -> Result { let node_c = CString::new(node.clone()).map_err(|e| format!("invalid node name: {e}"))?; let arg_json = serde_json::to_string(&arg) @@ -126,56 +150,126 @@ fn spawn_json_node_task( libc::free(out_ptr.cast()); } let payload = parse_c_callback_result(out_raw, &node)?; - Ok((node, payload)) + Ok(NodeExecutionJson { + node: node_for_result, + arg: arg_for_result, + outcome: payload, + }) })(); - let _ = tx.send(result); + let _ = tx.send(SchedulerEventJson::Node(result)); }) } -fn run_graph_scheduler_json( +async fn run_graph_scheduler_json( entry_point: String, finish_point: String, initial_state: Value, initial_input: Value, + engine: Engine, user_data: CUserData, callback: CNodeCallback, ) -> Result { - let (tx, rx) = mpsc::channel::), String>>(); + let (tx, mut rx) = tokio_mpsc::unbounded_channel::(); let state = Arc::new(Mutex::new(initial_state)); let user_data_bits = user_data.0; let tx_for_spawn = tx.clone(); let state_for_spawn = Arc::clone(&state); let state_for_merge = Arc::clone(&state); - let initial_arg = initial_input; - run_scheduler_loop( + let mut active: usize = 1; + let mut waiting: usize = 0; + spawn_json_node_task( entry_point, - &finish_point, - initial_arg, - move |node, arg| { - let snapshot = state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(); - spawn_json_node_task( - node, - arg, - snapshot, - tx_for_spawn.clone(), - user_data_bits, - callback, - ) - }, - move |_node_name, update| { - let mut guard = state_for_merge.lock().expect("state mutex poisoned"); - merge_json_update(&mut guard, update); - Ok(()) - }, - rx, + initial_input, + state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(), + tx_for_spawn.clone(), + user_data_bits, + callback, )?; + while active > 0 || waiting > 0 { + let evt = rx + .recv() + .await + .ok_or_else(|| "scheduler event channel closed".to_string())?; + match evt { + SchedulerEventJson::Node(result) => { + active = active.saturating_sub(1); + let exec = result?; + match exec.outcome { + NodeOutcome::Completed(node_result) => { + let mut guard = state_for_merge.lock().expect("state mutex poisoned"); + merge_json_update(&mut guard, node_result.update); + drop(guard); + if exec.node == finish_point { + break; + } + for send in node_result.sends { + active += 1; + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_json_node_task( + send.node, + send.arg, + snapshot, + tx_for_spawn.clone(), + user_data_bits, + callback, + )?; + } + } + NodeOutcome::Suspended { wait } => { + waiting += 1; + let tx_wait = tx_for_spawn.clone(); + let node = exec.node; + let arg = exec.arg; + let engine_for_wait = engine.clone(); + tokio::spawn(async move { + match engine_for_wait.wait_request_async(&wait).await { + Ok(event) => { + let _ = tx_wait.send(SchedulerEventJson::Resume { node, arg, event }); + } + Err(e) => { + let _ = tx_wait.send(SchedulerEventJson::WaitError(e)); + } + } + }); + } + } + } + SchedulerEventJson::Resume { node, arg, event } => { + waiting = waiting.saturating_sub(1); + active += 1; + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_json_node_task( + node, + wrap_resume_arg(arg, event), + snapshot, + tx_for_spawn.clone(), + user_data_bits, + callback, + )?; + } + SchedulerEventJson::WaitError(e) => return Err(e), + } + } let final_state = state.lock().expect("state mutex poisoned").clone(); Ok(final_state) } +fn wrap_resume_arg(arg: Value, event: WaitEvent) -> Value { + serde_json::json!({ + "__lg_resume_arg__": arg, + "__lg_resume_event__": event, + }) +} + #[no_mangle] pub extern "C" fn rc_engine_new() -> *mut Engine { Box::into_raw(Box::new(Engine::new())) @@ -278,7 +372,7 @@ pub unsafe extern "C" fn rc_wait_any_of_json( )) } }; - let result = (*ptr).wait_for_any_of(&any_of); + let result = run_loop_block_on((*ptr).wait_for_any_of_async(&any_of)); match result { Ok(event) => match serde_json::to_string(&event) { Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"event\":{s}}}")), @@ -343,15 +437,18 @@ pub unsafe extern "C" fn rc_run_graph_json( let (tx, rx) = mpsc::channel::>(); let user_data = CUserData(user_data); - let submit = run_loop_pool_execute(move || { + let run_engine = (*ptr).clone(); + let submit = run_loop_spawn(async move { let out = run_graph_scheduler_json( entry_point, finish_point, initial_state, initial_input, + run_engine, user_data, callback, - ); + ) + .await; let _ = tx.send(out); }); if let Err(e) = submit { diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs index 384f96121..d260b271a 100644 --- a/rust-core/src/lib_py.rs +++ b/rust-core/src/lib_py.rs @@ -1,7 +1,7 @@ #[cfg(feature = "python-bindings")] use crate::engine::{ - node_pool_execute, run_loop_pool_execute, run_scheduler_loop, AnyOfCondition, Engine, - NodeExecResult, SendPayload, WaitCondition, + node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition, Engine, NodeExecResult, + NodeOutcome, SendPayload, WaitCondition, WaitEvent, WaitRequest, }; #[cfg(feature = "python-bindings")] use pyo3::exceptions::PyValueError; @@ -17,6 +17,8 @@ use serde_json::Value; use std::sync::mpsc; #[cfg(feature = "python-bindings")] use std::sync::Arc; +#[cfg(feature = "python-bindings")] +use tokio::sync::mpsc as tokio_mpsc; #[cfg(feature = "python-bindings")] #[pyclass] @@ -58,9 +60,7 @@ impl PyRustEngine { fn wait_any_of_json(&self, any_of_json: &str) -> PyResult { let any_of: AnyOfCondition = serde_json::from_str(any_of_json) .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; - let event = self - .inner - .wait_for_any_of(&any_of) + let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) .map_err(PyValueError::new_err)?; serde_json::to_string(&event) .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) @@ -71,7 +71,8 @@ impl PyRustEngine { channel: channel.to_string(), n, }; - let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + let event = + run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; let event_json = serde_json::to_string(&event) .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; json_string_to_py_obj(py, &event_json) @@ -79,7 +80,8 @@ impl PyRustEngine { fn wait_timer(&self, py: Python<'_>, seconds: f64) -> PyResult> { let cond = WaitCondition::Timer { seconds }; - let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + let event = + run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; let event_json = serde_json::to_string(&event) .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; json_string_to_py_obj(py, &event_json) @@ -89,9 +91,7 @@ impl PyRustEngine { let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?; let any_of: AnyOfCondition = serde_json::from_str(&payload_json) .map_err(|e| PyValueError::new_err(format!("Invalid any_of payload: {e}")))?; - let event = self - .inner - .wait_for_any_of(&any_of) + let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) .map_err(PyValueError::new_err)?; let event_json = serde_json::to_string(&event) .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; @@ -101,7 +101,8 @@ impl PyRustEngine { fn wait_condition_json(&self, cond_json: &str) -> PyResult { let cond: WaitCondition = serde_json::from_str(cond_json) .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?; - let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + let event = + run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; serde_json::to_string(&event) .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) } @@ -120,9 +121,11 @@ impl PyRustEngine { let finish_point = finish_point.to_string(); let (done_tx, done_rx) = mpsc::channel::>(); let state_for_run = Arc::clone(&state); - run_loop_pool_execute(move || { + let engine = self.inner.clone(); + run_loop_spawn(async move { let run_result = - run_graph_scheduler(entry_point, finish_point, callback, state_for_run); + run_graph_scheduler(entry_point, finish_point, callback, state_for_run, engine) + .await; let _ = done_tx.send(run_result); }) .map_err(PyValueError::new_err)?; @@ -135,76 +138,166 @@ impl PyRustEngine { } #[cfg(feature = "python-bindings")] -fn run_graph_scheduler( +enum SchedulerEventPy { + Node(Result), + Resume { + node: String, + arg: Py, + event: WaitEvent, + }, + WaitError(String), +} + +struct NodeExecutionPy { + node: String, + arg: Py, + outcome: NodeOutcome, Py>, +} + +#[cfg(feature = "python-bindings")] +async fn run_graph_scheduler( entry_point: String, finish_point: String, callback: Arc>, state: Arc>, + engine: Engine, ) -> Result<(), String> { - let (tx, rx) = - mpsc::channel::, Py>), String>>(); + let (tx, mut rx) = tokio_mpsc::unbounded_channel::(); let initial_arg = Python::with_gil(|py| (*state).clone_ref(py)); let callback_for_spawn = Arc::clone(&callback); let state_for_spawn = Arc::clone(&state); let tx_for_spawn = tx.clone(); let state_for_merge = Arc::clone(&state); - run_scheduler_loop( + let mut active: usize = 1; + let mut waiting: usize = 0; + spawn_node_task( entry_point, - &finish_point, initial_arg, - move |node, arg| { - spawn_node_task( - node, - arg, - tx_for_spawn.clone(), - Arc::clone(&callback_for_spawn), - Arc::clone(&state_for_spawn), - ) - }, - move |node_name, update| { - Python::with_gil(|py| -> Result<(), String> { - if let Some(update) = update { - apply_update_to_state(py, state_for_merge.as_ref(), &update) - .map_err(|e| format!("state merge failed for `{node_name}`: {e}"))?; + tx_for_spawn.clone(), + Arc::clone(&callback_for_spawn), + Arc::clone(&state_for_spawn), + )?; + + while active > 0 || waiting > 0 { + let event = rx + .recv() + .await + .ok_or_else(|| "scheduler event channel closed".to_string())?; + match event { + SchedulerEventPy::Node(result) => { + active = active.saturating_sub(1); + let exec = result?; + match exec.outcome { + NodeOutcome::Completed(node_result) => { + Python::with_gil(|py| -> Result<(), String> { + if let Some(update) = node_result.update { + apply_update_to_state(py, state_for_merge.as_ref(), &update) + .map_err(|e| { + format!("state merge failed for `{}`: {e}", exec.node) + })?; + } + Ok(()) + })?; + if exec.node == finish_point { + break; + } + for send in node_result.sends { + active += 1; + spawn_node_task( + send.node, + send.arg, + tx_for_spawn.clone(), + Arc::clone(&callback_for_spawn), + Arc::clone(&state_for_spawn), + )?; + } + } + NodeOutcome::Suspended { wait } => { + waiting += 1; + let tx_wait = tx_for_spawn.clone(); + let node = exec.node; + let arg = exec.arg; + let engine_for_wait = engine.clone(); + tokio::spawn(async move { + let outcome = engine_for_wait.wait_request_async(&wait).await; + match outcome { + Ok(event) => { + let _ = tx_wait.send(SchedulerEventPy::Resume { node, arg, event }); + } + Err(e) => { + let _ = tx_wait.send(SchedulerEventPy::WaitError(e)); + } + } + }); + } } - Ok(()) - }) - }, - rx, - ) + } + SchedulerEventPy::Resume { node, arg, event } => { + waiting = waiting.saturating_sub(1); + active += 1; + let resume_arg = wrap_resume_arg(&arg, &event)?; + spawn_node_task( + node, + resume_arg, + tx_for_spawn.clone(), + Arc::clone(&callback_for_spawn), + Arc::clone(&state_for_spawn), + )?; + } + SchedulerEventPy::WaitError(e) => return Err(e), + } + } + Ok(()) } #[cfg(feature = "python-bindings")] fn spawn_node_task( node: String, arg: Py, - tx: mpsc::Sender, Py>), String>>, + tx: tokio_mpsc::UnboundedSender, callback: Arc>, state_for_task: Arc>, ) -> Result<(), String> { node_pool_execute(move || { + let node_for_result = node.clone(); + let arg_for_result = Python::with_gil(|py| arg.clone_ref(py)); let outcome = Python::with_gil( - |py| -> Result<(String, NodeExecResult, Py>), String> { + |py| -> Result { let callback_bound = callback.as_ref().bind(py); let payload_obj = callback_bound .call1((node.as_str(), arg, (*state_for_task).clone_ref(py))) .map_err(|e| format!("callback failed for node `{node}`: {e}"))?; - let payload = parse_node_exec_result(&payload_obj) + let payload = parse_node_outcome(py, &payload_obj) .map_err(|e| format!("invalid callback payload for `{node}`: {e}"))?; - Ok((node, payload)) + Ok(NodeExecutionPy { + node: node_for_result, + arg: arg_for_result, + outcome: payload, + }) }, ); - let _ = tx.send(outcome); + let _ = tx.send(SchedulerEventPy::Node(outcome)); }) } #[cfg(feature = "python-bindings")] -fn parse_node_exec_result( +fn parse_node_outcome( + py: Python<'_>, payload_obj: &Bound<'_, PyAny>, -) -> Result, Py>, String> { +) -> Result, Py>, String> { let payload_dict = payload_obj .downcast::() .map_err(|_| "payload must be a dict".to_string())?; + let suspended_item = payload_dict + .get_item("suspend") + .map_err(|e| format!("failed to read suspend: {e}"))?; + if let Some(wait_obj) = suspended_item { + let wait_json = py_obj_to_json_string(py, &wait_obj) + .map_err(|e| format!("failed to encode suspend payload: {e}"))?; + let wait: WaitRequest = + serde_json::from_str(&wait_json).map_err(|e| format!("invalid suspend payload: {e}"))?; + return Ok(NodeOutcome::Suspended { wait }); + } let update_item = payload_dict .get_item("update") @@ -242,7 +335,25 @@ fn parse_node_exec_result( sends.push(SendPayload { node, arg }); } - Ok(NodeExecResult { update, sends }) + Ok(NodeOutcome::Completed(NodeExecResult { update, sends })) +} + +#[cfg(feature = "python-bindings")] +fn wrap_resume_arg(arg: &Py, event: &WaitEvent) -> Result, String> { + Python::with_gil(|py| -> Result, String> { + let wrapper = PyDict::new(py); + wrapper + .set_item("__lg_resume_arg__", arg.clone_ref(py)) + .map_err(|e| format!("failed to set resume arg: {e}"))?; + let event_json = + serde_json::to_string(event).map_err(|e| format!("failed to encode wait event: {e}"))?; + let event_obj = + json_string_to_py_obj(py, &event_json).map_err(|e| format!("failed to parse event: {e}"))?; + wrapper + .set_item("__lg_resume_event__", event_obj.bind(py)) + .map_err(|e| format!("failed to set resume event: {e}"))?; + Ok(wrapper.unbind().into_any()) + }) } #[cfg(feature = "python-bindings")]