conmonrs/streaming_server/session/
pending_stdin.rs

1//! The stdin of a client which waits for the container, see
2//! [`PendingStdin`].
3
4use super::PENDING_STDIN_MAX_BYTES;
5use crate::streaming_server::remotecommand::Stdin;
6use bytes::{Bytes, BytesMut};
7use std::collections::VecDeque;
8
9/// The largest part of stdin which goes towards the container at once.
10/// Larger stdin gets split without copying it, so that the channels on the
11/// way, which count entries, hold at most this much per entry.
12pub(super) const STDIN_CHUNK_SIZE: usize = 64 * 1024;
13
14/// Stdin smaller than this gets merged with the stdin before it while it
15/// waits, so that a client which sends a byte per message does not make the
16/// server hold far more than the bytes it counts.
17pub(super) const STDIN_MERGE_SIZE: usize = 32 * 1024;
18
19/// What an entry of the pending stdin counts on top of its bytes.
20pub(super) const STDIN_ENTRY_OVERHEAD: usize = 64;
21
22/// Stdin of a client which the container did not take yet, bounded by
23/// [`PENDING_STDIN_MAX_BYTES`]. Small stdin gets merged and large stdin
24/// split into entries of at most [`STDIN_CHUNK_SIZE`], each counts
25/// [`STDIN_ENTRY_OVERHEAD`] on top of its bytes. Empty stdin carries nothing,
26/// and nothing takes stdin after the close, both get dropped. The entries
27/// of large stdin keep its message until the container took the last one,
28/// stdin goes to the container in order, so only the oldest message is
29/// partly taken.
30#[derive(Debug, Default)]
31pub(super) struct PendingStdin {
32    /// Stdin of at least [`STDIN_MERGE_SIZE`] bytes and merged small stdin,
33    /// the oldest first.
34    chunks: VecDeque<Bytes>,
35
36    /// Small stdin after the chunks, merged.
37    tail: BytesMut,
38
39    /// The bytes and the overhead of the chunks.
40    bytes: usize,
41
42    /// Whether the client closed its stdin.
43    pub(super) closed: bool,
44
45    /// Whether the close still has to be taken, after the stdin before it.
46    close_pending: bool,
47}
48
49impl PendingStdin {
50    pub(super) fn is_empty(&self) -> bool {
51        self.chunks.is_empty() && self.tail.is_empty() && !self.close_pending
52    }
53
54    /// What the pending stdin counts against [`PENDING_STDIN_MAX_BYTES`].
55    pub(super) fn len(&self) -> usize {
56        let tail = if self.tail.is_empty() {
57            0
58        } else {
59            self.tail.len() + STDIN_ENTRY_OVERHEAD
60        };
61        let close = if self.close_pending {
62            STDIN_ENTRY_OVERHEAD
63        } else {
64            0
65        };
66        self.bytes + tail + close
67    }
68
69    /// Whether the client has to wait, which stops reading the socket.
70    pub(super) fn is_full(&self) -> bool {
71        self.len() >= PENDING_STDIN_MAX_BYTES
72    }
73
74    pub(super) fn push(&mut self, stdin: Stdin) {
75        if self.closed {
76            return;
77        }
78        match stdin {
79            Stdin::Close => {
80                self.closed = true;
81                self.close_pending = true;
82            }
83            Stdin::Data(data) if data.len() < STDIN_MERGE_SIZE => {
84                if self.tail.len() + data.len() > STDIN_CHUNK_SIZE {
85                    self.flush_tail();
86                }
87                self.tail.extend_from_slice(&data);
88            }
89            Stdin::Data(mut data) => {
90                self.flush_tail();
91                while !data.is_empty() {
92                    let chunk = data.split_to(data.len().min(STDIN_CHUNK_SIZE));
93                    self.push_chunk(chunk);
94                }
95            }
96        }
97    }
98
99    /// The merged stdin becomes a chunk, so that later stdin goes behind it.
100    fn flush_tail(&mut self) {
101        if !self.tail.is_empty() {
102            let tail = self.take_tail();
103            self.push_chunk(tail);
104        }
105    }
106
107    /// The merged stdin in a buffer of its own. A part of the buffer of the
108    /// tail would keep all of it, which can be far more than the stdin it
109    /// counts. The tail keeps its buffer, of at most twice
110    /// [`STDIN_CHUNK_SIZE`], for the next stdin.
111    fn take_tail(&mut self) -> Bytes {
112        let tail = Bytes::copy_from_slice(&self.tail);
113        self.tail.clear();
114        tail
115    }
116
117    fn push_chunk(&mut self, chunk: Bytes) {
118        self.bytes += chunk.len() + STDIN_ENTRY_OVERHEAD;
119        self.chunks.push_back(chunk);
120    }
121
122    /// The oldest stdin, callers check that there is some.
123    pub(super) fn pop(&mut self) -> Stdin {
124        if let Some(chunk) = self.chunks.pop_front() {
125            self.bytes -= chunk.len() + STDIN_ENTRY_OVERHEAD;
126            return Stdin::Data(chunk);
127        }
128        if !self.tail.is_empty() {
129            return Stdin::Data(self.take_tail());
130        }
131        self.close_pending = false;
132        Stdin::Close
133    }
134
135    /// Drop the stdin, nothing takes it anymore.
136    pub(super) fn clear(&mut self) {
137        self.chunks.clear();
138        self.tail.clear();
139        self.bytes = 0;
140        self.close_pending = false;
141    }
142}
conmonrs/streaming_server/session/pending_stdin.rs

conmonrs/streaming_server/session/
pending_stdin.rs
