compio_io\compat/
sync_stream.rs

1use std::{
2    io::{self, BufRead, Read, Write},
3    mem::MaybeUninit,
4};
5
6use compio_buf::{BufResult, IntoInner, IoBufExt, IoBufMut, IoBufMutExt};
7
8use crate::{
9    buffer::Buffer,
10    util::{DEFAULT_BUF_SIZE, Splittable},
11};
12
13// 64MiB max
14pub(crate) const DEFAULT_MAX_BUFFER: usize = 64 * 1024 * 1024;
15
16#[derive(Debug)]
17struct SyncReadBuf {
18    buf: Buffer,
19    eof: bool,
20    base_capacity: usize,
21    max_buffer_size: usize,
22}
23
24impl SyncReadBuf {
25    pub fn new(start_capacity: usize, base_capacity: usize, max_buffer_size: usize) -> Self {
26        Self {
27            buf: Buffer::with_capacity(start_capacity),
28            eof: false,
29            base_capacity,
30            max_buffer_size,
31        }
32    }
33
34    pub fn is_eof(&self) -> bool {
35        self.eof
36    }
37
38    pub fn into_inner(mut self) -> Vec<u8> {
39        if self.buf.has_inner() {
40            let slice = self.buf.take_inner();
41            let begin = slice.begin();
42            let mut vec = slice.into_inner();
43            if begin > 0 {
44                vec.drain(..begin);
45            }
46            vec
47        } else {
48            Vec::new()
49        }
50    }
51
52    /// Returns the available bytes in the read buffer.
53    fn available_read(&self) -> io::Result<&[u8]> {
54        if self.buf.has_inner() {
55            Ok(self.buf.buffer())
56        } else {
57            Err(would_block("the read buffer is in use"))
58        }
59    }
60
61    /// Marks `amt` bytes as consumed from the read buffer.
62    ///
63    /// Resets the buffer when all data is consumed and shrinks capacity
64    /// if it has grown significantly beyond the base capacity.
65    pub fn consume(&mut self, amt: usize) {
66        let all_done = self.buf.advance(amt);
67
68        // Shrink oversized buffers back to base capacity
69        if all_done {
70            self.buf
71                .compact_to(self.base_capacity, self.max_buffer_size);
72        }
73    }
74
75    pub fn read_buf_uninit(&mut self, buf: &mut [MaybeUninit<u8>]) -> io::Result<usize> {
76        let available = self.fill_buf()?;
77
78        let to_read = available.len().min(buf.len());
79        buf[..to_read].write_copy_of_slice(&available[..to_read]);
80        self.consume(to_read);
81
82        Ok(to_read)
83    }
84
85    pub fn fill_buf(&mut self) -> io::Result<&[u8]> {
86        let available = self.available_read()?;
87
88        if available.is_empty() && !self.eof {
89            return Err(would_block("need to fill read buffer"));
90        }
91
92        Ok(available)
93    }
94
95    pub fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
96        let mut slice = self.fill_buf()?;
97        slice.read(buf).inspect(|res| {
98            self.consume(*res);
99        })
100    }
101
102    #[cfg(feature = "read_buf")]
103    pub fn read_buf(&mut self, mut buf: io::BorrowedCursor<'_, u8>) -> io::Result<()> {
104        let mut slice = self.fill_buf()?;
105        let old_written = buf.written();
106        slice.read_buf(buf.reborrow())?;
107        let len = buf.written() - old_written;
108        self.consume(len);
109        Ok(())
110    }
111
112    pub async fn fill_read_buf<S: crate::AsyncRead>(
113        &mut self,
114        stream: &mut S,
115    ) -> io::Result<usize> {
116        if self.eof {
117            return Ok(0);
118        }
119
120        // Compact buffer, move unconsumed data to the front
121        self.buf
122            .compact_to(self.base_capacity, self.max_buffer_size);
123
124        let read = self
125            .buf
126            .with(|mut inner| async {
127                let current_len = inner.buf_len();
128
129                if current_len >= self.max_buffer_size {
130                    return BufResult(
131                        Err(io::Error::new(
132                            io::ErrorKind::OutOfMemory,
133                            format!("read buffer size limit ({}) exceeded", self.max_buffer_size),
134                        )),
135                        inner,
136                    );
137                }
138
139                let capacity = inner.buf_capacity();
140                let available_space = capacity - current_len;
141
142                // If target space is less than base capacity, grow the buffer.
143                let target_space = self.base_capacity;
144                if available_space < target_space {
145                    let new_capacity = current_len + target_space;
146                    let _ = inner.reserve_exact(new_capacity - capacity);
147                }
148
149                let len = inner.buf_len();
150                let read_slice = inner.slice(len..);
151                stream.read(read_slice).await.into_inner()
152            })
153            .await?;
154        if read == 0 {
155            self.eof = true;
156        }
157        Ok(read)
158    }
159}
160
161#[derive(Debug)]
162struct SyncWriteBuf {
163    buf: Buffer,
164    base_capacity: usize,
165    max_buffer_size: usize,
166}
167
168impl SyncWriteBuf {
169    pub fn new(start_capacity: usize, base_capacity: usize, max_buffer_size: usize) -> Self {
170        Self {
171            buf: Buffer::with_capacity(start_capacity),
172            base_capacity,
173            max_buffer_size,
174        }
175    }
176
177    pub fn has_pending_write(&self) -> bool {
178        !self.buf.is_empty()
179    }
180
181    pub fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
182        if !self.buf.has_inner() {
183            return Err(would_block("the write buffer is in use"));
184        }
185        // Check if we should flush first
186        if self.buf.need_flush() && !self.buf.is_empty() {
187            return Err(would_block("need to flush write buffer"));
188        }
189
190        let written = self.buf.with_sync(|mut inner| {
191            let res = (|| {
192                if inner.buf_len() + buf.len() > self.max_buffer_size {
193                    let space = self.max_buffer_size - inner.buf_len();
194                    if space == 0 {
195                        Err(would_block("write buffer full, need to flush"))
196                    } else {
197                        inner.extend_from_slice(&buf[..space])?;
198                        Ok(space)
199                    }
200                } else {
201                    inner.extend_from_slice(buf)?;
202                    Ok(buf.len())
203                }
204            })();
205            BufResult(res, inner)
206        })?;
207
208        Ok(written)
209    }
210
211    pub async fn flush_write_buf<S: crate::AsyncWrite>(
212        &mut self,
213        stream: &mut S,
214    ) -> io::Result<usize> {
215        let flushed = self.buf.flush_to(stream).await?;
216        self.buf
217            .compact_to(self.base_capacity, self.max_buffer_size);
218        stream.flush().await?;
219        Ok(flushed)
220    }
221}
222
223/// A growable buffered stream adapter that bridges async I/O with sync traits.
224///
225/// # Buffer Growth Strategy
226///
227/// - **Read buffer**: Grows as needed to accommodate incoming data, up to
228///   `max_buffer_size`
229/// - **Write buffer**: Grows as needed for outgoing data, up to
230///   `max_buffer_size`
231/// - Both buffers shrink back to `base_capacity` when fully consumed and
232///   capacity exceeds 4x base
233///
234/// # Usage Pattern
235///
236/// The sync `Read` and `Write` implementations will return `WouldBlock` errors
237/// when buffers need servicing via the async methods:
238///
239/// - Call `fill_read_buf()` when `Read::read()` returns `WouldBlock`
240/// - Call `flush_write_buf()` when `Write::write()` returns `WouldBlock`
241///
242/// # Note on flush()
243///
244/// The `Write::flush()` method intentionally returns `Ok(())` without checking
245/// if there's buffered data. This is for compatibility with libraries like
246/// tungstenite that call `flush()` after every write. Actual flushing happens
247/// via the async `flush_write_buf()` method.
248#[derive(Debug)]
249pub struct SyncStream<S> {
250    inner: S,
251    read_buf: SyncReadBuf,
252    write_buf: SyncWriteBuf,
253}
254
255/// Read half of a [`SyncStream`] after splitting.
256#[derive(Debug)]
257pub struct SyncStreamReadHalf<S> {
258    inner: S,
259    read_buf: SyncReadBuf,
260}
261
262/// Write half of a [`SyncStream`] after splitting.
263#[derive(Debug)]
264pub struct SyncStreamWriteHalf<S> {
265    inner: S,
266    write_buf: SyncWriteBuf,
267}
268
269impl<S> SyncStream<S> {
270    /// Creates a new `SyncStream` with default buffer sizes.
271    ///
272    /// - Base capacity: 8KiB
273    /// - Max buffer size: 64MiB
274    pub fn new(stream: S) -> Self {
275        Self::with_capacity(DEFAULT_BUF_SIZE, stream)
276    }
277
278    /// Creates a new `SyncStream` with a custom base capacity.
279    ///
280    /// The maximum buffer size defaults to 64MiB.
281    pub fn with_capacity(base_capacity: usize, stream: S) -> Self {
282        Self::with_limits(base_capacity, DEFAULT_MAX_BUFFER, stream)
283    }
284
285    /// Creates a new `SyncStream` with custom base capacity and maximum
286    /// buffer size.
287    pub fn with_limits(base_capacity: usize, max_buffer_size: usize, stream: S) -> Self {
288        Self {
289            inner: stream,
290            read_buf: SyncReadBuf::new(base_capacity, base_capacity, max_buffer_size),
291            write_buf: SyncWriteBuf::new(base_capacity, base_capacity, max_buffer_size),
292        }
293    }
294
295    /// Returns a reference to the underlying stream.
296    pub fn get_ref(&self) -> &S {
297        &self.inner
298    }
299
300    /// Returns a mutable reference to the underlying stream.
301    pub fn get_mut(&mut self) -> &mut S {
302        &mut self.inner
303    }
304
305    /// Consumes the `SyncStream`, returning the underlying stream.
306    ///
307    /// Any buffered data is discarded. Use [`into_parts`](Self::into_parts)
308    /// if you need to preserve unread data.
309    pub fn into_inner(self) -> S {
310        self.inner
311    }
312
313    /// Consumes the `SyncStream`, returning the underlying stream and any
314    /// unread buffered data.
315    ///
316    /// If the read buffer is currently lent to an IO operation, the returned
317    /// `Vec` will be empty.
318    pub fn into_parts(self) -> (S, Vec<u8>) {
319        let remaining = self.read_buf.into_inner();
320        (self.inner, remaining)
321    }
322
323    /// Returns `true` if the stream has reached EOF.
324    pub fn is_eof(&self) -> bool {
325        self.read_buf.is_eof()
326    }
327
328    /// Pull some bytes from this source into the specified buffer.
329    pub fn read_buf_uninit(&mut self, buf: &mut [MaybeUninit<u8>]) -> io::Result<usize> {
330        self.read_buf.read_buf_uninit(buf)
331    }
332
333    /// Returns `true` if there is pending data in the write buffer that needs
334    /// to be flushed.
335    pub fn has_pending_write(&self) -> bool {
336        self.write_buf.has_pending_write()
337    }
338}
339
340impl<S> SyncStreamReadHalf<S> {
341    pub(crate) fn with_limits(base_capacity: usize, max_buffer_size: usize, stream: S) -> Self {
342        Self {
343            inner: stream,
344            read_buf: SyncReadBuf::new(base_capacity, base_capacity, max_buffer_size),
345        }
346    }
347
348    /// Returns a reference to the underlying stream.
349    pub fn get_ref(&self) -> &S {
350        &self.inner
351    }
352
353    /// Returns a mutable reference to the underlying stream.
354    pub fn get_mut(&mut self) -> &mut S {
355        &mut self.inner
356    }
357
358    /// Consumes the `SyncStreamReadHalf`, returning the underlying stream.
359    pub fn into_inner(self) -> S {
360        self.inner
361    }
362
363    /// Consumes the `SyncStream`, returning the underlying stream and any
364    /// unread buffered data.
365    ///
366    /// If the read buffer is currently lent to an IO operation, the returned
367    /// `Vec` will be empty.
368    pub fn into_parts(self) -> (S, Vec<u8>) {
369        let remaining = self.read_buf.into_inner();
370        (self.inner, remaining)
371    }
372
373    /// Returns `true` if the stream has reached EOF.
374    pub fn is_eof(&self) -> bool {
375        self.read_buf.is_eof()
376    }
377
378    /// Pull some bytes from this source into the specified buffer.
379    pub fn read_buf_uninit(&mut self, buf: &mut [MaybeUninit<u8>]) -> io::Result<usize> {
380        self.read_buf.read_buf_uninit(buf)
381    }
382}
383
384impl<S> SyncStreamWriteHalf<S> {
385    pub(crate) fn with_limits(base_capacity: usize, max_buffer_size: usize, stream: S) -> Self {
386        Self {
387            inner: stream,
388            write_buf: SyncWriteBuf::new(base_capacity, base_capacity, max_buffer_size),
389        }
390    }
391
392    /// Returns a reference to the underlying stream.
393    pub fn get_ref(&self) -> &S {
394        &self.inner
395    }
396
397    /// Returns a mutable reference to the underlying stream.
398    pub fn get_mut(&mut self) -> &mut S {
399        &mut self.inner
400    }
401
402    /// Consumes the `SyncStreamWriteHalf`, returning the underlying stream.
403    pub fn into_inner(self) -> S {
404        self.inner
405    }
406
407    /// Returns `true` if there is pending data in the write buffer that needs
408    /// to be flushed.
409    pub fn has_pending_write(&self) -> bool {
410        self.write_buf.has_pending_write()
411    }
412}
413
414impl<S> Read for SyncStream<S> {
415    /// Reads data from the internal buffer.
416    ///
417    /// Returns `WouldBlock` if the buffer is empty and not at EOF,
418    /// indicating that `fill_read_buf()` should be called.
419    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
420        self.read_buf.read(buf)
421    }
422
423    #[cfg(feature = "read_buf")]
424    fn read_buf(&mut self, buf: io::BorrowedCursor<'_, u8>) -> io::Result<()> {
425        self.read_buf.read_buf(buf)
426    }
427}
428
429impl<S> BufRead for SyncStream<S> {
430    fn fill_buf(&mut self) -> io::Result<&[u8]> {
431        self.read_buf.fill_buf()
432    }
433
434    fn consume(&mut self, amt: usize) {
435        self.read_buf.consume(amt);
436    }
437}
438
439impl<S> Read for SyncStreamReadHalf<S> {
440    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
441        self.read_buf.read(buf)
442    }
443
444    #[cfg(feature = "read_buf")]
445    fn read_buf(&mut self, buf: io::BorrowedCursor<'_, u8>) -> io::Result<()> {
446        self.read_buf.read_buf(buf)
447    }
448}
449
450impl<S> BufRead for SyncStreamReadHalf<S> {
451    fn fill_buf(&mut self) -> io::Result<&[u8]> {
452        self.read_buf.fill_buf()
453    }
454
455    fn consume(&mut self, amt: usize) {
456        self.read_buf.consume(amt);
457    }
458}
459
460impl<S> Write for SyncStream<S> {
461    /// Writes data to the internal buffer.
462    ///
463    /// Returns `WouldBlock` if the buffer needs flushing or has reached max
464    /// capacity. In the latter case, it may write partial data before
465    /// returning `WouldBlock`.
466    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
467        self.write_buf.write(buf)
468    }
469
470    /// Returns `Ok(())` without checking for buffered data.
471    ///
472    /// **Important**: This does NOT actually flush data to the underlying
473    /// stream. This behavior is intentional for compatibility with
474    /// libraries like tungstenite that call `flush()` after every write
475    /// operation. The actual async flush happens when `flush_write_buf()`
476    /// is called.
477    ///
478    /// This prevents spurious errors in sync code that expects `flush()` to
479    /// succeed after successfully buffering data.
480    fn flush(&mut self) -> io::Result<()> {
481        Ok(())
482    }
483}
484
485impl<S> Write for SyncStreamWriteHalf<S> {
486    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
487        self.write_buf.write(buf)
488    }
489
490    fn flush(&mut self) -> io::Result<()> {
491        Ok(())
492    }
493}
494
495fn would_block(msg: &str) -> io::Error {
496    io::Error::new(io::ErrorKind::WouldBlock, msg)
497}
498
499impl<S: crate::AsyncRead> SyncStream<S> {
500    /// Fills the read buffer by reading from the underlying async stream.
501    ///
502    /// This method:
503    /// 1. Compacts the buffer if there's unconsumed data
504    /// 2. Ensures there's space for at least `base_capacity` more bytes
505    /// 3. Reads data from the underlying stream
506    /// 4. Returns the number of bytes read (0 indicates EOF)
507    ///
508    /// # Errors
509    ///
510    /// Returns an error if:
511    /// - The read buffer has reached `max_buffer_size`
512    /// - The underlying stream returns an error
513    pub async fn fill_read_buf(&mut self) -> io::Result<usize> {
514        self.read_buf.fill_read_buf(&mut self.inner).await
515    }
516}
517
518impl<S: crate::AsyncRead> SyncStreamReadHalf<S> {
519    /// See [`SyncStream::fill_read_buf`].
520    pub async fn fill_read_buf(&mut self) -> io::Result<usize> {
521        self.read_buf.fill_read_buf(&mut self.inner).await
522    }
523}
524
525impl<S: crate::AsyncWrite> SyncStream<S> {
526    /// Flushes the write buffer to the underlying async stream.
527    ///
528    /// This method:
529    /// 1. Writes all buffered data to the underlying stream
530    /// 2. Calls `flush()` on the underlying stream
531    /// 3. Returns the total number of bytes flushed
532    ///
533    /// On error, any unwritten data remains in the buffer and can be retried.
534    ///
535    /// # Errors
536    ///
537    /// Returns an error if the underlying stream returns an error.
538    /// In this case, the buffer retains any data that wasn't successfully
539    /// written.
540    pub async fn flush_write_buf(&mut self) -> io::Result<usize> {
541        self.write_buf.flush_write_buf(&mut self.inner).await
542    }
543}
544
545impl<S: crate::AsyncWrite> SyncStreamWriteHalf<S> {
546    /// See [`SyncStream::flush_write_buf`].
547    pub async fn flush_write_buf(&mut self) -> io::Result<usize> {
548        self.write_buf.flush_write_buf(&mut self.inner).await
549    }
550}
551
552impl<S: Splittable> Splittable for SyncStream<S> {
553    type ReadHalf = SyncStreamReadHalf<S::ReadHalf>;
554    type WriteHalf = SyncStreamWriteHalf<S::WriteHalf>;
555
556    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
557        let (r, w) = self.inner.split();
558        let read_half = SyncStreamReadHalf {
559            inner: r,
560            read_buf: self.read_buf,
561        };
562        let write_half = SyncStreamWriteHalf {
563            inner: w,
564            write_buf: self.write_buf,
565        };
566        (read_half, write_half)
567    }
568}
compio_io\compat/sync_stream.rs

compio_io\compat/
sync_stream.rs
