use std::time::Duration;
use super::{
error::BspError,
graph::{EdgeId, NodeId},
run::PendingSuspension,
};
pub trait BspInterceptor: Send + Sync {
fn on_superstep_start(&self, _step: u64, _frontier: &[NodeId]) {}
fn on_node_start(&self, _node_id: NodeId, _name: &str) {}
fn on_node_complete(&self, _node_id: NodeId, _name: &str, _duration: Duration) {}
fn on_node_error(&self, _node_id: NodeId, _name: &str, _err: &BspError) {}
fn on_node_retry(&self, _node_id: NodeId, _attempt: u32, _wait: Duration, _err: &BspError) {}
fn on_edge_taken(&self, _source: NodeId, _target: NodeId, _condition_label: &str) {}
fn on_loop_iteration(&self, _edge: EdgeId, _iteration: u64, _max: Option<u64>) {}
fn on_suspended(&self, _suspensions: &[PendingSuspension]) {}
fn on_complete(&self, _steps: u64) {}
}
#[derive(Debug, Clone)]
pub enum BspEvent {
SuperstepStart { step: u64, frontier: Vec<NodeId> },
NodeStart { node_id: NodeId, name: String },
NodeComplete {
node_id: NodeId,
name: String,
duration: Duration,
},
NodeError {
node_id: NodeId,
name: String,
error: BspError,
},
NodeRetry {
node_id: NodeId,
attempt: u32,
wait: Duration,
error: BspError,
},
EdgeTaken {
source: NodeId,
target: NodeId,
condition_label: String,
},
LoopIteration {
edge: EdgeId,
iteration: u64,
max: Option<u64>,
},
Suspended { suspensions: Vec<PendingSuspension> },
Complete { steps: u64 },
}
#[cfg(feature = "tokio")]
pub struct ChannelInterceptor {
tx: tokio::sync::broadcast::Sender<BspEvent>,
}
#[cfg(feature = "tokio")]
impl ChannelInterceptor {
#[must_use]
pub fn new(capacity: usize) -> Self {
let (tx, _) = tokio::sync::broadcast::channel(capacity);
Self { tx }
}
#[must_use]
pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<BspEvent> {
self.tx.subscribe()
}
fn send(&self, event: BspEvent) {
let _ = self.tx.send(event);
}
}
#[cfg(feature = "tokio")]
impl BspInterceptor for ChannelInterceptor {
fn on_superstep_start(&self, step: u64, frontier: &[NodeId]) {
self.send(BspEvent::SuperstepStart {
step,
frontier: frontier.to_vec(),
});
}
fn on_node_start(&self, node_id: NodeId, name: &str) {
self.send(BspEvent::NodeStart {
node_id,
name: name.to_string(),
});
}
fn on_node_complete(&self, node_id: NodeId, name: &str, duration: Duration) {
self.send(BspEvent::NodeComplete {
node_id,
name: name.to_string(),
duration,
});
}
fn on_node_error(&self, node_id: NodeId, name: &str, err: &BspError) {
self.send(BspEvent::NodeError {
node_id,
name: name.to_string(),
error: err.clone(),
});
}
fn on_node_retry(&self, node_id: NodeId, attempt: u32, wait: Duration, err: &BspError) {
self.send(BspEvent::NodeRetry {
node_id,
attempt,
wait,
error: err.clone(),
});
}
fn on_edge_taken(&self, source: NodeId, target: NodeId, condition_label: &str) {
self.send(BspEvent::EdgeTaken {
source,
target,
condition_label: condition_label.to_string(),
});
}
fn on_loop_iteration(&self, edge: EdgeId, iteration: u64, max: Option<u64>) {
self.send(BspEvent::LoopIteration {
edge,
iteration,
max,
});
}
fn on_suspended(&self, suspensions: &[PendingSuspension]) {
self.send(BspEvent::Suspended {
suspensions: suspensions.to_vec(),
});
}
fn on_complete(&self, steps: u64) {
self.send(BspEvent::Complete { steps });
}
}
#[cfg(test)]
mod tests {
use super::*;
struct LegacyInterceptor;
impl BspInterceptor for LegacyInterceptor {
fn on_node_start(&self, _node_id: NodeId, _name: &str) {}
fn on_node_complete(&self, _node_id: NodeId, _name: &str, _duration: Duration) {}
fn on_node_error(&self, _node_id: NodeId, _name: &str, _err: &BspError) {}
}
#[test]
fn test_default_methods_are_noops() {
let interceptor = LegacyInterceptor;
interceptor.on_superstep_start(0, &[1, 2]);
interceptor.on_node_retry(1, 1, Duration::ZERO, &BspError::Deadlock);
interceptor.on_edge_taken(1, 2, "always");
interceptor.on_loop_iteration(0, 1, Some(30));
interceptor.on_suspended(&[]);
interceptor.on_complete(3);
}
#[cfg(feature = "tokio")]
#[test]
fn test_channel_interceptor_broadcasts() {
let interceptor = ChannelInterceptor::new(16);
let mut rx = interceptor.subscribe();
interceptor.on_edge_taken(1, 2, "always");
interceptor.on_complete(7);
let first = rx.try_recv().unwrap();
let second = rx.try_recv().unwrap();
assert!(matches!(
first,
BspEvent::EdgeTaken {
source: 1,
target: 2,
..
}
));
assert!(matches!(second, BspEvent::Complete { steps: 7 }));
}
}