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

net_io/
session.rs
