feat(s2n-quic-dc): update MTU on dc path when MTU is updated (#2327)

* feat(s2n-quic-dc): update MTU on dc path when MTU is updated

* remove std

* add tests and update map

* typo

* add test

* add clone test

* use fetch_min

* add entry to the handshakingpath

dc/s2n-quic-dc/src/path/secret/map.rs

@@ -298,7 +298,7 @@ impl Map {
  ) -> Option<(seal::Once, Credentials, ApplicationParams)> {
  let state = self.state.peers.get_by_key(&peer)?;
  let (sealer, credentials) = state.uni_sealer();
- Some((sealer, credentials, state.parameters))
+ Some((sealer, credentials, state.parameters.clone()))
  }
 
  pub fn open_once(
@@ -319,7 +319,7 @@ impl Map {
  let state = self.state.peers.get_by_key(&peer)?;
  let keys = state.bidi_local(features);
 
- Some((keys, state.parameters))
+ Some((keys, state.parameters.clone()))
  }
 
  pub fn pair_for_credentials(
@@ -330,7 +330,7 @@ impl Map {
  ) -> Option<(Bidirectional, ApplicationParams)> {
  let state = self.pre_authentication(credentials, control_out)?;
 
- let params = state.parameters;
+ let params = state.parameters.clone();
  let keys = state.bidi_remote(self.clone(), credentials, features);
 
  Some((keys, params))
@@ -684,13 +684,13 @@ impl Entry {
  secret: schedule::Secret,
  sender: sender::State,
  receiver: receiver::State,
- mut parameters: ApplicationParams,
+ parameters: ApplicationParams,
  rehandshake_time: Duration,
  ) -> Self {
  // clamp max datagram size to a well-known value
- parameters.max_datagram_size = parameters
+ parameters
  .max_datagram_size
- .min(crate::stream::MAX_DATAGRAM_SIZE as _);
+ .fetch_min(crate::stream::MAX_DATAGRAM_SIZE as _, Ordering::Relaxed);
 
  assert!(rehandshake_time.as_secs() <= u32::MAX as u64);
  Self {
@@ -911,6 +911,7 @@ pub struct HandshakingPath {
  parameters: ApplicationParams,
  endpoint_type: s2n_quic_core::endpoint::Type,
  secret: Option<schedule::Secret>,
+ entry: Option<Arc<Entry>>,
  map: Map,
 }
 
@@ -924,9 +925,10 @@ impl HandshakingPath {
  Self {
  peer: connection_info.remote_address.clone().into(),
  dc_version: connection_info.dc_version,
- parameters: connection_info.application_params,
+ parameters: connection_info.application_params.clone(),
  endpoint_type,
  secret: None,
+ entry: None,
  map,
  }
  }
@@ -1018,12 +1020,22 @@ impl dc::Path for HandshakingPath {
  .expect("peer tokens are only received after secrets are ready"),
  sender,
  receiver,
- self.parameters,
+ self.parameters.clone(),
  self.map.state.rehandshake_period,
  );
  let entry = Arc::new(entry);
+ self.entry = Some(entry.clone());
  self.map.insert(entry);
  }
+
+ fn on_mtu_updated(&mut self, mtu: u16) {
+ if let Some(entry) = self.entry.as_ref() {
+ entry
+ .parameters
+ .max_datagram_size
+ .store(mtu, Ordering::Relaxed);
+ }
+ }
 }
 
 #[cfg(test)]

dc/s2n-quic-dc/src/stream/endpoint.rs

@@ -19,7 +19,10 @@ use s2n_quic_core::{
  inet::{ExplicitCongestionNotification, SocketAddress},
  varint::VarInt,
 };
-use std::{io, sync::Arc};
+use std::{
+ io,
+ sync::{atomic::Ordering, Arc},
+};
 use tracing::{debug_span, Instrument as _};
 
 type Result<T = (), E = io::Error> = core::result::Result<T, E>;
@@ -193,7 +196,7 @@ where
  let flow = flow::non_blocking::State::new(flow_offset);
 
  let path = send::path::Info {
- max_datagram_size: parameters.max_datagram_size,
+ max_datagram_size: parameters.max_datagram_size.load(Ordering::Relaxed),
  send_quantum,
  ecn: ExplicitCongestionNotification::Ect0,
  next_expected_control_packet: VarInt::ZERO,

dc/s2n-quic-dc/src/stream/send/state.rs

@@ -42,7 +42,10 @@ use s2n_quic_core::{
  varint::VarInt,
 };
 use slotmap::SlotMap;
-use std::collections::{BinaryHeap, VecDeque};
+use std::{
+ collections::{BinaryHeap, VecDeque},
+ sync::atomic::Ordering,
+};
 use tracing::{debug, trace};
 
 pub mod probe;
@@ -118,7 +121,7 @@ pub struct PeerActivity {
 impl State {
  #[inline]
  pub fn new(stream_id: stream::Id, params: &ApplicationParams) -> Self {
- let max_datagram_size = params.max_datagram_size;
+ let max_datagram_size = params.max_datagram_size.load(Ordering::Relaxed);
  let initial_max_data = params.remote_max_data;
  let local_max_data = params.local_send_max_data;
 

quic/s2n-quic-core/src/dc.rs

@@ -11,7 +11,11 @@ use crate::{
  transport::parameters::{DcSupportedVersions, InitialFlowControlLimits},
  varint::VarInt,
 };
-use core::{num::NonZeroU32, time::Duration};
+use core::{
+ num::NonZeroU32,
+ sync::atomic::{AtomicU16, Ordering},
+ time::Duration,
+};
 
 mod disabled;
 mod traits;
@@ -91,25 +95,37 @@ impl<'a> DatagramInfo<'a> {
 }
 
 /// Various settings relevant to the dc path
-#[derive(Clone, Copy, Debug)]
+#[derive(Debug)]
 #[non_exhaustive]
 pub struct ApplicationParams {
- pub max_datagram_size: u16,
+ pub max_datagram_size: AtomicU16,
  pub remote_max_data: VarInt,
  pub local_send_max_data: VarInt,
  pub local_recv_max_data: VarInt,
  // Actually a Duration, stored as milliseconds to shrink this struct
  pub max_idle_timeout: Option<NonZeroU32>,
 }
 
+impl Clone for ApplicationParams {
+ fn clone(&self) -> Self {
+ Self {
+ max_datagram_size: AtomicU16::new(self.max_datagram_size.load(Ordering::Relaxed)),
+ remote_max_data: self.remote_max_data,
+ local_send_max_data: self.local_send_max_data,
+ local_recv_max_data: self.local_recv_max_data,
+ max_idle_timeout: self.max_idle_timeout,
+ }
+ }
+}
+
 impl ApplicationParams {
  pub fn new(
  max_datagram_size: u16,
  peer_flow_control_limits: &InitialFlowControlLimits,
  limits: &Limits,
  ) -> Self {
  Self {
- max_datagram_size,
+ max_datagram_size: AtomicU16::new(max_datagram_size),
  remote_max_data: peer_flow_control_limits.max_data,
  local_send_max_data: limits.initial_stream_limits().max_data_bidi_local,
  local_recv_max_data: limits.initial_stream_limits().max_data_bidi_remote,
@@ -125,3 +141,52 @@ impl ApplicationParams {
  Some(Duration::from_millis(self.max_idle_timeout?.get() as u64))
  }
 }
+
+#[cfg(test)]
+mod tests {
+ use crate::{
+ connection::Limits, dc::ApplicationParams, transport::parameters::InitialFlowControlLimits,
+ varint::VarInt,
+ };
+ use std::{sync::atomic::Ordering, time::Duration};
+
+ #[test]
+ fn clone() {
+ let initial_flow_control_limits = InitialFlowControlLimits {
+ max_data: VarInt::from_u32(2222),
+ ..Default::default()
+ };
+
+ let limits = Limits {
+ bidirectional_local_data_window: 1234.try_into().unwrap(),
+ bidirectional_remote_data_window: 6789.try_into().unwrap(),
+ max_idle_timeout: Duration::from_millis(999).try_into().unwrap(),
+ ..Default::default()
+ };
+
+ let params = ApplicationParams::new(9000, &initial_flow_control_limits, &limits);
+
+ assert_eq!(9000, params.max_datagram_size.load(Ordering::Relaxed));
+ assert_eq!(limits.max_idle_timeout(), params.max_idle_timeout());
+ assert_eq!(1234, params.local_send_max_data.as_u64());
+ assert_eq!(6789, params.local_recv_max_data.as_u64());
+ assert_eq!(2222, params.remote_max_data.as_u64());
+
+ let cloned_params = params.clone();
+
+ assert_eq!(
+ params.max_datagram_size.load(Ordering::Relaxed),
+ cloned_params.max_datagram_size.load(Ordering::Relaxed)
+ );
+ assert_eq!(params.max_idle_timeout, cloned_params.max_idle_timeout);
+ assert_eq!(
+ params.local_send_max_data,
+ cloned_params.local_send_max_data
+ );
+ assert_eq!(
+ params.local_recv_max_data,
+ cloned_params.local_recv_max_data
+ );
+ assert_eq!(params.remote_max_data, cloned_params.remote_max_data);
+ }
+}

quic/s2n-quic-core/src/dc/disabled.rs

@@ -44,4 +44,8 @@ impl Path for () {
  ) {
  unimplemented!()
  }
+
+ fn on_mtu_updated(&mut self, _mtu: u16) {
+ unimplemented!()
+ }
 }

quic/s2n-quic-core/src/dc/testing.rs

@@ -10,7 +10,7 @@ use crate::{
 };
 use core::{num::NonZeroU32, time::Duration};
 use std::sync::{
- atomic::{AtomicU8, Ordering},
+ atomic::{AtomicU16, AtomicU8, Ordering},
  Arc,
 };
 
@@ -36,14 +36,19 @@ pub struct MockDcPath {
  pub on_peer_stateless_reset_tokens_count: u8,
  pub stateless_reset_tokens: Vec<stateless_reset::Token>,
  pub peer_stateless_reset_tokens: Vec<stateless_reset::Token>,
+ pub mtu: u16,
 }
 
 impl dc::Endpoint for MockDcEndpoint {
  type Path = MockDcPath;
 
- fn new_path(&mut self, _connection_info: &ConnectionInfo) -> Option<Self::Path> {
+ fn new_path(&mut self, connection_info: &ConnectionInfo) -> Option<Self::Path> {
  Some(MockDcPath {
  stateless_reset_tokens: self.stateless_reset_tokens.clone(),
+ mtu: connection_info
+ .application_params
+ .max_datagram_size
+ .load(Ordering::Relaxed),
  ..Default::default()
  })
  }
@@ -76,10 +81,15 @@ impl dc::Path for MockDcPath {
  self.peer_stateless_reset_tokens
  .extend(stateless_reset_tokens);
  }
+
+ fn on_mtu_updated(&mut self, mtu: u16) {
+ self.mtu = mtu
+ }
 }
 
+#[allow(clippy::declare_interior_mutable_const)]
 pub const TEST_APPLICATION_PARAMS: ApplicationParams = ApplicationParams {
- max_datagram_size: 1472,
+ max_datagram_size: AtomicU16::new(1472),
  remote_max_data: VarInt::from_u32(1u32 << 25),
  local_send_max_data: VarInt::from_u32(1u32 << 25),
  local_recv_max_data: VarInt::from_u32(1u32 << 25),

quic/s2n-quic-core/src/dc/traits.rs

@@ -45,6 +45,9 @@ pub trait Path: 'static + Send {
  &mut self,
  stateless_reset_tokens: impl Iterator<Item = &'a stateless_reset::Token>,
  );
+
+ /// Called when the MTU has been updated for the path
+ fn on_mtu_updated(&mut self, mtu: u16);
 }
 
 impl<P: Path> Path for Option<P> {
@@ -69,4 +72,11 @@ impl<P: Path> Path for Option<P> {
  path.on_peer_stateless_reset_tokens(stateless_reset_tokens)
  }
  }
+
+ #[inline]
+ fn on_mtu_updated(&mut self, max_datagram_size: u16) {
+ if let Some(path) = self {
+ path.on_mtu_updated(max_datagram_size)
+ }
+ }
 }