1mod activation;
18mod checkpoint_data;
19pub mod checkpointer;
20pub mod command;
21pub mod compiled;
22pub mod config;
23pub mod graph;
24pub mod matrix_hook;
25pub mod pending_writes;
26pub mod phase_store;
27mod pregel;
28pub mod registry;
29pub mod retry;
30pub mod snapshot;
31
32pub use checkpoint_data::CheckpointData;
34pub use checkpointer::{CheckpointMeta, Checkpointer, InMemoryCheckpointer, PendingWrite};
35pub use command::Command;
36pub use compiled::{CompiledGraph, ExecutionOutcome};
37pub use config::GraphConfig;
38pub use graph::StateGraph;
39pub use matrix_hook::{
40 ConvergenceRecorder, DefaultMatrixHook, MatrixHook, MatrixHookHandle, RoutingResolver,
41};
42pub use pending_writes::PendingWrites;
43pub use phase_store::{PhaseStateStore, PhaseStoreError};
44pub use registry::GraphRegistry;
45pub use retry::{RetryPolicy, with_retry};
46pub use snapshot::StateSnapshot;
47
48pub use pe_core::types::{END, START};
50
51#[cfg(test)]
55#[allow(dead_code)]
56pub(crate) mod tests {
57 use pe_core::node::{NodeContext, NodeFn, NodeFuture, NodeResult};
58 use pe_core::state::{State, StateUpdate};
59 use serde::{Deserialize, Serialize};
60
61 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
64 pub struct TestState {
65 pub messages: Vec<String>,
66 pub counter: u32,
67 pub thread_id: String,
68 }
69
70 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
71 pub struct TestUpdate {
72 pub messages: Option<Vec<String>>,
73 pub counter: Option<u32>,
74 }
75
76 impl StateUpdate for TestUpdate {}
77
78 impl State for TestState {
79 type Update = TestUpdate;
80
81 fn apply(&mut self, update: TestUpdate) {
82 if let Some(msgs) = update.messages {
84 self.messages.extend(msgs);
85 }
86 if let Some(c) = update.counter {
88 self.counter = c;
89 }
90 }
91 }
92
93 impl TestState {
94 pub fn new() -> Self {
95 Self {
96 messages: Vec::new(),
97 counter: 0,
98 thread_id: "test-thread".into(),
99 }
100 }
101 }
102
103 impl TestUpdate {
104 pub fn with_message(msg: impl Into<String>) -> Self {
105 Self {
106 messages: Some(vec![msg.into()]),
107 counter: None,
108 }
109 }
110
111 pub fn with_counter(n: u32) -> Self {
112 Self {
113 messages: None,
114 counter: Some(n),
115 }
116 }
117 }
118
119 pub struct AppendNode {
123 node_name: &'static str,
124 message: &'static str,
125 }
126
127 impl AppendNode {
128 pub fn new(name: &'static str, message: &'static str) -> Self {
129 Self {
130 node_name: name,
131 message,
132 }
133 }
134 }
135
136 impl NodeFn<TestState> for AppendNode {
137 fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
138 let msg = self.message.to_string();
139 Box::pin(async move { NodeResult::Update(TestUpdate::with_message(msg)) })
140 }
141
142 fn name(&self) -> &str {
143 self.node_name
144 }
145 }
146}