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