net_io/
session.rs

1//! Sans-IO sessions: a protocol engine reacts to bytes and returns what to do next.
2
3use std::collections::VecDeque;
4use std::fmt;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::{Duration, Instant};
8
9use crate::buf::{Buffer, BufferPool};
10use crate::completer::Completer;
11use crate::error::{CloseReason, LimitKind, NetError, RefusalCause, TimeoutKind};
12use crate::gate::KillHandle;
13use crate::id::WorkerId;
14use crate::info::{ConnectionInfo, GateTag};
15use crate::limits::SizeClass;
16use crate::link::Link;
17
18/// A sans-IO protocol engine bound to one connection.
19///
20/// net-io reads, calls the session with the bytes, and does what the returned
21/// [`Step`] says. Sessions are called only on the worker that owns the
22/// connection, never concurrently.
23///
24/// # Examples
25///
26/// ```no_run
27/// use net_io::{Input, Session, SessionContext, Step};
28///
29/// struct Echo;
30/// impl Session for Echo {
31///     fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
32///         // Copy as much as the output budget takes; the rest stays in `input`,
33///         // and net-io calls on_data again after the flush because the input changed.
34///         let taken = cx.output().push_partial(input.bytes());
35///         input.consume(taken);
36///         Step::Write
37///     }
38/// }
39///
40/// let server = net_io::serve_sessions("127.0.0.1:7002", |_info: &net_io::ConnectionInfo| Echo)?;
41/// server.wait_for_stop();
42/// # Ok::<(), net_io::NetError>(())
43/// ```
44pub trait Session: Send + 'static {
45    /// Once, after admission and the handshake.
46    fn on_open(&mut self, cx: &mut SessionContext<'_>) -> Step {
47        let _ = cx;
48        Step::Read
49    }
50
51    /// Bytes arrived; `input` holds everything not yet consumed.
52    fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step;
53
54    /// After a `Step::Write` flushed all pending output.
55    fn on_writable(&mut self, cx: &mut SessionContext<'_>) -> Step {
56        let _ = cx;
57        Step::Read
58    }
59
60    /// Another thread posted a completion through this connection's [`Completer`].
61    fn on_wake(&mut self, cx: &mut SessionContext<'_>, token: u32, value: u64) -> Step {
62        let _ = (cx, token, value);
63        Step::Abort
64    }
65
66    /// A deadline passed. `Close` or `Abort` closes with the timeout's own reason
67    /// and close style; anything else re-arms with the same limits.
68    fn on_timeout(&mut self, cx: &mut SessionContext<'_>, which: TimeoutKind) -> Step {
69        let _ = (cx, which);
70        Step::Close
71    }
72
73    /// The peer sent FIN; the session may keep writing (half-close).
74    fn on_peer_closed(&mut self, cx: &mut SessionContext<'_>) -> Step {
75        let _ = cx;
76        Step::Close
77    }
78
79    /// The gate denied a write and no byte of the current response has left:
80    /// write the protocol's refusal for `code` into `output`.
81    fn on_denied(&mut self, code: u16, output: &mut Output) {
82        let _ = (code, output);
83    }
84
85    /// Exactly once, with the one close reason.
86    fn on_close(&mut self, info: &ConnectionInfo, reason: CloseReason) {
87        let _ = (info, reason);
88    }
89}
90
91/// Creates one session per admitted connection. Any `Fn(&ConnectionInfo) -> S` is one.
92///
93/// # Examples
94///
95/// ```
96/// use net_io::{ConnectionInfo, Input, RefusalCause, Session, SessionContext, SessionFactory, Step};
97///
98/// struct Discard;
99/// impl Session for Discard {
100///     fn on_data(&mut self, _cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
101///         let n = input.bytes().len();
102///         input.consume(n);
103///         Step::Read
104///     }
105/// }
106///
107/// struct OnlyLoopback;
108/// impl SessionFactory for OnlyLoopback {
109///     type Session = Discard;
110///     fn open(&self, info: &ConnectionInfo) -> Result<Discard, RefusalCause> {
111///         match info.client_addr() {
112///             Some(addr) if addr.ip().is_loopback() => Ok(Discard),
113///             _ => Err(RefusalCause::Session),
114///         }
115///     }
116/// }
117///
118/// let info = ConnectionInfo::for_testing("10.0.0.1:1".parse().unwrap());
119/// assert!(OnlyLoopback.open(&info).is_err());
120/// ```
121pub trait SessionFactory: Send + Sync + 'static {
122    /// The session type.
123    type Session: Session;
124
125    /// A session for this connection.
126    ///
127    /// # Errors
128    ///
129    /// Refusing closes the connection with `CloseReason::Refused(cause)`.
130    fn open(&self, info: &ConnectionInfo) -> Result<Self::Session, RefusalCause>;
131}
132
133impl<F, S> SessionFactory for F
134where
135    F: Fn(&ConnectionInfo) -> S + Send + Sync + 'static,
136    S: Session,
137{
138    type Session = S;
139    fn open(&self, info: &ConnectionInfo) -> Result<S, RefusalCause> {
140        Ok(self(info))
141    }
142}
143
144/// What the session wants next.
145///
146/// ```
147/// assert_ne!(net_io::Step::Close, net_io::Step::Abort);
148/// ```
149#[non_exhaustive]
150#[derive(Clone, Copy, Debug, PartialEq, Eq)]
151pub enum Step {
152    /// Wait for more input. Output pushed in this call is not sent: return
153    /// `Write` (or `Close`) to flush it first.
154    Read,
155    /// Flush pending output, then `on_writable`.
156    Write,
157    /// No reads, no idle timer, until a completion or the park deadline.
158    Park {
159        /// The deadline; `None` uses the park timeout.
160        until: Option<Instant>,
161    },
162    /// Graceful: flush, FIN, bounded linger.
163    Close,
164    /// RST now.
165    Abort,
166}
167
168/// Which deadlines run.
169///
170/// ```
171/// assert_eq!(net_io::Phase::default(), net_io::Phase::Idle);
172/// ```
173#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
174pub enum Phase {
175    /// Between exchanges: the idle timeout.
176    #[default]
177    Idle,
178    /// Receiving an exchange: the read timeout, the inbound rate, the exchange deadline.
179    Receiving,
180    /// Sending: the write timeout, the outbound rate, the exchange deadline.
181    Sending,
182}
183
184/// Why a swap was refused.
185///
186/// ```
187/// assert_ne!(net_io::SwapError::NotEmpty, net_io::SwapError::Busy);
188/// ```
189#[non_exhaustive]
190#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191pub enum SwapError {
192    /// The buffers are of different size classes.
193    ClassMismatch,
194    /// The incoming buffer has data.
195    NotEmpty,
196    /// A driver operation holds the buffer.
197    Busy,
198}
199
200impl fmt::Display for SwapError {
201    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
202        f.write_str(match self {
203            SwapError::ClassMismatch => "the buffers are of different size classes",
204            SwapError::NotEmpty => "the incoming buffer is not empty",
205            SwapError::Busy => "a driver operation holds the buffer",
206        })
207    }
208}
209
210impl std::error::Error for SwapError {}
211
212/// The connection's read buffer: a zero-copy view of what was read and not yet consumed.
213///
214/// ```
215/// let mut input = net_io::Input::from_bytes(b"GET / HTTP/1.1\r\n");
216/// assert_eq!(input.bytes(), b"GET / HTTP/1.1\r\n");
217/// input.consume(4);
218/// assert_eq!(input.bytes(), b"/ HTTP/1.1\r\n");
219/// ```
220#[derive(Debug)]
221pub struct Input {
222    pub(crate) buf: Buffer,
223    pub(crate) start: usize,
224    pub(crate) changed: bool,
225}
226
227impl Input {
228    pub(crate) fn new(buf: Buffer) -> Input {
229        Input {
230            buf,
231            start: 0,
232            changed: false,
233        }
234    }
235
236    /// An input holding `bytes`, for tests of sessions (allocates).
237    #[must_use]
238    pub fn from_bytes(bytes: &[u8]) -> Input {
239        let class = (bytes.len() as u32)
240            .max(1024)
241            .next_power_of_two()
242            .max(SizeClass::SMALL.bytes());
243        let class = SizeClass::new(class.min(1 << 20)).unwrap_or(SizeClass::LARGE);
244        let mut pool = BufferPool::standalone(u64::from(class.bytes()));
245        let mut input = Input::new(match pool.take(class) {
246            Some(b) => b,
247            None => return Input::new(Buffer::detached(class)),
248        });
249        let n = bytes.len().min(input.buf.capacity());
250        input.buf.spare()[..n].copy_from_slice(&bytes[..n]);
251        input.buf.set_len(n);
252        input
253    }
254
255    /// The unconsumed bytes.
256    #[must_use]
257    pub fn bytes(&self) -> &[u8] {
258        &self.buf.as_slice()[self.start..]
259    }
260
261    /// Consumes `n` bytes (saturating at the length).
262    pub fn consume(&mut self, n: usize) {
263        let n = n.min(self.buf.len() - self.start);
264        self.start += n;
265        if n > 0 {
266            self.changed = true;
267        }
268        if self.start == self.buf.len() {
269            self.buf.clear();
270            self.start = 0;
271        }
272    }
273
274    /// Whether the buffer is full: the session must consume or grow it.
275    #[must_use]
276    pub fn is_full(&self) -> bool {
277        self.buf.len() == self.buf.capacity() && self.start == 0
278    }
279
280    /// The buffer's size class.
281    #[must_use]
282    pub fn class(&self) -> SizeClass {
283        self.buf.class()
284    }
285
286    /// Exchanges the read buffer with an empty `other` of the same class in O(1):
287    /// `other` now holds the unconsumed bytes, and the connection reads into the
288    /// buffer it was given.
289    ///
290    /// # Errors
291    ///
292    /// [`SwapError::ClassMismatch`], [`SwapError::NotEmpty`].
293    ///
294    /// ```
295    /// use net_io::{BufferPool, Input};
296    /// let mut input = Input::from_bytes(b"body");
297    /// let mut pool = BufferPool::standalone(1 << 20);
298    /// let mut spare = pool.take(input.class()).unwrap();
299    /// input.swap(&mut spare).unwrap();
300    /// assert_eq!(spare.as_slice(), b"body");
301    /// assert!(input.bytes().is_empty());
302    /// ```
303    pub fn swap(&mut self, other: &mut Buffer) -> Result<(), SwapError> {
304        if other.class() != self.buf.class() {
305            return Err(SwapError::ClassMismatch);
306        }
307        if !other.is_empty() {
308            return Err(SwapError::NotEmpty);
309        }
310        if self.start > 0 {
311            self.buf.compact(self.start);
312            self.start = 0;
313        }
314        std::mem::swap(&mut self.buf, other);
315        self.changed = true;
316        Ok(())
317    }
318
319    /// Room to read into, compacting first when the free tail is under a quarter.
320    pub(crate) fn spare_for_read(&mut self) -> &mut [u8] {
321        let free = self.buf.capacity() - self.buf.len();
322        if self.start > 0 && free < self.buf.capacity() / 4 {
323            self.buf.compact(self.start);
324            self.start = 0;
325        }
326        self.buf.spare()
327    }
328
329    pub(crate) fn filled(&mut self, n: usize) {
330        let len = self.buf.len() + n;
331        self.buf.set_len(len);
332        self.changed = true;
333    }
334}
335
336/// The output budget was reached; return `Step::Write` and continue after the flush.
337///
338/// ```
339/// assert_eq!(net_io::OutputFull.to_string(), "the connection's output budget is full");
340/// ```
341#[derive(Clone, Copy, Debug, PartialEq, Eq)]
342pub struct OutputFull;
343
344impl fmt::Display for OutputFull {
345    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
346        f.write_str("the connection's output budget is full")
347    }
348}
349
350impl std::error::Error for OutputFull {}
351
352#[derive(Debug)]
353pub(crate) enum Chunk {
354    Pooled(Buffer, usize),
355    Static(&'static [u8], usize),
356}
357
358impl Chunk {
359    pub(crate) fn remaining(&self) -> &[u8] {
360        match self {
361            Chunk::Pooled(b, off) => &b.as_slice()[*off..],
362            Chunk::Static(s, off) => &s[*off..],
363        }
364    }
365}
366
367/// What the session wants written: a bounded queue of chunks.
368///
369/// ```
370/// let mut out = net_io::Output::standalone(64 * 1024);
371/// out.push(b"hello ").unwrap();
372/// out.push_static(b"world").unwrap();
373/// assert_eq!(out.pending(), 11);
374/// ```
375#[derive(Debug)]
376pub struct Output {
377    pub(crate) chunks: VecDeque<Chunk>,
378    free: Vec<Buffer>,
379    budget: Arc<AtomicU64>,
380    limit: u32,
381    used: u32,
382    wipe: bool,
383    pub(crate) response_open: bool,
384    pub(crate) response_tag: Option<GateTag>,
385    pub(crate) bytes_since_begin: u64,
386}
387
388const OUTPUT_CLASS: SizeClass = SizeClass::MEDIUM;
389const MAX_CHUNKS: usize = 64;
390
391impl Output {
392    pub(crate) fn new(budget: Arc<AtomicU64>, limit: u32, wipe: bool) -> Output {
393        Output {
394            chunks: VecDeque::with_capacity(MAX_CHUNKS),
395            free: Vec::new(),
396            budget,
397            limit,
398            used: 0,
399            wipe,
400            response_open: false,
401            response_tag: None,
402            bytes_since_begin: 0,
403        }
404    }
405
406    /// An output with its own budget, for tests of sessions (allocates).
407    #[must_use]
408    pub fn standalone(limit: u32) -> Output {
409        Output::new(Arc::new(AtomicU64::new(u64::from(limit))), limit, true)
410    }
411
412    fn take_buffer(&mut self) -> Option<Buffer> {
413        if let Some(b) = self.free.pop() {
414            return Some(b);
415        }
416        let bytes = u64::from(OUTPUT_CLASS.bytes());
417        let mut left = self.budget.load(Ordering::Relaxed);
418        loop {
419            if left < bytes {
420                return None;
421            }
422            match self.budget.compare_exchange_weak(
423                left,
424                left - bytes,
425                Ordering::AcqRel,
426                Ordering::Relaxed,
427            ) {
428                Ok(_) => return Some(Buffer::detached(OUTPUT_CLASS)),
429                Err(actual) => left = actual,
430            }
431        }
432    }
433
434    fn room(&self) -> usize {
435        let tail = match self.chunks.back() {
436            Some(Chunk::Pooled(b, _)) => b.capacity() - b.len(),
437            _ => 0,
438        };
439        let more = (self.limit.saturating_sub(self.used) / OUTPUT_CLASS.bytes()) as usize
440            * OUTPUT_CLASS.bytes() as usize;
441        tail + more
442    }
443
444    /// Copies as much of `bytes` as the budget allows; returns how many were taken.
445    pub fn push_partial(&mut self, mut bytes: &[u8]) -> usize {
446        let mut taken = 0;
447        while !bytes.is_empty() {
448            if let Some(Chunk::Pooled(b, _)) = self.chunks.back_mut() {
449                let spare = b.spare();
450                if !spare.is_empty() {
451                    let n = spare.len().min(bytes.len());
452                    spare[..n].copy_from_slice(&bytes[..n]);
453                    let len = b.len() + n;
454                    b.set_len(len);
455                    bytes = &bytes[n..];
456                    taken += n;
457                    continue;
458                }
459            }
460            if self.chunks.len() >= MAX_CHUNKS
461                || self.used.saturating_add(OUTPUT_CLASS.bytes()) > self.limit
462            {
463                break;
464            }
465            match self.take_buffer() {
466                Some(b) => {
467                    self.used += OUTPUT_CLASS.bytes();
468                    self.chunks.push_back(Chunk::Pooled(b, 0));
469                }
470                None => break,
471            }
472        }
473        taken
474    }
475
476    /// Copies all of `bytes`, or nothing.
477    ///
478    /// # Errors
479    ///
480    /// [`OutputFull`] when the budget cannot take all of it.
481    pub fn push(&mut self, bytes: &[u8]) -> Result<(), OutputFull> {
482        if bytes.len() > self.room() {
483            return Err(OutputFull);
484        }
485        let n = self.push_partial(bytes);
486        if n == bytes.len() {
487            Ok(())
488        } else {
489            Err(OutputFull)
490        }
491    }
492
493    /// Moves an owned buffer in (no copy).
494    ///
495    /// # Errors
496    ///
497    /// [`OutputFull`] when the budget or the chunk limit is reached.
498    pub fn push_buffer(&mut self, buffer: Buffer) -> Result<(), OutputFull> {
499        let size = u32::try_from(buffer.capacity()).unwrap_or(u32::MAX);
500        if self.chunks.len() >= MAX_CHUNKS || self.used.saturating_add(size) > self.limit {
501            return Err(OutputFull);
502        }
503        self.used += size;
504        self.chunks.push_back(Chunk::Pooled(buffer, 0));
505        Ok(())
506    }
507
508    /// Queues a static slice (no copy, not counted against the budget).
509    ///
510    /// # Errors
511    ///
512    /// [`OutputFull`] when the chunk limit is reached.
513    pub fn push_static(&mut self, bytes: &'static [u8]) -> Result<(), OutputFull> {
514        if self.chunks.len() >= MAX_CHUNKS {
515            return Err(OutputFull);
516        }
517        if !bytes.is_empty() {
518            self.chunks.push_back(Chunk::Static(bytes, 0));
519        }
520        Ok(())
521    }
522
523    /// Marks the start of a response (for the gate's denial boundary).
524    pub fn begin_response(&mut self) {
525        self.response_open = true;
526        self.response_tag = None;
527        self.bytes_since_begin = 0;
528    }
529
530    /// Marks the start of a response and tags it for the gate (`WriteCheck::response`).
531    pub fn begin_tagged_response(&mut self, tag: GateTag) {
532        self.begin_response();
533        self.response_tag = Some(tag);
534    }
535
536    /// Marks the end of the response; clears the tag.
537    pub fn end_response(&mut self) {
538        self.response_open = false;
539        self.response_tag = None;
540    }
541
542    /// Bytes queued and not yet sent.
543    #[must_use]
544    pub fn pending(&self) -> usize {
545        self.chunks.iter().map(|c| c.remaining().len()).sum()
546    }
547
548    /// Whether nothing is queued.
549    #[must_use]
550    pub fn is_empty(&self) -> bool {
551        self.chunks.is_empty()
552    }
553
554    /// Overwrites every pending chunk net-io owns with zeroes and drops it.
555    pub fn wipe_all(&mut self) {
556        while let Some(chunk) = self.chunks.pop_front() {
557            if let Chunk::Pooled(mut b, _) = chunk {
558                b.wipe();
559                self.release_chunk(b);
560            }
561        }
562        self.used = 0;
563    }
564
565    fn release_chunk(&mut self, mut b: Buffer) {
566        self.used = self
567            .used
568            .saturating_sub(u32::try_from(b.capacity()).unwrap_or(u32::MAX));
569        if self.wipe {
570            b.wipe();
571        } else {
572            b.clear();
573        }
574        if b.class() == OUTPUT_CLASS {
575            self.free.push(b);
576        } else {
577            // A buffer moved in with push_buffer goes back to the shared budget.
578            self.budget
579                .fetch_add(b.capacity() as u64, Ordering::Release);
580        }
581    }
582
583    /// Marks `n` bytes as sent, releasing finished chunks.
584    pub(crate) fn advance(&mut self, mut n: usize) {
585        while let Some(front) = self.chunks.front_mut() {
586            let rem = front.remaining().len();
587            if n < rem {
588                match front {
589                    Chunk::Pooled(_, off) | Chunk::Static(_, off) => *off += n,
590                }
591                return;
592            }
593            n -= rem;
594            if let Some(Chunk::Pooled(b, _)) = self.chunks.pop_front() {
595                self.release_chunk(b);
596            }
597        }
598    }
599
600    /// Resets for the next connection.
601    pub(crate) fn reset(&mut self) {
602        self.wipe_all();
603        self.response_open = false;
604        self.response_tag = None;
605        self.bytes_since_begin = 0;
606    }
607}
608
609/// Everything a session may touch, for one call.
610pub struct SessionContext<'a> {
611    pub(crate) link: &'a mut Link,
612    pub(crate) output: &'a mut Output,
613    pub(crate) pool: &'a mut BufferPool,
614    pub(crate) completer: &'a Completer,
615    pub(crate) draining: bool,
616    pub(crate) pressure: f32,
617    pub(crate) worker: WorkerId,
618    pub(crate) grow_to: Option<SizeClass>,
619}
620
621impl SessionContext<'_> {
622    /// The connection's facts.
623    #[must_use]
624    pub fn info(&self) -> &ConnectionInfo {
625        &self.link.info
626    }
627
628    /// The output queue.
629    pub fn output(&mut self) -> &mut Output {
630        self.output
631    }
632
633    /// The worker's clock.
634    #[must_use]
635    pub fn now(&self) -> Instant {
636        self.link.now()
637    }
638
639    /// Chooses which deadlines run.
640    pub fn set_phase(&mut self, phase: crate::session::Phase) {
641        self.link.set_phase(phase);
642    }
643
644    /// The deadline for the current exchange (`None` clears it).
645    pub fn set_exchange_deadline(&mut self, at: Option<Instant>) {
646        self.link.set_exchange_deadline(at);
647    }
648
649    /// Overrides this connection's idle timeout from its next wait for an
650    /// exchange on; `None` restores the listener's limit. For example, a
651    /// protocol shortens keep-alive while the worker pool is nearly full.
652    /// Values below the timer tick are raised to one tick.
653    ///
654    /// ```no_run
655    /// use std::time::Duration;
656    /// use net_io::{Input, Session, SessionContext, Step};
657    ///
658    /// struct KeepAlive;
659    /// impl Session for KeepAlive {
660    ///     fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
661    ///         let n = input.bytes().len();
662    ///         input.consume(n);
663    ///         let short = (cx.pressure() > 0.75).then_some(Duration::from_secs(2));
664    ///         cx.set_idle_timeout(short);
665    ///         Step::Write
666    ///     }
667    /// }
668    /// ```
669    pub fn set_idle_timeout(&mut self, timeout: Option<Duration>) {
670        self.link.set_idle_timeout(timeout);
671    }
672
673    /// Whether drain started: finish the current exchange, then close.
674    #[must_use]
675    pub fn draining(&self) -> bool {
676        self.draining
677    }
678
679    /// 0.0 to 1.0: busy workers (blocking driver); 1.0 while the control line sheds.
680    #[must_use]
681    pub fn pressure(&self) -> f32 {
682        self.pressure
683    }
684
685    /// A handle other threads can use to abort this connection.
686    #[must_use]
687    pub fn kill_handle(&self) -> KillHandle {
688        KillHandle {
689            slot: self.link.slot.clone(),
690            generation: self.link.generation,
691            id: self.link.info.id,
692        }
693    }
694
695    /// This worker's completion queue.
696    #[must_use]
697    pub fn completer(&self) -> Completer {
698        self.completer.clone()
699    }
700
701    /// The worker that runs this connection.
702    #[must_use]
703    pub fn worker(&self) -> WorkerId {
704        self.worker
705    }
706
707    /// This worker's buffer pool.
708    pub fn buffers(&mut self) -> &mut BufferPool {
709        self.pool
710    }
711
712    /// Returns a buffer (for example one taken by `Input::swap`); it is wiped.
713    pub fn release(&mut self, buffer: Buffer) {
714        self.pool.put(buffer);
715    }
716
717    /// Moves the unconsumed input into a buffer of `class` after this call returns.
718    ///
719    /// # Errors
720    ///
721    /// `Limit { which: ReadBuffer }` above the maximum read class.
722    pub fn grow_input(&mut self, class: SizeClass) -> Result<(), NetError> {
723        if class > self.link.limits.read_class_max {
724            return Err(NetError::Limit {
725                which: LimitKind::ReadBuffer,
726            });
727        }
728        self.grow_to = Some(class);
729        Ok(())
730    }
731}
732
733impl fmt::Debug for SessionContext<'_> {
734    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
735        f.debug_struct("SessionContext")
736            .field("id", &self.link.info.id)
737            .field("draining", &self.draining)
738            .finish()
739    }
740}
741
742#[cfg(test)]
743mod tests {
744    use super::*;
745
746    #[test]
747    fn test_input_consume_saturates_and_resets() {
748        let mut i = Input::from_bytes(b"abc");
749        i.consume(10);
750        assert!(i.bytes().is_empty());
751        assert_eq!(i.start, 0);
752    }
753
754    #[test]
755    fn test_input_swap_rules() {
756        let mut i = Input::from_bytes(b"xyz");
757        let mut pool = BufferPool::standalone(1 << 22);
758        let mut wrong = pool.take(SizeClass::LARGE).unwrap();
759        assert_eq!(i.swap(&mut wrong), Err(SwapError::ClassMismatch));
760        let mut full = pool.take(i.class()).unwrap();
761        full.set_len(1);
762        assert_eq!(i.swap(&mut full), Err(SwapError::NotEmpty));
763        let mut empty = pool.take(i.class()).unwrap();
764        i.consume(1);
765        i.swap(&mut empty).unwrap();
766        assert_eq!(empty.as_slice(), b"yz");
767    }
768
769    #[test]
770    fn test_output_push_all_or_nothing_and_partial() {
771        let mut o = Output::standalone(OUTPUT_CLASS.bytes());
772        let big = vec![7u8; OUTPUT_CLASS.bytes() as usize + 1];
773        assert_eq!(o.push(&big), Err(OutputFull));
774        assert!(o.is_empty());
775        assert_eq!(o.push_partial(&big), OUTPUT_CLASS.bytes() as usize);
776        assert_eq!(o.pending(), OUTPUT_CLASS.bytes() as usize);
777    }
778
779    #[test]
780    fn test_output_advance_releases_chunks() {
781        let mut o = Output::standalone(1 << 20);
782        o.push(b"hello").unwrap();
783        o.push_static(b" world").unwrap();
784        o.advance(7);
785        assert_eq!(o.pending(), 4);
786        assert_eq!(o.chunks.len(), 1);
787        assert_eq!(o.used, 0);
788        o.advance(4);
789        assert!(o.is_empty());
790    }
791
792    #[test]
793    fn test_output_response_markers() {
794        let mut o = Output::standalone(1 << 20);
795        o.begin_tagged_response(GateTag(3));
796        assert_eq!(o.response_tag, Some(GateTag(3)));
797        o.end_response();
798        assert_eq!(o.response_tag, None);
799    }
800
801    /// Property: whatever mix of pushes is accepted and however the sends split,
802    /// the bytes that leave are the accepted bytes in order, and the budget returns.
803    #[test]
804    fn prop_output_stream_matches_accepted_pushes() {
805        use crate::testing::Gen;
806        static TEXT: [u8; 4096] = {
807            let mut t = [0u8; 4096];
808            let mut i = 0;
809            while i < t.len() {
810                t[i] = (i % 251) as u8;
811                i += 1;
812            }
813            t
814        };
815        let mut g = Gen::new(0x0047_0041);
816        for round in 0..if cfg!(miri) { 3 } else { 300 } {
817            let mut o = Output::standalone(64 * 1024);
818            let mut expected = Vec::new();
819            let mut sent = Vec::new();
820            for _ in 0..200 {
821                match g.below(4) {
822                    0 => {
823                        let n = g.below(40_000) as usize;
824                        let data: Vec<u8> = (0..n).map(|_| g.next_u64() as u8).collect();
825                        let took = o.push_partial(&data);
826                        expected.extend_from_slice(&data[..took]);
827                    }
828                    1 => {
829                        let n = g.below(3000) as usize;
830                        let data: Vec<u8> = (0..n).map(|_| g.next_u64() as u8).collect();
831                        if o.push(&data).is_ok() {
832                            expected.extend_from_slice(&data);
833                        }
834                    }
835                    2 => {
836                        let start = g.below(TEXT.len() as u64) as usize;
837                        let piece = &TEXT[start..];
838                        if o.push_static(piece).is_ok() {
839                            expected.extend_from_slice(piece);
840                        }
841                    }
842                    _ => {
843                        // A send of a random size, taken from the front chunks as a vectored write would.
844                        let mut want = g.below(70_000) as usize;
845                        let mut taken = 0;
846                        for chunk in o.chunks.iter() {
847                            if want == 0 {
848                                break;
849                            }
850                            let r = chunk.remaining();
851                            let n = r.len().min(want);
852                            sent.extend_from_slice(&r[..n]);
853                            want -= n;
854                            taken += n;
855                        }
856                        o.advance(taken);
857                    }
858                }
859            }
860            while !o.is_empty() {
861                let n = o.chunks.front().map_or(0, |c| c.remaining().len());
862                if let Some(c) = o.chunks.front() {
863                    sent.extend_from_slice(c.remaining());
864                }
865                o.advance(n);
866            }
867            assert!(sent == expected, "round {round}: stream differs");
868            assert_eq!(o.used, 0, "round {round}: budget not returned");
869            assert_eq!(o.pending(), 0);
870        }
871    }
872}
net_io/session.rs

net_io/
session.rs
