use alloc::{collections::TryReserveError, vec::Vec};
use crate::{FlowObservation, FlowObserver};
const _: () = assert!(core::mem::size_of::<usize>() <= core::mem::size_of::<u64>());
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FlowAttributionFamily {
CircuitTurn,
PipeSend,
ParticipantStep,
SwitchboardRoute,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FlowAttributionPhase {
Began,
Completed,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct FlowAttribution {
family: FlowAttributionFamily,
phase: FlowAttributionPhase,
subject: &'static str,
message: &'static str,
transport: Option<&'static str>,
observations: u64,
representation_bytes: u64,
saturated: bool,
}
impl FlowAttribution {
pub const fn family(self) -> FlowAttributionFamily {
self.family
}
pub const fn phase(self) -> FlowAttributionPhase {
self.phase
}
pub const fn subject(self) -> &'static str {
self.subject
}
pub const fn message(self) -> &'static str {
self.message
}
pub const fn transport(self) -> Option<&'static str> {
self.transport
}
pub const fn observations(self) -> u64 {
self.observations
}
pub const fn representation_bytes(self) -> u64 {
self.representation_bytes
}
pub const fn saturated(self) -> bool {
self.saturated
}
fn record(&mut self, representation_bytes: usize) {
accumulate(&mut self.observations, 1, &mut self.saturated);
accumulate(
&mut self.representation_bytes,
representation_bytes as u64,
&mut self.saturated,
);
}
}
#[derive(Debug)]
pub enum AttributingFlowObserverCreationError {
StorageUnavailable { source: TryReserveError },
}
pub struct AttributingFlowObserver {
maximum_attributions: usize,
attributions: Vec<FlowAttribution>,
full: bool,
}
impl AttributingFlowObserver {
pub fn try_new(
maximum_attributions: usize,
) -> Result<Self, AttributingFlowObserverCreationError> {
let mut attributions = Vec::new();
attributions
.try_reserve_exact(maximum_attributions)
.map_err(
|source| AttributingFlowObserverCreationError::StorageUnavailable { source },
)?;
Ok(Self {
maximum_attributions,
attributions,
full: false,
})
}
pub fn attributions(&self) -> &[FlowAttribution] {
&self.attributions
}
pub const fn full(&self) -> bool {
self.full
}
pub fn reset(&mut self) {
self.attributions.clear();
self.full = false;
}
fn record(&mut self, observation: FlowObservation) {
let (family, phase, subject, message, transport, representation_bytes) =
attribution_identity(observation);
if let Some(attribution) = self.attributions.iter_mut().find(|attribution| {
attribution.family == family
&& attribution.phase == phase
&& attribution.subject == subject
&& attribution.message == message
&& attribution.transport == transport
}) {
attribution.record(representation_bytes);
return;
}
if self.attributions.len() == self.maximum_attributions {
self.full = true;
return;
}
self.attributions.push(FlowAttribution {
family,
phase,
subject,
message,
transport,
observations: 1,
representation_bytes: representation_bytes as u64,
saturated: false,
});
}
}
impl FlowObserver for AttributingFlowObserver {
fn observe(&mut self, observation: FlowObservation) {
self.record(observation);
}
}
fn attribution_identity(
observation: FlowObservation,
) -> (
FlowAttributionFamily,
FlowAttributionPhase,
&'static str,
&'static str,
Option<&'static str>,
usize,
) {
match observation {
FlowObservation::CircuitTurnBegan {
circuit,
input,
input_representation_bytes,
} => (
FlowAttributionFamily::CircuitTurn,
FlowAttributionPhase::Began,
circuit,
input,
None,
input_representation_bytes,
),
FlowObservation::CircuitTurnCompleted {
circuit,
output,
output_representation_bytes,
} => (
FlowAttributionFamily::CircuitTurn,
FlowAttributionPhase::Completed,
circuit,
output,
None,
output_representation_bytes,
),
FlowObservation::PipeSendBegan {
participant,
input,
transport,
input_representation_bytes,
} => (
FlowAttributionFamily::PipeSend,
FlowAttributionPhase::Began,
participant,
input,
Some(transport),
input_representation_bytes,
),
FlowObservation::PipeSendCompleted {
participant,
transport,
outcome,
output_representation_bytes,
} => (
FlowAttributionFamily::PipeSend,
FlowAttributionPhase::Completed,
participant,
outcome,
Some(transport),
output_representation_bytes,
),
FlowObservation::ParticipantStepBegan {
participant,
input,
input_representation_bytes,
} => (
FlowAttributionFamily::ParticipantStep,
FlowAttributionPhase::Began,
participant,
input,
None,
input_representation_bytes,
),
FlowObservation::ParticipantStepCompleted {
participant,
outcome,
output_representation_bytes,
} => (
FlowAttributionFamily::ParticipantStep,
FlowAttributionPhase::Completed,
participant,
outcome,
None,
output_representation_bytes,
),
FlowObservation::SwitchboardRouteBegan {
switchboard,
input,
input_representation_bytes,
} => (
FlowAttributionFamily::SwitchboardRoute,
FlowAttributionPhase::Began,
switchboard,
input,
None,
input_representation_bytes,
),
FlowObservation::SwitchboardRouteCompleted {
switchboard,
output,
output_representation_bytes,
} => (
FlowAttributionFamily::SwitchboardRoute,
FlowAttributionPhase::Completed,
switchboard,
output,
None,
output_representation_bytes,
),
}
}
fn accumulate(target: &mut u64, increment: u64, saturated: &mut bool) {
match target.checked_add(increment) {
Some(total) => *target = total,
None => {
*target = u64::MAX;
*saturated = true;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn every_observation_retains_its_complete_identity() -> Result<(), &'static str> {
let mut observer = AttributingFlowObserver::try_new(8).map_err(|_| "creation failed")?;
for observation in [
FlowObservation::CircuitTurnBegan {
circuit: "circuit",
input: "circuit-in",
input_representation_bytes: 1,
},
FlowObservation::CircuitTurnCompleted {
circuit: "circuit",
output: "circuit-out",
output_representation_bytes: 2,
},
FlowObservation::PipeSendBegan {
participant: "pipe",
input: "pipe-in",
transport: "transport",
input_representation_bytes: 3,
},
FlowObservation::PipeSendCompleted {
participant: "pipe",
transport: "transport",
outcome: "pipe-out",
output_representation_bytes: 4,
},
FlowObservation::ParticipantStepBegan {
participant: "participant",
input: "participant-in",
input_representation_bytes: 5,
},
FlowObservation::ParticipantStepCompleted {
participant: "participant",
outcome: "participant-out",
output_representation_bytes: 6,
},
FlowObservation::SwitchboardRouteBegan {
switchboard: "switchboard",
input: "switchboard-in",
input_representation_bytes: 7,
},
FlowObservation::SwitchboardRouteCompleted {
switchboard: "switchboard",
output: "switchboard-out",
output_representation_bytes: 8,
},
] {
observer.observe(observation);
}
let [
circuit_began,
circuit_completed,
pipe_began,
pipe_completed,
participant_began,
participant_completed,
switchboard_began,
switchboard_completed,
] = observer.attributions()
else {
return Err("unexpected attribution count");
};
assert_eq!(
(circuit_began.family(), circuit_began.phase()),
(
FlowAttributionFamily::CircuitTurn,
FlowAttributionPhase::Began
)
);
assert_eq!(circuit_completed.message(), "circuit-out");
assert_eq!(pipe_began.transport(), Some("transport"));
assert_eq!(pipe_completed.representation_bytes(), 4);
assert_eq!(participant_began.subject(), "participant");
assert_eq!(participant_completed.observations(), 1);
assert_eq!(switchboard_began.message(), "switchboard-in");
assert_eq!(switchboard_completed.representation_bytes(), 8);
Ok(())
}
#[test]
fn observations_are_aggregated_by_their_complete_identity() -> Result<(), &'static str> {
let mut observer = AttributingFlowObserver::try_new(2).map_err(|_| "creation failed")?;
for observation in [
FlowObservation::ParticipantStepBegan {
participant: "connection",
input: "receive",
input_representation_bytes: 8,
},
FlowObservation::ParticipantStepBegan {
participant: "connection",
input: "receive",
input_representation_bytes: 8,
},
FlowObservation::ParticipantStepCompleted {
participant: "connection",
outcome: "decoded",
output_representation_bytes: 16,
},
] {
observer.observe(observation);
}
let [began, completed] = observer.attributions() else {
return Err("unexpected attribution count");
};
assert_eq!(began.observations(), 2);
assert_eq!(began.representation_bytes(), 16);
assert_eq!(completed.observations(), 1);
assert_eq!(completed.representation_bytes(), 16);
assert!(!observer.full());
Ok(())
}
#[test]
fn capacity_exhaustion_is_explicit_and_reset_preserves_capacity() -> Result<(), &'static str> {
let mut observer = AttributingFlowObserver::try_new(1).map_err(|_| "creation failed")?;
observer.observe(FlowObservation::CircuitTurnBegan {
circuit: "first",
input: "input",
input_representation_bytes: 1,
});
observer.observe(FlowObservation::CircuitTurnBegan {
circuit: "second",
input: "input",
input_representation_bytes: 1,
});
assert!(observer.full());
assert_eq!(observer.attributions().len(), 1);
observer.reset();
assert!(!observer.full());
assert!(observer.attributions().is_empty());
Ok(())
}
#[test]
fn overflowing_counts_saturate_explicitly() {
let mut attribution = FlowAttribution {
family: FlowAttributionFamily::ParticipantStep,
phase: FlowAttributionPhase::Began,
subject: "participant",
message: "input",
transport: None,
observations: u64::MAX,
representation_bytes: u64::MAX,
saturated: false,
};
attribution.record(1);
assert_eq!(attribution.observations(), u64::MAX);
assert_eq!(attribution.representation_bytes(), u64::MAX);
assert!(attribution.saturated());
}
}