Lochnair-Labs/irtt-rs
2
0 Fork 0
Code Issues 1 Pull Requests 4 Actions 8 Packages Projects Releases 6 Wiki Activity

refactor(session): consolidate client session lifecycle

Browse Source
This commit is contained in:
Nils Andreas Svee committed 2026-07-31 11:10:23 +07:00
1 parent dc827908d9
commit 4b4adc518c
1 file changed
+106 -108
Download Patch File Download Diff File Expand all files Collapse all files
crates/irtt-client/src/session/machine.rs
+106 -108
View File
@@ -29,18 +29,18 @@ pub(crate) struct SessionMachine {
config: ClientConfig, config: ClientConfig,
remote: std::net::SocketAddr, remote: std::net::SocketAddr,
requested: Params, requested: Params,
negotiated: Option<NegotiatedParams>, state: MachineState,
phase: ClientPhase,
session: Option<ActiveSession>,
final_packets_sent: u64,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug)]
enum ClientPhase { enum MachineState {
Connected, Connected,
Open { token: u64 }, Open(Box<ActiveSession>),
NoTestCompleted, NoTestCompleted,
Closed { source: CloseSource }, Closed {
source: CloseSource,
packets_sent: u64,
},
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -51,6 +51,8 @@ enum CloseSource {
#[derive(Debug)] #[derive(Debug)]
struct ActiveSession { struct ActiveSession {
token: u64,
negotiated: NegotiatedParams,
next_wire_seq: u32, next_wire_seq: u32,
highest_received_seq: Option<u32>, highest_received_seq: Option<u32>,
packets_sent: u64, packets_sent: u64,
@@ -91,10 +93,7 @@ impl SessionMachine {
config, config,
remote, remote,
requested, requested,
negotiated: None, state: MachineState::Connected,
phase: ClientPhase::Connected,
session: None,
final_packets_sent: 0,
}) })
} }
@@ -107,11 +106,11 @@ impl SessionMachine {
} }
pub(crate) fn open_packet(&self) -> Result<Vec<u8>, ClientError> { pub(crate) fn open_packet(&self) -> Result<Vec<u8>, ClientError> {
match self.phase { match self.state {
ClientPhase::Connected => {} MachineState::Connected => {}
ClientPhase::Open { .. } => return Err(ClientError::AlreadyOpen), MachineState::Open(_) => return Err(ClientError::AlreadyOpen),
ClientPhase::Closed { .. } => return Err(ClientError::AlreadyClosed), MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed),
ClientPhase::NoTestCompleted => return Err(ClientError::AlreadyCompleted), MachineState::NoTestCompleted => return Err(ClientError::AlreadyCompleted),
} }
let request = OpenRequest { let request = OpenRequest {
@@ -161,7 +160,9 @@ impl SessionMachine {
} }
pub(crate) fn next_send_deadline(&self) -> Option<Instant> { pub(crate) fn next_send_deadline(&self) -> Option<Instant> {
let session = self.session.as_ref()?; let MachineState::Open(session) = &self.state else {
return None;
};
if session.sending_done { if session.sending_done {
return None; return None;
} }
@@ -205,18 +206,16 @@ impl SessionMachine {
where where
F: FnOnce(&[u8]) -> Result<SendProbeResult, ClientError>, F: FnOnce(&[u8]) -> Result<SendProbeResult, ClientError>,
{ {
let token = match self.phase { let (config, state) = (&self.config, &mut self.state);
ClientPhase::Open { token } => token, let probe_timeout = config.probe_timeout;
ClientPhase::Closed { .. } => return Err(ClientError::AlreadyClosed), let hmac_key = config.hmac_key.as_deref();
ClientPhase::Connected => return Err(ClientError::NotOpen), let session = match state {
ClientPhase::NoTestCompleted => return Err(ClientError::AlreadyCompleted), MachineState::Open(session) => session,
MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed),
MachineState::Connected => return Err(ClientError::NotOpen),
MachineState::NoTestCompleted => return Err(ClientError::AlreadyCompleted),
}; };
let session = self
.session
.as_mut()
.expect("session must exist when phase is Open");
if session.sending_done { if session.sending_done {
return Ok(vec![]); return Ok(vec![]);
} }
@@ -232,35 +231,25 @@ impl SessionMachine {
session.pending.check_capacity()?; session.pending.check_capacity()?;
let negotiated = self
.negotiated
.as_ref()
.expect("negotiated must exist when Open");
let wire_seq = session.next_wire_seq; let wire_seq = session.next_wire_seq;
let scheduled_at = scheduled_at_override.unwrap_or(session.next_send_at); let scheduled_at = scheduled_at_override.unwrap_or(session.next_send_at);
let request = EchoRequest { let request = EchoRequest {
token, token: session.token,
sequence: wire_seq, sequence: wire_seq,
params: negotiated.params.clone(), params: session.negotiated.params.clone(),
payload: vec![], payload: vec![],
}; };
let packet = encode_echo_request(&request, self.config.hmac_key.as_deref())?; let packet = encode_echo_request(&request, hmac_key)?;
let send_result = send(&packet)?; let send_result = send(&packet)?;
let session = self
.session
.as_mut()
.expect("session must exist when phase is Open");
let pending = PendingProbe { let pending = PendingProbe {
wire_seq, wire_seq,
sent_at: send_result.sent_at, sent_at: send_result.sent_at,
timeout_at: send_result timeout_at: send_result
.sent_at .sent_at
.mono .mono
.checked_add(self.config.probe_timeout) .checked_add(probe_timeout)
.ok_or(ClientError::DurationOverflow)?, .ok_or(ClientError::DurationOverflow)?,
}; };
session.pending.insert(pending)?; session.pending.insert(pending)?;
@@ -274,11 +263,7 @@ impl SessionMachine {
counter: "packets_sent", counter: "packets_sent",
})?; })?;
let negotiated = self let interval_ns = u64::try_from(session.negotiated.params.interval_ns)
.negotiated
.as_ref()
.expect("negotiated must exist when Open");
let interval_ns = u64::try_from(negotiated.params.interval_ns)
.expect("validated positive negotiated interval"); .expect("validated positive negotiated interval");
let interval = Duration::from_nanos(interval_ns); let interval = Duration::from_nanos(interval_ns);
let (scheduled_at, next_send_at) = if skip_missed_slots { let (scheduled_at, next_send_at) = if skip_missed_slots {
@@ -292,10 +277,6 @@ impl SessionMachine {
) )
}; };
let timer_error = instant_abs_diff(send_result.sent_at.mono, scheduled_at); let timer_error = instant_abs_diff(send_result.sent_at.mono, scheduled_at);
let session = self
.session
.as_mut()
.expect("session must exist when phase is Open");
session.next_send_at = next_send_at; session.next_send_at = next_send_at;
if let Some(end) = session.end_mono { if let Some(end) = session.end_mono {
@@ -321,12 +302,7 @@ impl SessionMachine {
now: ClientTimestamp, now: ClientTimestamp,
meta: ReceiveMeta, meta: ReceiveMeta,
) -> Result<Vec<ClientEvent>, ClientError> { ) -> Result<Vec<ClientEvent>, ClientError> {
match self.phase { self.open_session()?;
ClientPhase::Open { .. } => {}
ClientPhase::Closed { .. } => return Err(ClientError::AlreadyClosed),
ClientPhase::Connected => return Err(ClientError::NotOpen),
ClientPhase::NoTestCompleted => return Err(ClientError::AlreadyCompleted),
}
let Some(reply) = self.decode_received_packet(packet) else { let Some(reply) = self.decode_received_packet(packet) else {
return Ok(vec![ClientEvent::Warning { return Ok(vec![ClientEvent::Warning {
@@ -342,17 +318,7 @@ impl SessionMachine {
&mut self, &mut self,
now: Instant, now: Instant,
) -> Result<Vec<ClientEvent>, ClientError> { ) -> Result<Vec<ClientEvent>, ClientError> {
match self.phase { let session = self.open_session_mut()?;
ClientPhase::Open { .. } => {}
ClientPhase::Closed { .. } => return Err(ClientError::AlreadyClosed),
ClientPhase::Connected => return Err(ClientError::NotOpen),
ClientPhase::NoTestCompleted => return Err(ClientError::AlreadyCompleted),
}
let session = self
.session
.as_mut()
.expect("session must exist when phase is Open");
let expired = session.pending.drain_expired(now); let expired = session.pending.drain_expired(now);
let mut events = Vec::with_capacity(expired.len()); let mut events = Vec::with_capacity(expired.len());
@@ -372,10 +338,10 @@ impl SessionMachine {
where where
F: FnOnce(&[u8]) -> Result<(), ClientError>, F: FnOnce(&[u8]) -> Result<(), ClientError>,
{ {
let token = match self.phase { let token = match &self.state {
ClientPhase::Open { token } => token, MachineState::Open(session) => session.token,
ClientPhase::Closed { .. } => return Err(ClientError::AlreadyClosed), MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed),
ClientPhase::Connected | ClientPhase::NoTestCompleted => { MachineState::Connected | MachineState::NoTestCompleted => {
return Err(ClientError::NotOpen) return Err(ClientError::NotOpen)
} }
}; };
@@ -393,38 +359,40 @@ impl SessionMachine {
} }
pub(crate) fn is_run_complete(&self) -> bool { pub(crate) fn is_run_complete(&self) -> bool {
let Some(session) = self.session.as_ref() else { match &self.state {
return matches!( MachineState::Open(session) => session.sending_done && session.pending.len() == 0,
self.phase, MachineState::NoTestCompleted | MachineState::Closed { .. } => true,
ClientPhase::Closed { .. } | ClientPhase::NoTestCompleted MachineState::Connected => false,
); }
};
session.sending_done && session.pending.len() == 0
} }
pub(crate) fn is_peer_closed(&self) -> bool { pub(crate) fn is_peer_closed(&self) -> bool {
matches!( matches!(
self.phase, self.state,
ClientPhase::Closed { MachineState::Closed {
source: CloseSource::Peer source: CloseSource::Peer,
..
} }
) )
} }
pub(crate) fn has_timed_out_metadata(&self) -> bool { pub(crate) fn has_timed_out_metadata(&self) -> bool {
self.session matches!(
.as_ref() &self.state,
.is_some_and(|session| session.timed_out.len() > 0) MachineState::Open(session) if session.timed_out.len() > 0
)
} }
pub(crate) fn packets_sent(&self) -> u64 { pub(crate) fn packets_sent(&self) -> u64 {
self.session match &self.state {
.as_ref() MachineState::Open(session) => session.packets_sent,
.map_or(self.final_packets_sent, |session| session.packets_sent) MachineState::Closed { packets_sent, .. } => *packets_sent,
MachineState::Connected | MachineState::NoTestCompleted => 0,
}
} }
pub(crate) fn is_open(&self) -> bool { pub(crate) fn is_open(&self) -> bool {
matches!(self.phase, ClientPhase::Open { .. }) matches!(self.state, MachineState::Open(_))
} }
fn accept_normal_open<F>( fn accept_normal_open<F>(
@@ -443,8 +411,6 @@ impl SessionMachine {
self.config.negotiation_policy, self.config.negotiation_policy,
)?; )?;
before_normal_open(&negotiated)?; before_normal_open(&negotiated)?;
self.negotiated = Some(negotiated.clone());
self.phase = ClientPhase::Open { token };
let end_mono = if negotiated.params.duration_ns > 0 { let end_mono = if negotiated.params.duration_ns > 0 {
Some(negotiated_end_mono( Some(negotiated_end_mono(
@@ -455,7 +421,9 @@ impl SessionMachine {
None None
}; };
self.session = Some(ActiveSession { self.state = MachineState::Open(Box::new(ActiveSession {
token,
negotiated: negotiated.clone(),
next_wire_seq: 0, next_wire_seq: 0,
highest_received_seq: None, highest_received_seq: None,
packets_sent: 0, packets_sent: 0,
@@ -466,7 +434,7 @@ impl SessionMachine {
timed_out: TimedOutMap::new(self.config.max_pending_probes), timed_out: TimedOutMap::new(self.config.max_pending_probes),
completed: CompletedSet::new(self.config.max_pending_probes), completed: CompletedSet::new(self.config.max_pending_probes),
sending_done: false, sending_done: false,
}); }));
let event = ClientEvent::SessionStarted { let event = ClientEvent::SessionStarted {
remote: self.remote, remote: self.remote,
@@ -493,8 +461,7 @@ impl SessionMachine {
reply.params, reply.params,
self.config.negotiation_policy, self.config.negotiation_policy,
)?; )?;
self.negotiated = Some(negotiated.clone()); self.state = MachineState::NoTestCompleted;
self.phase = ClientPhase::NoTestCompleted;
let event = ClientEvent::NoTestCompleted { let event = ClientEvent::NoTestCompleted {
remote: self.remote, remote: self.remote,
negotiated: negotiated.clone(), negotiated: negotiated.clone(),
@@ -508,12 +475,16 @@ impl SessionMachine {
} }
fn decode_received_packet(&self, packet: &[u8]) -> Option<EchoReply> { fn decode_received_packet(&self, packet: &[u8]) -> Option<EchoReply> {
let negotiated = self let session = self
.negotiated .open_session()
.as_ref() .expect("decode_received_packet is only called for an open session");
.expect("negotiated must exist when Open");
decode_echo_reply(packet, &negotiated.params, self.config.hmac_key.as_deref()).ok() decode_echo_reply(
packet,
&session.negotiated.params,
self.config.hmac_key.as_deref(),
)
.ok()
} }
fn process_echo_reply( fn process_echo_reply(
@@ -523,10 +494,10 @@ impl SessionMachine {
now: ClientTimestamp, now: ClientTimestamp,
meta: ReceiveMeta, meta: ReceiveMeta,
) -> Result<Vec<ClientEvent>, ClientError> { ) -> Result<Vec<ClientEvent>, ClientError> {
let token = match self.phase { let token = self
ClientPhase::Open { token } => token, .open_session()
_ => unreachable!(), .expect("process_echo_reply is only called for an open session")
}; .token;
if reply.token != token { if reply.token != token {
return Ok(vec![ClientEvent::Warning { return Ok(vec![ClientEvent::Warning {
kind: WarningKind::WrongToken, kind: WarningKind::WrongToken,
@@ -541,7 +512,9 @@ impl SessionMachine {
let wire_seq = reply.sequence; let wire_seq = reply.sequence;
let should_close = flags::has(reply.flags, flags::FLAG_CLOSE); let should_close = flags::has(reply.flags, flags::FLAG_CLOSE);
let mut events = { let mut events = {
let session = self.session.as_mut().expect("session must exist when Open"); let session = self
.open_session_mut()
.expect("process_echo_reply is only called for an open session");
if let Some(pending) = session.pending.remove(wire_seq) { if let Some(pending) = session.pending.remove(wire_seq) {
let rtt = compute_rtt(&pending.sent_at, &now, &reply.timestamps); let rtt = compute_rtt(&pending.sent_at, &now, &reply.timestamps);
@@ -658,11 +631,36 @@ impl SessionMachine {
} }
fn transition_to_closed(&mut self, source: CloseSource) { fn transition_to_closed(&mut self, source: CloseSource) {
if let Some(mut session) = self.session.take() { let packets_sent = match &mut self.state {
MachineState::Open(session) => {
session.timed_out.clear(); session.timed_out.clear();
self.final_packets_sent = session.packets_sent; session.packets_sent
}
MachineState::Closed { packets_sent, .. } => *packets_sent,
MachineState::Connected | MachineState::NoTestCompleted => 0,
};
self.state = MachineState::Closed {
source,
packets_sent,
};
}
fn open_session(&self) -> Result<&ActiveSession, ClientError> {
match &self.state {
MachineState::Open(session) => Ok(session),
MachineState::Closed { .. } => Err(ClientError::AlreadyClosed),
MachineState::Connected => Err(ClientError::NotOpen),
MachineState::NoTestCompleted => Err(ClientError::AlreadyCompleted),
}
}
fn open_session_mut(&mut self) -> Result<&mut ActiveSession, ClientError> {
match &mut self.state {
MachineState::Open(session) => Ok(session),
MachineState::Closed { .. } => Err(ClientError::AlreadyClosed),
MachineState::Connected => Err(ClientError::NotOpen),
MachineState::NoTestCompleted => Err(ClientError::AlreadyCompleted),
} }
self.phase = ClientPhase::Closed { source };
} }
} }