Skip to main content

rama_core/stream/
bytes_freeze.rs

1//! [`BytesFreeze`] — adapter that freezes [`BytesMut`] decoder output to
2//! [`Bytes`] without disturbing the [`Sink`] side.
3//!
4//! Most byte-oriented codecs in `tokio-util` (`BytesCodec`,
5//! `LengthDelimitedCodec`, …) decode to [`BytesMut`] but encode [`Bytes`].
6//! That asymmetry surfaces whenever you want to bridge two such
7//! [`Stream`] + [`Sink`] pairs through
8//! [`super::StreamForwardService`] — the symmetric item-type bound `T`
9//! forces both sides to agree on one type for stream items *and* sink
10//! input. Wrapping each side in [`BytesFreeze`] aligns them on `T = Bytes`.
11
12use core::pin::Pin;
13use core::task::{Context, Poll};
14
15use bytes::{Bytes, BytesMut};
16use futures::{Sink, Stream};
17
18/// Wraps a duplex [`Stream`] + [`Sink`] that decodes to [`BytesMut`] and
19/// accepts [`Bytes`] in its sink, exposing a uniform `Stream<Bytes>` +
20/// `Sink<Bytes>` on top.
21///
22/// Cheap zero-copy: [`BytesMut::freeze`] converts in place. The sink path
23/// is a direct pass-through.
24///
25/// # Example
26///
27/// ```no_run
28/// use rama_core::stream::BytesFreeze;
29/// use rama_core::stream::codec::{Framed, LengthDelimitedCodec};
30/// use tokio::net::TcpStream;
31///
32/// # async fn _example(tcp: TcpStream) {
33/// let framed = Framed::new(tcp, LengthDelimitedCodec::builder()
34///     .length_field_type::<u16>()
35///     .new_codec());
36/// // `framed`'s Stream yields BytesMut, its Sink takes Bytes.
37/// // BytesFreeze unifies both onto Bytes.
38/// let aligned = BytesFreeze::new(framed);
39/// # let _ = aligned;
40/// # }
41/// ```
42#[derive(Debug)]
43pub struct BytesFreeze<S> {
44    inner: S,
45}
46
47impl<S> BytesFreeze<S> {
48    /// Wrap `inner` in a [`BytesFreeze`] adapter.
49    pub fn new(inner: S) -> Self {
50        Self { inner }
51    }
52
53    /// Borrow the wrapped duplex.
54    #[must_use]
55    pub fn get_ref(&self) -> &S {
56        &self.inner
57    }
58
59    /// Mutably borrow the wrapped duplex.
60    pub fn get_mut(&mut self) -> &mut S {
61        &mut self.inner
62    }
63
64    /// Unwrap and return the inner duplex.
65    pub fn into_inner(self) -> S {
66        self.inner
67    }
68}
69
70impl<S, E> Stream for BytesFreeze<S>
71where
72    S: Stream<Item = Result<BytesMut, E>> + Unpin,
73{
74    type Item = Result<Bytes, E>;
75
76    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
77        Pin::new(&mut self.inner)
78            .poll_next(cx)
79            .map(|opt| opt.map(|res| res.map(BytesMut::freeze)))
80    }
81}
82
83impl<S, E> Sink<Bytes> for BytesFreeze<S>
84where
85    S: Sink<Bytes, Error = E> + Unpin,
86{
87    type Error = E;
88
89    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
90        Pin::new(&mut self.inner).poll_ready(cx)
91    }
92
93    fn start_send(mut self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
94        Pin::new(&mut self.inner).start_send(item)
95    }
96
97    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
98        Pin::new(&mut self.inner).poll_flush(cx)
99    }
100
101    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
102        Pin::new(&mut self.inner).poll_close(cx)
103    }
104}