use crate::carrier::{
AttemptGeneration, BoundedFrameQueue, Carrier, CarrierAction, CarrierCloseCause,
CarrierCounters, CarrierEvent, CarrierLostCause, CarrierMachine, CarrierMachineInput,
CarrierRefusal, CarrierSendOutcome, CarrierSendRefusal, CarrierState, ConnectionKey, EpisodeId,
FrameDecoder, OpaqueBytes, decode_envelope, step,
};
pub trait EnvelopeShape {
fn decode(envelope: &[u8]) -> Result<Vec<OpaqueBytes>, CarrierRefusal>;
}
#[derive(Debug)]
pub struct NativeShape;
impl EnvelopeShape for NativeShape {
fn decode(envelope: &[u8]) -> Result<Vec<OpaqueBytes>, CarrierRefusal> {
decode_envelope(envelope)?;
let mut decoder = FrameDecoder::new();
let split_sizes = [1, 2, 5, 13];
let mut offset = 0;
let mut split_index = 0;
let mut frames = Vec::new();
while offset < envelope.len() {
let end = (offset + split_sizes[split_index % split_sizes.len()]).min(envelope.len());
frames.extend(decoder.push(&envelope[offset..end])?);
offset = end;
split_index += 1;
}
decoder.finish()?;
Ok(frames)
}
}
#[derive(Debug)]
pub struct WasmShape;
impl EnvelopeShape for WasmShape {
fn decode(envelope: &[u8]) -> Result<Vec<OpaqueBytes>, CarrierRefusal> {
decode_envelope(envelope).map(|payload| vec![payload])
}
}
#[derive(Debug)]
pub struct MockAdapter<S> {
shape: std::marker::PhantomData<S>,
machine: CarrierMachine,
counters: CarrierCounters,
events: Vec<CarrierEvent>,
queue: Option<BoundedFrameQueue>,
fatal_refusal: Option<CarrierRefusal>,
}
pub type NativeShapedMock = MockAdapter<NativeShape>;
pub type WasmShapedMock = MockAdapter<WasmShape>;
impl<S: EnvelopeShape> MockAdapter<S> {
pub fn new(episode: EpisodeId) -> Self {
Self {
shape: std::marker::PhantomData,
machine: CarrierMachine::new(episode),
counters: CarrierCounters::default(),
events: Vec::new(),
queue: None,
fatal_refusal: None,
}
}
pub fn apply(&mut self, input: CarrierMachineInput) -> Vec<CarrierAction> {
let transition = step(self.machine, input);
self.machine = transition.machine;
self.record_actions(&transition.actions);
self.refresh_queue();
transition.actions
}
pub fn receive_envelope(
&mut self,
key: ConnectionKey,
envelope: &[u8],
) -> Result<(), CarrierRefusal> {
if let Some(refusal) = self.fatal_refusal {
return Err(refusal);
}
match S::decode(envelope) {
Ok(frames) => {
for bytes in frames {
self.apply(CarrierMachineInput::FrameReceived { key, bytes });
}
Ok(())
}
Err(refusal) => {
self.fatal_refusal = Some(refusal);
self.apply(CarrierMachineInput::Refused { key, refusal });
self.apply(CarrierMachineInput::EstablishedFate {
key,
cause: CarrierLostCause::MalformedFrame(refusal),
});
Err(refusal)
}
}
}
pub fn flush_one(&mut self) -> bool {
let Some(queue) = self.queue.as_mut() else {
return false;
};
if queue.front().is_none() {
return false;
}
let actions = queue.confirm_handoff(&self.counters);
self.record_actions(&actions);
true
}
pub fn current_attempt(&self) -> Option<AttemptGeneration> {
match self.machine.state() {
CarrierState::Connected {
episode: _,
generation: _,
key: _,
}
| CarrierState::Stopped {
episode: _,
cause: _,
} => None,
CarrierState::Parked {
episode: _,
cause: _,
attempt,
} => match attempt {
crate::carrier::AttemptSlot::Idle => None,
crate::carrier::AttemptSlot::InFlight { generation } => Some(generation),
},
}
}
pub fn state(&self) -> CarrierState {
self.machine.state()
}
pub fn counters(&self) -> &CarrierCounters {
&self.counters
}
pub fn events(&self) -> &[CarrierEvent] {
&self.events
}
pub fn clear_events(&mut self) {
self.events.clear();
}
fn record_actions(&mut self, actions: &[CarrierAction]) {
self.counters.record_actions(actions);
for action in actions {
match action {
CarrierAction::EmitEvent(event) => self.events.push(event.clone()),
CarrierAction::StartAttempt { generation: _ }
| CarrierAction::CancelAttempt { generation: _ }
| CarrierAction::CloseSocket { key: _, cause: _ } => {}
}
}
}
fn refresh_queue(&mut self) {
match self.machine.state() {
CarrierState::Connected {
episode: _,
generation: _,
key,
} => {
if self.queue.is_none() {
self.queue = Some(BoundedFrameQueue::new(key));
}
}
CarrierState::Parked {
episode: _,
cause: _,
attempt: _,
}
| CarrierState::Stopped {
episode: _,
cause: _,
} => self.queue = None,
}
}
}
impl<S: EnvelopeShape> Carrier for MockAdapter<S> {
fn send(&mut self, key: ConnectionKey, bytes: OpaqueBytes) -> CarrierSendOutcome {
let current = match self.machine.state() {
CarrierState::Connected {
episode: _,
generation: _,
key,
} => Some(key),
CarrierState::Parked {
episode: _,
cause: _,
attempt: _,
}
| CarrierState::Stopped {
episode: _,
cause: _,
} => None,
};
if current != Some(key) {
return CarrierSendOutcome::Refused(CarrierSendRefusal::ConnectionNotCurrent { key });
}
let Some(queue) = self.queue.as_mut() else {
return CarrierSendOutcome::Refused(CarrierSendRefusal::ConnectionNotCurrent { key });
};
let result = queue.enqueue(bytes);
let outcome = result.outcome;
self.record_actions(&result.actions);
if let CarrierSendOutcome::Refused(CarrierSendRefusal::QueueCapacityExceeded {
queued_frames,
queued_bytes,
}) = outcome
{
let refusal = CarrierRefusal::QueueCapacityExceeded {
queued_frames,
queued_bytes,
};
self.apply(CarrierMachineInput::EstablishedFate {
key,
cause: CarrierLostCause::LocalClose(CarrierCloseCause::Refused(refusal)),
});
}
outcome
}
fn close(&mut self, key: ConnectionKey, cause: CarrierCloseCause) {
self.apply(CarrierMachineInput::ExplicitClose {
key: Some(key),
cause,
});
}
}