mod activation;
mod checkpoint_data;
pub mod checkpointer;
pub mod command;
pub mod compiled;
pub mod config;
pub mod graph;
pub mod matrix_hook;
pub mod pending_writes;
pub mod phase_store;
mod pregel;
pub mod registry;
pub mod retry;
pub mod snapshot;
pub use checkpoint_data::CheckpointData;
pub use checkpointer::{CheckpointMeta, Checkpointer, InMemoryCheckpointer, PendingWrite};
pub use command::Command;
pub use compiled::{CompiledGraph, ExecutionOutcome};
pub use config::GraphConfig;
pub use graph::StateGraph;
pub use matrix_hook::{
ConvergenceRecorder, DefaultMatrixHook, MatrixHook, MatrixHookHandle, RoutingResolver,
};
pub use pending_writes::PendingWrites;
pub use phase_store::{PhaseStateStore, PhaseStoreError};
pub use registry::GraphRegistry;
pub use retry::{RetryPolicy, with_retry};
pub use snapshot::StateSnapshot;
pub use pe_core::types::{END, START};
#[cfg(test)]
#[allow(dead_code)]
pub(crate) mod tests {
use pe_core::node::{NodeContext, NodeFn, NodeFuture, NodeResult};
use pe_core::state::{State, StateUpdate};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TestState {
pub messages: Vec<String>,
pub counter: u32,
pub thread_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TestUpdate {
pub messages: Option<Vec<String>>,
pub counter: Option<u32>,
}
impl StateUpdate for TestUpdate {}
impl State for TestState {
type Update = TestUpdate;
fn apply(&mut self, update: TestUpdate) {
if let Some(msgs) = update.messages {
self.messages.extend(msgs);
}
if let Some(c) = update.counter {
self.counter = c;
}
}
}
impl TestState {
pub fn new() -> Self {
Self {
messages: Vec::new(),
counter: 0,
thread_id: "test-thread".into(),
}
}
}
impl TestUpdate {
pub fn with_message(msg: impl Into<String>) -> Self {
Self {
messages: Some(vec![msg.into()]),
counter: None,
}
}
pub fn with_counter(n: u32) -> Self {
Self {
messages: None,
counter: Some(n),
}
}
}
pub struct AppendNode {
node_name: &'static str,
message: &'static str,
}
impl AppendNode {
pub fn new(name: &'static str, message: &'static str) -> Self {
Self {
node_name: name,
message,
}
}
}
impl NodeFn<TestState> for AppendNode {
fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
let msg = self.message.to_string();
Box::pin(async move { NodeResult::Update(TestUpdate::with_message(msg)) })
}
fn name(&self) -> &str {
self.node_name
}
}
}