use std::num::NonZeroUsize;
use bytes::BufMut;
use tracing::trace;
use crate::packet::BufLen;
/// The buffer in which to write datagrams for [`Connection::poll_transmit`]
///
/// The `poll_transmit` function writes zero or more datagrams to a buffer. Multiple
/// datagrams are possible in case GSO (Generic Segmentation Offload) is supported.
///
/// This buffer tracks datagrams being written to it. There is always a "current" datagram,
/// which is started by calling [`TransmitBuf::start_new_datagram`]. Writing to the buffer
/// is done through the [`BufMut`] interface.
///
/// Usually a datagram contains one QUIC packet, though QUIC-TRANSPORT 12.2 Coalescing
/// Packets allows for placing multiple packets into a single datagram provided all but the
/// last packet uses long headers. This is normally used during connection setup where often
/// the initial, handshake and sometimes even a 1-RTT packet can be coalesced into a single
/// datagram.
///
/// Inside a single packet multiple QUIC frames are written.
///
/// The buffer managed here is passed straight to the OS' `sendmsg` call (or variant) once
/// `poll_transmit` returns. So needs to contain the datagrams as they are sent on the
/// wire.
///
/// [`Connection::poll_transmit`]: super::Connection::poll_transmit
#[derive(Debug)]
pub(super) struct TransmitBuf<'a> {
/// The buffer itself, packets are written to this buffer
buf: &'a mut Vec<u8>,
/// Offset into the buffer at which the current datagram starts
///
/// Note that when coalescing packets this might be before the start of the current
/// packet.
datagram_start: usize,
/// The maximum offset allowed to be used for the current datagram in the buffer
///
/// The first and last datagram in a batch are allowed to be smaller then the maximum
/// size. All datagrams in between need to be exactly this size.
buf_capacity: usize,
/// The maximum number of datagrams allowed to write into [`TransmitBuf::buf`]
max_datagrams: NonZeroUsize,
/// The number of datagrams already (partially) written into the buffer
///
/// Incremented by a call to [`TransmitBuf::start_new_datagram`].
pub(super) num_datagrams: usize,
/// The segment size of this GSO batch
///
/// The segment size is the size of each datagram in the GSO batch, only the last
/// datagram in the batch may be smaller.
///
/// For the first datagram this is set to the maximum size a datagram is allowed to be:
/// the current path MTU. After the first datagram is finished this is reduced to the
/// size of the first datagram and can no longer change.
segment_size: usize,
}
impl<'a> TransmitBuf<'a> {
pub(super) fn new(buf: &'a mut Vec<u8>, max_datagrams: NonZeroUsize, mtu: usize) -> Self {
buf.clear();
Self {
buf,
datagram_start: 0,
buf_capacity: 0,
max_datagrams,
num_datagrams: 0,
segment_size: mtu,
}
}
/// Starts a datagram with a custom datagram size
///
/// This is a specialized version of [`TransmitBuf::start_new_datagram`] which sets the
/// datagram size. Useful for e.g. PATH_CHALLENGE, tail-loss probes or MTU probes.
///
/// After the first datagram you can never increase the segment size. If you decrease
/// the size of a datagram in a batch, it must be the last datagram of the batch.
pub(super) fn start_new_datagram_with_size(&mut self, datagram_size: usize) {
// Only reserve space for this datagram, usually it is the last one in the batch.
let max_capacity_hint = datagram_size;
self.new_datagram_inner(datagram_size, max_capacity_hint)
}
/// Starts a new datagram in the transmit buffer
///
/// If this starts the second datagram the segment size will be set to the size of the
/// first datagram.
///
/// If the underlying buffer does not have enough capacity yet this will allocate enough
/// capacity for all the datagrams allowed in a single batch. Use
/// [`TransmitBuf::start_new_datagram_with_size`] if you know you will need less.
pub(super) fn start_new_datagram(&mut self) {
// We reserve the maximum space for sending `max_datagrams` upfront to avoid any
// reallocations if more datagrams have to be appended later on. Benchmarks have
// shown a 5-10% throughput improvement compared to continuously resizing the
// datagram buffer. While this will lead to over-allocation for small transmits
// (e.g. purely containing ACKs), modern memory allocators (e.g. mimalloc and
// jemalloc) will pool certain allocation sizes and therefore this is still rather
// efficient.
let max_capacity_hint = self.max_datagrams.get() * self.segment_size;
self.new_datagram_inner(self.segment_size, max_capacity_hint)
}
fn new_datagram_inner(&mut self, datagram_size: usize, max_capacity_hint: usize) {
debug_assert!(self.num_datagrams < self.max_datagrams.into());
if self.num_datagrams == 1 {
// Set the segment size to the size of the first datagram.
self.segment_size = self.buf.len();
}
if self.num_datagrams >= 1 {
debug_assert!(datagram_size <= self.segment_size);
if datagram_size < self.segment_size {
// If this is a GSO batch and this datagram is smaller than the segment
// size, this must be the last datagram in the batch.
self.max_datagrams = NonZeroUsize::MIN.saturating_add(self.num_datagrams);
}
}
self.datagram_start = self.buf.len();
debug_assert_eq!(
self.datagram_start % self.segment_size,
0,
"datagrams in a GSO batch must be aligned to the segment size"
);
self.buf_capacity = self.datagram_start + datagram_size;
if self.buf_capacity > self.buf.capacity() {
self.buf
.reserve_exact(max_capacity_hint.saturating_sub(self.buf.capacity()));
}
self.num_datagrams += 1;
}
/// Clips the segment size to the current size
///
/// Only valid for the first datagram, when the datagram might be smaller than the
/// segment size. Needed before estimating the available space in the next datagram
/// based on [`TransmitBuf::segment_size`].
///
/// Use [`TransmitBuf::start_new_datagram_with_size`] if you need to reduce the size of
/// the last datagram in a batch.
pub(super) fn clip_segment_size(&mut self) {
debug_assert_eq!(self.num_datagrams, 1);
if self.buf.len() < self.segment_size {
trace!(
segment_size = self.buf.len(),
prev_segment_size = self.segment_size,
"clipped datagram size"
);
}
self.segment_size = self.buf.len();
self.buf_capacity = self.buf.len();
}
/// Returns the GSO segment size
///
/// This is also the maximum size datagrams are allowed to be. The first and last
/// datagram in a batch are allowed to be smaller however. After the first datagram the
/// segment size is clipped to the size of the first datagram.
///
/// If the last datagram was created using [`TransmitBuf::start_new_datagram_with_size`]
/// the the segment size will be greater than the current datagram is allowed to be.
/// Thus [`TransmitBuf::datagram_remaining_mut`] should be used if you need to know the
/// amount of data that can be written into the datagram.
pub(super) fn segment_size(&self) -> usize {
self.segment_size
}
/// Returns the number of datagrams written into the buffer
///
/// The last datagram is not necessarily finished yet.
pub(super) fn num_datagrams(&self) -> usize {
self.num_datagrams
}
/// Returns the maximum number of datagrams allowed to be written into the buffer
pub(super) fn max_datagrams(&self) -> NonZeroUsize {
self.max_datagrams
}
/// Returns the start offset of the current datagram in the buffer
///
/// In other words, this offset contains the first byte of the current datagram.
pub(super) fn datagram_start_offset(&self) -> usize {
self.datagram_start
}
/// Returns the maximum offset in the buffer allowed for the current datagram
///
/// The first and last datagram in a batch are allowed to be smaller then the maximum
/// size. All datagrams in between need to be exactly this size.
pub(super) fn datagram_max_offset(&self) -> usize {
self.buf_capacity
}
/// Returns the number of bytes that may still be written into this datagram
pub(super) fn datagram_remaining_mut(&self) -> usize {
self.buf_capacity.saturating_sub(self.buf.len())
}
/// Returns `true` if the buffer did not have anything written into it
pub(super) fn is_empty(&self) -> bool {
self.len() == 0
}
/// The number of bytes written into the buffer so far
pub(super) fn len(&self) -> usize {
self.buf.len()
}
/// Returns the already written bytes in the buffer
pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
self.buf.as_mut_slice()
}
}
unsafe impl BufMut for TransmitBuf<'_> {
fn remaining_mut(&self) -> usize {
self.buf.remaining_mut()
}
unsafe fn advance_mut(&mut self, cnt: usize) {
unsafe { self.buf.advance_mut(cnt) };
}
fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
self.buf.chunk_mut()
}
}
impl BufLen for TransmitBuf<'_> {
fn len(&self) -> usize {
self.len()
}
}
use NonZeroUsize;
use BufMut;
use trace;
use crateBufLen;
/// The buffer in which to write datagrams for [`Connection::poll_transmit`]
///
/// The `poll_transmit` function writes zero or more datagrams to a buffer. Multiple
/// datagrams are possible in case GSO (Generic Segmentation Offload) is supported.
///
/// This buffer tracks datagrams being written to it. There is always a "current" datagram,
/// which is started by calling [`TransmitBuf::start_new_datagram`]. Writing to the buffer
/// is done through the [`BufMut`] interface.
///
/// Usually a datagram contains one QUIC packet, though QUIC-TRANSPORT 12.2 Coalescing
/// Packets allows for placing multiple packets into a single datagram provided all but the
/// last packet uses long headers. This is normally used during connection setup where often
/// the initial, handshake and sometimes even a 1-RTT packet can be coalesced into a single
/// datagram.
///
/// Inside a single packet multiple QUIC frames are written.
///
/// The buffer managed here is passed straight to the OS' `sendmsg` call (or variant) once
/// `poll_transmit` returns. So needs to contain the datagrams as they are sent on the
/// wire.
///
/// [`Connection::poll_transmit`]: super::Connection::poll_transmit
pub
unsafe