Appearance
Custom Stream Processing
Datum's lowest-level graph extension point is the GraphStage / GraphStageLogic API. Built-in graph stages such as Buffer, TakeWhile, and several junctions use this protocol when a stage needs direct control over push, pull, completion, and failure.
GraphStage shape and spec
A custom stage implements GraphStage:
type Shapedeclares the stage ports, such asFlowShape<In, Out>for one inlet and one outlet.allocate_shapecreates typedInlet<T>andOutlet<T>ports.stage_specreturns aStageSpec; useStageSpec::opaque(name, inlets, outlets)for a stage that runs customGraphStageLogic.create_logicinstalls the runtime handlers.
For a one-in/one-out processing stage, the shape is FlowShape<In, Out>, exposed through GraphFlowShape when you return the shape from GraphDsl.
GraphStageLogic
GraphStageLogic stores per-materialization state for ports and handlers. Install handlers with set_handler(&inlet, Box<dyn InHandler>), set_out_handler(&outlet, Box<dyn OutHandler>), and set_timer_handler(Box<dyn TimerHandler>) when the stage owns timers.
InHandler callbacks:
on_pushruns when an inlet has an available element.on_upstream_finishruns when upstream completes.on_upstream_failureruns when upstream fails.
OutHandler callbacks:
on_pullruns when downstream requests an element.on_downstream_finishruns when downstream cancels.keep_handlercontrols whether a handler remains installed after callback dispatch.
Port and stage operations
Use port operations to model backpressure explicitly:
| Operation | Effect |
|---|---|
pull(&inlet) | Ask upstream for the next element |
grab(&inlet) | Take the element made available by on_push |
push(&outlet, item) | Emit one element when the outlet has demand |
offer(&inlet, item) | Supply an element to an inlet that has been pulled |
complete(&outlet) | Complete one outlet |
fail(&outlet, error) | Fail one outlet |
complete_stage() | Close all ports |
fail_stage(error) | Fail the whole stage |
Helpers reduce handler boilerplate:
emitandemit_multiplepush now if there is demand, otherwise install temporary output handlers that emit later.readandread_npull and collect input before running a callback.get_async_callbackreturns anAsyncCallback;invoke_without_logicqueues work for the fused executor to drain on the stage logic thread.schedule_once,schedule_periodically, andschedule_periodically_with_initial_delayinstall keyed stage timers.cancel_timercancels by key, andis_timer_activereports the active state.set_log_callbackandlog(level, message)provide stage-local logging hooks.
Worked stage pattern
The implementation pattern below mirrors Datum's in-crate one-in/one-out stages.
rust
use datum::{
AnyInlet, AnyOutlet, FlowShape, GraphStage, GraphStageLogic, InHandler, Inlet, OutHandler,
Outlet, PortAllocator, StageSpec, StreamResult,
};
use datum::graph::Shape;
struct AddOne;
impl GraphStage for AddOne {
type Shape = FlowShape<u64, u64>;
fn name(&self) -> &str {
"AddOne"
}
fn allocate_shape(&self, allocator: &mut PortAllocator) -> Self::Shape {
FlowShape::new(
allocator.inlet("add-one.in"),
allocator.outlet("add-one.out"),
)
}
fn stage_spec(&self, shape: &Self::Shape) -> StageSpec {
StageSpec::opaque(self.name(), shape.inlets(), shape.outlets())
}
fn create_logic(&self, shape: &Self::Shape) -> GraphStageLogic {
struct In {
inlet: Inlet<u64>,
outlet: Outlet<u64>,
}
impl InHandler for In {
fn on_push(
&mut self,
logic: &mut GraphStageLogic,
_inlet: AnyInlet,
) -> StreamResult<()> {
let value = logic.grab(&self.inlet)?;
logic.emit(&self.outlet, value + 1)?;
if !logic.has_been_pulled(&self.inlet) {
logic.pull(&self.inlet)?;
}
Ok(())
}
fn on_upstream_finish(
&mut self,
logic: &mut GraphStageLogic,
_inlet: AnyInlet,
) -> StreamResult<()> {
logic.complete(&self.outlet)
}
}
struct Out {
inlet: Inlet<u64>,
}
impl OutHandler for Out {
fn on_pull(
&mut self,
logic: &mut GraphStageLogic,
_outlet: AnyOutlet,
) -> StreamResult<()> {
if !logic.has_been_pulled(&self.inlet) {
logic.pull(&self.inlet)?;
}
Ok(())
}
}
let mut logic = GraphStageLogic::new(shape);
logic.set_log_callback(|_level, _message| {});
logic.log("debug", "AddOne initialized");
logic.pull(&shape.inlet()).unwrap();
logic
.set_handler(
&shape.inlet(),
Box::new(In {
inlet: shape.inlet(),
outlet: shape.outlet(),
}),
)
.unwrap();
logic
.set_out_handler(
&shape.outlet(),
Box::new(Out {
inlet: shape.inlet(),
}),
)
.unwrap();
logic
}
}Stage timers
GraphStageLogic timers are owned by a single materialized stage. Runtime timer callbacks enqueue a stage-timer event; the fused executor delivers that event back to TimerHandler::on_timer on the stage logic thread, serialized with ordinary push/pull handlers. Timers are keyed by string, can be replaced or canceled by key, and are canceled automatically when the stage completes or fails.
rust
use datum::graph::Shape;
use datum::{
AnyOutlet, FlowShape, GraphDsl, GraphFlowShape, GraphStage, GraphStageLogic, OutHandler,
Outlet, PortAllocator, StageSpec, StreamResult, TimerHandler,
};
use std::time::Duration;
struct DelayedOne;
impl GraphStage for DelayedOne {
type Shape = FlowShape<(), u64>;
fn name(&self) -> &str {
"DelayedOne"
}
fn allocate_shape(&self, allocator: &mut PortAllocator) -> Self::Shape {
FlowShape::new(
allocator.inlet("delayed-one.in"),
allocator.outlet("delayed-one.out"),
)
}
fn stage_spec(&self, shape: &Self::Shape) -> StageSpec {
StageSpec::opaque(self.name(), shape.inlets(), shape.outlets())
}
fn create_logic(&self, shape: &Self::Shape) -> GraphStageLogic {
struct Out {
armed: bool,
}
impl OutHandler for Out {
fn on_pull(
&mut self,
logic: &mut GraphStageLogic,
_outlet: AnyOutlet,
) -> StreamResult<()> {
if !self.armed {
self.armed = true;
logic.schedule_once("emit", Duration::from_millis(1))?;
}
Ok(())
}
}
struct Timer {
outlet: Outlet<u64>,
}
impl TimerHandler for Timer {
fn on_timer(&mut self, logic: &mut GraphStageLogic, key: &str) -> StreamResult<()> {
assert_eq!(key, "emit");
assert!(!logic.is_timer_active("emit"));
logic.push(&self.outlet, 1)?;
logic.complete(&self.outlet)
}
}
let mut logic = GraphStageLogic::new(shape);
logic
.set_out_handler(&shape.outlet(), Box::new(Out { armed: false }))
.unwrap();
logic.set_timer_handler(Box::new(Timer {
outlet: shape.outlet(),
}));
logic
}
}
let graph = GraphDsl::try_create(|builder| {
let delayed = builder.add(DelayedOne);
Ok(GraphFlowShape::new(delayed.inlet(), delayed.outlet()))
})
.unwrap();
let values = graph.run_with_input([()]).unwrap();Monitoring and termination
For observability around ordinary flows, use the linear helpers before reaching for a custom stage. Flow::monitor observes each emitted element through a closure. Flow::watch_termination runs a callback when the flow materializes and preserves the original materialized value.
rust
use datum::{Flow, NotUsed, Sink, Source};
use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
};
let observed = Arc::new(Mutex::new(Vec::new()));
let materialized = Arc::new(AtomicBool::new(false));
let flow = Flow::<u64, u64, NotUsed>::identity()
.monitor({
let observed = Arc::clone(&observed);
move |item| observed.lock().unwrap().push(*item)
})
.watch_termination({
let materialized = Arc::clone(&materialized);
move |_mat| {
materialized.store(true, Ordering::SeqCst);
}
});
let values: Vec<u64> = Source::from_iter(1_u64..=3)
.via(flow)
.run_with(Sink::collect())
.unwrap()
.wait()
.unwrap();Next steps
- Working with Graphs —
GraphDsl, shapes, and junctions - Configuration — runtime timers and async-boundary configuration
- Operator Reference — graph-stage API index