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}