Skip to content

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 Shape declares the stage ports, such as FlowShape<In, Out> for one inlet and one outlet.
  • allocate_shape creates typed Inlet<T> and Outlet<T> ports.
  • stage_spec returns a StageSpec; use StageSpec::opaque(name, inlets, outlets) for a stage that runs custom GraphStageLogic.
  • create_logic installs 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_push runs when an inlet has an available element.
  • on_upstream_finish runs when upstream completes.
  • on_upstream_failure runs when upstream fails.

OutHandler callbacks:

  • on_pull runs when downstream requests an element.
  • on_downstream_finish runs when downstream cancels.
  • keep_handler controls whether a handler remains installed after callback dispatch.

Port and stage operations

Use port operations to model backpressure explicitly:

OperationEffect
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:

  • emit and emit_multiple push now if there is demand, otherwise install temporary output handlers that emit later.
  • read and read_n pull and collect input before running a callback.
  • get_async_callback returns an AsyncCallback; invoke_without_logic queues work for the fused executor to drain on the stage logic thread.
  • schedule_once, schedule_periodically, and schedule_periodically_with_initial_delay install keyed stage timers. cancel_timer cancels by key, and is_timer_active reports the active state.
  • set_log_callback and log(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