monomorphize QUIC packet send slots (#4812)

Summary:
Pull Request resolved: #4812

Make `PacketIo` expose its borrowed send slot through a generic associated type, and make `Endpoint` retain its concrete packet I/O type. This removes the per-batch `Box<dyn PacketSendSlot>` allocation and virtual dispatch while preserving driver-owned slot lifetimes.

Reviewed By: shayne-fletcher

Differential Revision: D117573677

Author: mariusae

chrysalis-transport-quiche/src/driver.rs

@@ -150,6 +150,7 @@ use crate::command::EndpointCommands;
 use crate::identity::EndpointIdentity;
 use crate::identity::certificate_pid;
 use crate::io::PacketIo;
+use crate::io::PacketSendSlot;
 
 mod connection;
 mod network;
@@ -461,9 +462,9 @@ fn complete_unknown(submission: Submission, completions: &mut Vec<Completion>) {
 }
 
 /// Single-owner quiche endpoint state machine.
-pub struct Endpoint {
+pub struct Endpoint<I> {
     driver: DriverId,
-    io: Box<dyn PacketIo>,
+    io: I,
     network: Network,
     submission_queue: EndpointSubmissionReceiver<EndpointCommand>,
     completion_sender: CompletionSender,
@@ -481,9 +482,9 @@ pub struct Endpoint {
     next_statistics_update: Instant,
 }
 
-impl Endpoint {
+impl<I: PacketIo> Endpoint<I> {
     /// Constructs a client endpoint and its thread-safe application handles.
-    pub fn client<I: PacketIo + 'static>(
+    pub fn client(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -507,7 +508,7 @@ impl Endpoint {
     }
 
     /// Constructs a client endpoint with explicit admission limits.
-    pub fn client_with_limits<I: PacketIo + 'static>(
+    pub fn client_with_limits(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -532,7 +533,7 @@ impl Endpoint {
     }
 
     /// Constructs a server endpoint and its thread-safe application handles.
-    pub fn server<I: PacketIo + 'static>(
+    pub fn server(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -556,7 +557,7 @@ impl Endpoint {
     }
 
     /// Constructs a server endpoint with explicit admission limits.
-    pub fn server_with_limits<I: PacketIo + 'static>(
+    pub fn server_with_limits(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -581,7 +582,7 @@ impl Endpoint {
     }
 
     /// Constructs a duplex endpoint that can initiate and accept connections.
-    pub fn duplex<I: PacketIo + 'static>(
+    pub fn duplex(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -606,7 +607,7 @@ impl Endpoint {
     }
 
     /// Constructs a duplex endpoint with explicit admission limits.
-    pub fn duplex_with_limits<I: PacketIo + 'static>(
+    pub fn duplex_with_limits(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -633,7 +634,7 @@ impl Endpoint {
     }
 
     /// Constructs a duplex endpoint whose CIDs use `routing_pid` independently of its identity.
-    pub fn duplex_routed<I: PacketIo + 'static>(
+    pub fn duplex_routed(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -659,7 +660,7 @@ impl Endpoint {
     }
 
     /// Constructs a routed duplex endpoint with explicit admission limits.
-    pub fn duplex_routed_with_limits<I: PacketIo + 'static>(
+    pub fn duplex_routed_with_limits(
         driver: DriverId,
         io: I,
         identity: EndpointIdentity,
@@ -684,7 +685,7 @@ impl Endpoint {
         )
     }
 
-    fn new<I: PacketIo + 'static>(
+    fn new(
         driver: DriverId,
         io: I,
         _identity: EndpointIdentity,
@@ -717,7 +718,7 @@ impl Endpoint {
         (
             Self {
                 driver,
-                io: Box::new(io),
+                io,
                 network: Network::new(
                     driver,
                     routing_pid,
@@ -836,7 +837,7 @@ impl Endpoint {
         }
         self.flush_completions();
         self.network
-            .queue_packets(self.io.as_mut(), &mut self.new_completions)?;
+            .queue_packets(&mut self.io, &mut self.new_completions)?;
 
         let mut wait = self.network.next_timeout(maximum_wait);
         if let Some(deadline) = self.shutdown_deadline {
@@ -850,7 +851,7 @@ impl Endpoint {
         self.update_statistics();
         self.expire_shutdown();
         self.network
-            .queue_packets(self.io.as_mut(), &mut self.new_completions)?;
+            .queue_packets(&mut self.io, &mut self.new_completions)?;
         self.stage_completions();
         self.flush_completions();
         self.finish_shutdown_if_ready();

chrysalis-transport-quiche/src/driver/network.rs

@@ -8,6 +8,7 @@
 
 use super::connection::ConnectionState;
 use super::*;
+use crate::io::PacketSendSlot;
 
 pub(super) struct Network {
     driver: DriverId,
@@ -372,9 +373,9 @@ impl Network {
             .saturating_sub(established_servers);
     }
 
-    pub(super) fn queue_packets(
+    pub(super) fn queue_packets<I: PacketIo>(
         &mut self,
-        io: &mut dyn PacketIo,
+        io: &mut I,
         _completions: &mut Vec<Completion>,
     ) -> Result<(), Error> {
         let segment_size = io.segment_size();

chrysalis-transport-quiche/src/driver/tests.rs

@@ -99,8 +99,8 @@ fn drain(handle: &EndpointHandle, events: &mut Vec<Completion>) {
 }
 
 fn drive_until(
-    client: &mut Endpoint,
-    server: &mut Endpoint,
+    client: &mut Endpoint<UdpDriver>,
+    server: &mut Endpoint<UdpDriver>,
     client_handle: &EndpointHandle,
     server_handle: &EndpointHandle,
     client_events: &mut Vec<Completion>,
@@ -118,9 +118,9 @@ fn drive_until(
 }
 
 fn endpoint_pair() -> (
-    Endpoint,
+    Endpoint<UdpDriver>,
     EndpointHandle,
-    Endpoint,
+    Endpoint<UdpDriver>,
     EndpointHandle,
     std::net::SocketAddr,
     Pid,
@@ -131,9 +131,9 @@ fn endpoint_pair() -> (
 fn endpoint_pair_with_completion_capacity(
     completion_capacity: usize,
 ) -> (
-    Endpoint,
+    Endpoint<UdpDriver>,
     EndpointHandle,
-    Endpoint,
+    Endpoint<UdpDriver>,
     EndpointHandle,
     std::net::SocketAddr,
     Pid,

chrysalis-transport-quiche/src/io.rs

@@ -35,16 +35,16 @@ pub trait PacketSendSlot {
     fn buffer_mut(&mut self) -> &mut [u8];
 
     /// Queues the initialized prefix for transmission.
-    fn submit(
-        self: Box<Self>,
-        length: usize,
-        destination: SocketAddr,
-        send_at: Instant,
-    ) -> io::Result<()>;
+    fn submit(self, length: usize, destination: SocketAddr, send_at: Instant) -> io::Result<()>;
 }
 
 /// Runtime-neutral packet I/O consumed by one quiche endpoint thread.
 pub trait PacketIo: Send {
+    /// Driver-owned transmit allocation borrowed by [`Self::try_send_slot`].
+    type SendSlot<'a>: PacketSendSlot
+    where
+        Self: 'a;
+
     /// Returns whether an authenticated upstream boundary validated source addresses.
     fn peer_addresses_validated(&self) -> bool {
         false
@@ -63,7 +63,7 @@ pub trait PacketIo: Send {
     fn notifier(&self) -> Arc<dyn Notifier>;
 
     /// Borrows a free transmit allocation.
-    fn try_send_slot(&mut self) -> Option<Box<dyn PacketSendSlot + '_>>;
+    fn try_send_slot(&mut self) -> Option<Self::SendSlot<'_>>;
 
     /// Waits for packet I/O, pacing, or an application notification.
     fn poll(&mut self, timeout: Duration) -> io::Result<()>;
@@ -78,25 +78,19 @@ pub trait PacketIo: Send {
     ) -> Result<(), Error>;
 }
 
-struct UringSendSlot<'a>(SendSlot<'a>);
-
-impl PacketSendSlot for UringSendSlot<'_> {
+impl PacketSendSlot for SendSlot<'_> {
     fn buffer_mut(&mut self) -> &mut [u8] {
-        self.0.buffer_mut()
+        SendSlot::buffer_mut(self)
     }
 
-    fn submit(
-        self: Box<Self>,
-        length: usize,
-        destination: SocketAddr,
-        send_at: Instant,
-    ) -> io::Result<()> {
-        let Self(slot) = *self;
-        slot.submit(length, destination, send_at)
+    fn submit(self, length: usize, destination: SocketAddr, send_at: Instant) -> io::Result<()> {
+        SendSlot::submit(self, length, destination, send_at)
     }
 }
 
 impl PacketIo for UdpDriver {
+    type SendSlot<'a> = SendSlot<'a>;
+
     fn local_addr(&self) -> io::Result<SocketAddr> {
         UdpDriver::local_addr(self)
     }
@@ -113,9 +107,8 @@ impl PacketIo for UdpDriver {
         Arc::new(UdpDriver::notifier(self))
     }
 
-    fn try_send_slot(&mut self) -> Option<Box<dyn PacketSendSlot + '_>> {
+    fn try_send_slot(&mut self) -> Option<Self::SendSlot<'_>> {
         UdpDriver::try_send_slot(self)
-            .map(|slot| Box::new(UringSendSlot(slot)) as Box<dyn PacketSendSlot>)
     }
 
     fn poll(&mut self, timeout: Duration) -> io::Result<()> {

chrysalis-transport-tokio/src/lib.rs

@@ -685,9 +685,9 @@ impl Transport {
         )
     }
 
-    fn spawn(
+    fn spawn<I: PacketIo + 'static>(
         pid: Pid,
-        mut endpoint: Endpoint,
+        mut endpoint: Endpoint<I>,
         handle: EndpointHandle,
         notifier: Arc<TokioNotifier>,
         completion_capacity: NonZeroUsize,

chrysalis-transport/src/packet_io.rs

@@ -632,7 +632,7 @@ impl RoutedUdpPacketIo {
     }
 }
 
-struct RoutedUdpSendSlot<'a> {
+pub(crate) struct RoutedUdpSendSlot<'a> {
     slot: Option<SendSlot<'a>>,
     router: &'a Router,
     addresses: &'a CarrierAddressBook,
@@ -649,12 +649,7 @@ impl PacketSendSlot for RoutedUdpSendSlot<'_> {
             .buffer_mut()
     }
 
-    fn submit(
-        mut self: Box<Self>,
-        length: usize,
-        peer: SocketAddr,
-        send_at: Instant,
-    ) -> io::Result<()> {
+    fn submit(mut self, length: usize, peer: SocketAddr, send_at: Instant) -> io::Result<()> {
         let mut slot = self
             .slot
             .take()
@@ -719,6 +714,8 @@ impl PacketSendSlot for RoutedUdpSendSlot<'_> {
 }
 
 impl PacketIo for RoutedUdpPacketIo {
+    type SendSlot<'a> = RoutedUdpSendSlot<'a>;
+
     fn peer_addresses_validated(&self) -> bool {
         false
     }
@@ -739,21 +736,21 @@ impl PacketIo for RoutedUdpPacketIo {
         self.wake.clone()
     }
 
-    fn try_send_slot(&mut self) -> Option<Box<dyn PacketSendSlot + '_>> {
+    fn try_send_slot(&mut self) -> Option<Self::SendSlot<'_>> {
         let router = self.router.as_ref();
         let addresses = &self.addresses;
         let fallback = self.fallback.as_deref();
         let fallback_transmits = &mut self.fallback_transmits;
         let segment_size = self.segment_size;
         let slot = self.udp.try_send_slot()?;
-        Some(Box::new(RoutedUdpSendSlot {
+        Some(RoutedUdpSendSlot {
             slot: Some(slot),
             router,
             addresses,
             fallback,
             fallback_transmits,
             segment_size,
-        }))
+        })
     }
 
     fn poll(&mut self, timeout: Duration) -> io::Result<()> {
@@ -1106,7 +1103,7 @@ impl<T: DatagramSocket> CarrierPacketIo<T> {
     }
 }
 
-struct CarrierSendSlot<'a, T> {
+pub(crate) struct CarrierSendSlot<'a, T> {
     io: &'a mut CarrierPacketIo<T>,
     buffer: Option<Vec<u8>>,
 }
@@ -1119,7 +1116,7 @@ impl<T> PacketSendSlot for CarrierSendSlot<'_, T> {
     }
 
     fn submit(
-        mut self: Box<Self>,
+        mut self,
         length: usize,
         destination: SocketAddr,
         send_at: Instant,
@@ -1161,6 +1158,11 @@ impl<T> Drop for CarrierSendSlot<'_, T> {
 }
 
 impl<T: DatagramSocket> PacketIo for CarrierPacketIo<T> {
+    type SendSlot<'a>
+        = CarrierSendSlot<'a, T>
+    where
+        T: 'a;
+
     fn local_addr(&self) -> io::Result<SocketAddr> {
         Ok(self.local)
     }
@@ -1177,12 +1179,12 @@ impl<T: DatagramSocket> PacketIo for CarrierPacketIo<T> {
         self.wake.clone()
     }
 
-    fn try_send_slot(&mut self) -> Option<Box<dyn PacketSendSlot + '_>> {
+    fn try_send_slot(&mut self) -> Option<Self::SendSlot<'_>> {
         let buffer = self.free_transmit.pop()?;
-        Some(Box::new(CarrierSendSlot {
+        Some(CarrierSendSlot {
             io: self,
             buffer: Some(buffer),
-        }))
+        })
     }
 
     fn poll(&mut self, timeout: Duration) -> io::Result<()> {