WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 00:22:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

tokio-rewrite

Browse Source
This commit is contained in:
Quanzheng Long
2026-03-13 23:16:40 -07:00
parent c2a9661d2a
commit 1b3d075dbb
10 changed files with 626 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph-go/advancedgraph/graph.go
+31 -3
View File
@@ -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
}
langgraph-go/advancedgraph/rust_engine.go
+14
View File
@@ -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:
langgraph-go/advancedgraph/types.go
+20 -1
View File
@@ -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)
libs/langgraph/langgraph/advanced_graph/state.py
+57 -12
View File
@@ -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 == "<lambda>":
libs/langgraph/tests/advanced-graph/test_run_pool_size.py
+55
View File
@@ -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
rust-core/Cargo.lock
Generated
+28
View File
@@ -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"
rust-core/Cargo.toml
+1
View File
@@ -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"] }
rust-core/src/engine.rs
+129 -85
View File
@@ -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<A> {
pub node: String,
pub arg: A,
@@ -46,6 +57,11 @@ pub struct NodeExecResult<U, A> {
pub sends: Vec<SendPayload<A>>,
}
pub enum NodeOutcome<U, A> {
Completed(NodeExecResult<U, A>),
Suspended { wait: WaitRequest },
}
type Task = Box<dyn FnOnce() + Send + 'static>;
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::<usize>().ok());
parsed.unwrap_or(default).max(min)
}
struct ThreadPool {
tx: mpsc::Sender<Task>,
_workers: Vec<thread::JoinHandle<()>>,
@@ -122,9 +145,10 @@ where
debug_log("node_pool_execute() called");
static NODE_POOL: OnceLock<ThreadPool> = 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<ThreadPool> = 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<F>(future: F) -> Result<(), String>
where
F: Future<Output = ()> + Send + 'static,
{
run_runtime().spawn(future);
Ok(())
}
pub fn run_loop_block_on<F>(future: F) -> F::Output
where
F: Future,
{
run_runtime().block_on(future)
}
fn run_runtime() -> &'static Runtime {
static RUNTIME: OnceLock<Runtime> = 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<U: Send + 'static, A: Send + 'static, FSpawn, FMerge>(
@@ -227,8 +281,8 @@ pub fn merge_json_update(state: &mut Value, update: Option<Value>) {
#[derive(Clone, Default)]
pub struct Engine {
channels: Arc<Mutex<HashMap<String, VecDeque<serde_json::Value>>>>,
channel_notify: Arc<Condvar>,
channels: Arc<StdMutex<HashMap<String, VecDeque<serde_json::Value>>>>,
channel_notify: Arc<Notify>,
}
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<WaitEvent, String> {
debug_log(&format!("Engine::wait_for(cond={cond:?})"));
pub async fn wait_request_async(&self, wait: &WaitRequest) -> Result<WaitEvent, String> {
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<WaitEvent, String> {
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<WaitEvent, String> {
pub async fn wait_for_any_of_async(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
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<WaitEvent, String> {
run_loop_block_on(self.wait_for_async(cond))
}
pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
run_loop_block_on(self.wait_for_any_of_async(any_of))
}
fn try_take_channel_event(
&self,
channel: &str,
n: usize,
) -> Result<Option<WaitEvent>, 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),
}))
}
}
rust-core/src/lib_c.rs
+135 -38
View File
@@ -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<NodeExecResultJsonWire>,
#[serde(default)]
suspend: Option<WaitRequest>,
#[serde(default)]
error: Option<String>,
}
@@ -63,7 +66,7 @@ fn into_c_ptr(s: String) -> *mut c_char {
fn parse_c_callback_result(
raw: String,
node_name: &str,
) -> Result<NodeExecResult<Value, Value>, String> {
) -> Result<NodeOutcome<Value, Value>, 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<NodeExecutionJson, String>),
Resume {
node: String,
arg: Value,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionJson {
node: String,
arg: Value,
outcome: NodeOutcome<Value, Value>,
}
fn spawn_json_node_task(
node: String,
arg: Value,
state_snapshot: Value,
tx: mpsc::Sender<Result<(String, NodeExecResult<Value, Value>), String>>,
tx: tokio_mpsc::UnboundedSender<SchedulerEventJson>,
user_data_bits: libc::c_ulong,
callback: CNodeCallback,
) -> Result<(), String> {
node_pool_execute(move || {
let result = (|| -> Result<(String, NodeExecResult<Value, Value>), String> {
let node_for_result = node.clone();
let arg_for_result = arg.clone();
let result = (|| -> Result<NodeExecutionJson, String> {
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<Value, String> {
let (tx, rx) = mpsc::channel::<Result<(String, NodeExecResult<Value, Value>), String>>();
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventJson>();
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::<Result<Value, String>>();
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 {
rust-core/src/lib_py.rs
+156 -45
View File
@@ -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<String> {
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<Py<PyAny>> {
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<String> {
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::<Result<(), String>>();
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<NodeExecutionPy, String>),
Resume {
node: String,
arg: Py<PyAny>,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionPy {
node: String,
arg: Py<PyAny>,
outcome: NodeOutcome<Py<PyAny>, Py<PyAny>>,
}
#[cfg(feature = "python-bindings")]
async fn run_graph_scheduler(
entry_point: String,
finish_point: String,
callback: Arc<Py<PyAny>>,
state: Arc<Py<PyAny>>,
engine: Engine,
) -> Result<(), String> {
let (tx, rx) =
mpsc::channel::<Result<(String, NodeExecResult<Py<PyAny>, Py<PyAny>>), String>>();
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventPy>();
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<PyAny>,
tx: mpsc::Sender<Result<(String, NodeExecResult<Py<PyAny>, Py<PyAny>>), String>>,
tx: tokio_mpsc::UnboundedSender<SchedulerEventPy>,
callback: Arc<Py<PyAny>>,
state_for_task: Arc<Py<PyAny>>,
) -> 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<PyAny>, Py<PyAny>>), String> {
|py| -> Result<NodeExecutionPy, String> {
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<NodeExecResult<Py<PyAny>, Py<PyAny>>, String> {
) -> Result<NodeOutcome<Py<PyAny>, Py<PyAny>>, String> {
let payload_dict = payload_obj
.downcast::<PyDict>()
.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<PyAny>, event: &WaitEvent) -> Result<Py<PyAny>, String> {
Python::with_gil(|py| -> Result<Py<PyAny>, 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")]