codec: Add `Framed::with_capacity` (#2215)

2 files changed, 47 insertions, 0 deletions

diff --git a/tokio-util/src/codec/framed.rs b/tokio-util/src/codec/framed.rs
index c770212b..e42ed28c 100644
--- a/tokio-util/src/codec/framed.rs
+++ b/tokio-util/src/codec/framed.rs
@@ -97,6 +97,39 @@ where
             inner: framed_read2(framed_write2(Fuse { io: inner, codec })),
         }
     }
+
+    /// Provides a [`Stream`] and [`Sink`] interface for reading and writing to this
+    /// I/O object, using [`Decoder`] and [`Encoder`] to read and write the raw data,
+    /// with a specific read buffer initial capacity.
+    ///
+    /// Raw I/O objects work with byte sequences, but higher-level code usually
+    /// wants to batch these into meaningful chunks, called "frames". This
+    /// method layers framing on top of an I/O object, by using the codec
+    /// traits to handle encoding and decoding of messages frames. Note that
+    /// the incoming and outgoing frame types may be distinct.
+    ///
+    /// This function returns a *single* object that is both [`Stream`] and
+    /// [`Sink`]; grouping this into a single object is often useful for layering
+    /// things like gzip or TLS, which require both read and write access to the
+    /// underlying object.
+    ///
+    /// If you want to work more directly with the streams and sink, consider
+    /// calling [`split`] on the `Framed` returned by this method, which will
+    /// break them into separate objects, allowing them to interact more easily.
+    ///
+    /// [`Stream`]: tokio::stream::Stream
+    /// [`Sink`]: futures_sink::Sink
+    /// [`Decode`]: crate::codec::Decoder
+    /// [`Encoder`]: crate::codec::Encoder
+    /// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split
+    pub fn with_capacity(inner: T, codec: U, capacity: usize) -> Framed<T, U> {
+        Framed {
+            inner: framed_read2_with_buffer(
+                framed_write2(Fuse { io: inner, codec }),
+                BytesMut::with_capacity(capacity),
+            ),
+        }
+    }
 }
 
 impl<T, U> Framed<T, U> {
diff --git a/tokio-util/src/codec/framed_read.rs b/tokio-util/src/codec/framed_read.rs
index ab1ed92b..e7798c32 100644
--- a/tokio-util/src/codec/framed_read.rs
+++ b/tokio-util/src/codec/framed_read.rs
@@ -50,6 +50,20 @@ where
             }),
         }
     }
+
+    /// Creates a new `FramedRead` with the given `decoder` and a buffer of `capacity`
+    /// initial size.
+    pub fn with_capacity(inner: T, decoder: D, capacity: usize) -> FramedRead<T, D> {
+        FramedRead {
+            inner: framed_read2_with_buffer(
+                Fuse {
+                    io: inner,
+                    codec: decoder,
+                },
+                BytesMut::with_capacity(capacity),
+            ),
+        }
+    }
 }
 
 impl<T, D> FramedRead<T, D> {