├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── benches └── bench.rs ├── examples └── encoder.rs ├── rustfmt.toml └── src ├── chunked.rs ├── iter.rs ├── lib.rs ├── loosely.rs ├── strictly.rs └── tests.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | **/*.rs.bk 3 | Cargo.lock 4 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "chunked-bytes" 3 | version = "0.3.0" 4 | license = "MIT" 5 | authors = ["Mikhail Zabaluev "] 6 | edition = "2018" 7 | 8 | repository = "https://github.com/mzabaluev/chunked-bytes" 9 | description = """ 10 | A rope-like non-contiguous buffer for efficient data structure serialization 11 | and vectored output. 12 | """ 13 | readme = "README.md" 14 | keywords = ["buffers", "rope", "zero-copy", "io"] 15 | categories = ["network-programming", "data-structures"] 16 | 17 | [dependencies] 18 | bytes = "1.0" 19 | 20 | [dev-dependencies] 21 | futures = { version = "0.3", features = ["std"], default-features = false } 22 | generic-tests = "0.1.1" 23 | pin-project = "1.0" 24 | 25 | [dev-dependencies.tokio] 26 | version = "1.1" 27 | features = ["rt-multi-thread", "macros"] 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020 Mikhail Zabaluev 2 | 3 | Permission is hereby granted, free of charge, to any 4 | person obtaining a copy of this software and associated 5 | documentation files (the "Software"), to deal in the 6 | Software without restriction, including without 7 | limitation the rights to use, copy, modify, merge, 8 | publish, distribute, sublicense, and/or sell copies of 9 | the Software, and to permit persons to whom the Software 10 | is furnished to do so, subject to the following 11 | conditions: 12 | 13 | The above copyright notice and this permission notice 14 | shall be included in all copies or substantial portions 15 | of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF 18 | ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED 19 | TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 20 | PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT 21 | SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 22 | CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 23 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR 24 | IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 25 | DEALINGS IN THE SOFTWARE. 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Chunked Bytes 2 | 3 | This crate provides two variants of `ChunkedBytes`, a non-contiguous buffer 4 | for efficient data structure serialization and vectored output. 5 | 6 | In network programming, there is often a need to serialize data structures 7 | into a wire protocol representation and send the resulting sequence of bytes 8 | to an output object, such as a socket. For a variety of reasons, developers 9 | of protocol implementations prefer to serialize the data into 10 | an intermediate buffer in memory rather than deal with output objects directly 11 | in either synchronous or asynchronous form, or both. When a contiguous 12 | buffer like `Vec` is used for this, reallocations to fit a larger length of 13 | serialized data may adversely impact performance, while evaluating the required 14 | length to pre-allocate beforehand may be cumbersome or difficult in the 15 | context. The single contiguous buffer also forms a boundary for write requests, 16 | creating a need for copy-back schemes to avoid inefficiently fragmented writes 17 | of tail data. 18 | 19 | Another important use case is passing network data through. If some of the data 20 | is received into [`Bytes`](https://docs.rs/bytes) handles, it should be possible 21 | to inject the data into the output stream without extra copying. 22 | 23 | Enter `ChunkedBytes`, containers that can be used to coalesce data added as 24 | byte slices via the `BufMut` interface, as well as possible `Bytes` input, 25 | into a sequence of chunks suitable for implementations of the 26 | [`Write::write_vectored`][write_vectored] method. This design aims to deliver 27 | good performance regardless of the size of buffered data and with no need 28 | to pre-allocate sufficient capacity for it. 29 | 30 | [write_vectored]: https://doc.rust-lang.org/stable/std/io/trait.Write.html#method.write_vectored 31 | 32 | ## License 33 | 34 | This project is licensed under the [MIT license](LICENSE). 35 | 36 | ### Contribution 37 | 38 | Unless you explicitly state otherwise, any contribution intentionally submitted 39 | for inclusion in `chunked-bytes` by you, shall be licensed as MIT, without any 40 | additional terms or conditions. 41 | -------------------------------------------------------------------------------- /benches/bench.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | #![feature(maybe_uninit_slice)] 3 | #![feature(write_all_vectored)] 4 | 5 | extern crate test; 6 | 7 | use bytes::{Buf, BufMut, Bytes, BytesMut}; 8 | use chunked_bytes::{loosely, strictly}; 9 | 10 | use std::cmp::min; 11 | use std::io::{self, IoSlice, Write}; 12 | use std::ptr; 13 | use test::Bencher; 14 | 15 | /// Imitates default TCP socket buffer size on Linux 16 | const BUF_SIZE: usize = 16 * 1024; 17 | 18 | trait BenchBuf: Buf + BufMut { 19 | fn construct() -> Self; 20 | fn construct_with_profile(chunk_size: usize, cnt: usize) -> Self; 21 | fn put_bytes(&mut self, bytes: Bytes); 22 | 23 | fn produce(&mut self, mut cnt: usize) { 24 | while cnt != 0 { 25 | let dst = self.chunk_mut(); 26 | let write_len = min(cnt, dst.len()); 27 | unsafe { 28 | ptr::write_bytes(dst.as_mut_ptr(), 0, write_len); 29 | self.advance_mut(write_len); 30 | } 31 | cnt -= write_len; 32 | } 33 | } 34 | 35 | fn consume_vectored(&mut self, mut cnt: usize) { 36 | // Do what TcpStream does 37 | loop { 38 | let mut slices = [IoSlice::new(&[]); 64]; 39 | let n = self.chunks_vectored(&mut slices); 40 | if n == 0 { 41 | break; 42 | } 43 | let mut sink = io::sink(); 44 | let total_len = sink.write_vectored(&mut slices[..n]).unwrap(); 45 | if cnt <= total_len { 46 | self.advance(cnt); 47 | break; 48 | } else { 49 | self.advance(total_len); 50 | cnt -= total_len; 51 | } 52 | } 53 | } 54 | } 55 | 56 | #[allow(dead_code)] 57 | fn write_vectored_into_black_box(buf: &mut B, mut cnt: usize) { 58 | // Do what TcpStream does 59 | loop { 60 | let mut slices = [IoSlice::new(&[]); 64]; 61 | let n = buf.chunks_vectored(&mut slices); 62 | if n == 0 { 63 | break; 64 | } 65 | let mut total_len = 0; 66 | for i in 0..n { 67 | let slice = &*slices[i]; 68 | total_len += slice.len(); 69 | // This has heavy impact on the benchmarks. Vectored output calls 70 | // to the OS would not make that big of a difference. 71 | for p in slice { 72 | test::black_box(*p); 73 | } 74 | } 75 | if cnt <= total_len { 76 | buf.advance(cnt); 77 | break; 78 | } else { 79 | buf.advance(total_len); 80 | cnt -= total_len; 81 | } 82 | } 83 | } 84 | 85 | impl BenchBuf for loosely::ChunkedBytes { 86 | fn construct() -> Self { 87 | Self::with_chunk_size_hint(BUF_SIZE) 88 | } 89 | 90 | fn construct_with_profile(chunk_size: usize, cnt: usize) -> Self { 91 | Self::with_profile(chunk_size, cnt) 92 | } 93 | 94 | fn put_bytes(&mut self, bytes: Bytes) { 95 | self.put_bytes(bytes) 96 | } 97 | } 98 | 99 | impl BenchBuf for strictly::ChunkedBytes { 100 | fn construct() -> Self { 101 | Self::with_chunk_size_limit(BUF_SIZE) 102 | } 103 | 104 | fn construct_with_profile(chunk_size: usize, cnt: usize) -> Self { 105 | Self::with_profile(chunk_size, cnt) 106 | } 107 | 108 | fn put_bytes(&mut self, bytes: Bytes) { 109 | self.put_bytes(bytes) 110 | } 111 | } 112 | 113 | impl BenchBuf for BytesMut { 114 | fn construct() -> Self { 115 | BytesMut::with_capacity(BUF_SIZE) 116 | } 117 | 118 | fn construct_with_profile(chunk_size: usize, cnt: usize) -> Self { 119 | BytesMut::with_capacity(chunk_size * cnt) 120 | } 121 | 122 | fn put_bytes(&mut self, bytes: Bytes) { 123 | self.put(bytes) 124 | } 125 | } 126 | 127 | #[generic_tests::define] 128 | mod benches { 129 | use super::*; 130 | 131 | #[bench] 132 | fn clean_pass_through(b: &mut Bencher) { 133 | let mut buf = B::construct(); 134 | let prealloc_cap = buf.chunk_mut().len(); 135 | b.iter(|| { 136 | buf.produce(prealloc_cap); 137 | buf.consume_vectored(prealloc_cap); 138 | }); 139 | } 140 | 141 | fn pump_through_staggered(b: &mut Bencher, carry_over: usize) { 142 | let mut buf = B::construct(); 143 | let prealloc_cap = buf.chunk_mut().len(); 144 | buf.produce(prealloc_cap); 145 | b.iter(|| { 146 | buf.consume_vectored(prealloc_cap - carry_over); 147 | buf.produce(prealloc_cap - carry_over); 148 | }); 149 | } 150 | 151 | #[bench] 152 | fn staggered_copy_back(b: &mut Bencher) { 153 | pump_through_staggered::(b, BUF_SIZE * 2 / 3); 154 | } 155 | 156 | #[bench] 157 | fn staggered_new_alloc(b: &mut Bencher) { 158 | pump_through_staggered::(b, (BUF_SIZE * 2 + 2) / 3 + 1); 159 | } 160 | 161 | fn pump_pressured( 162 | inflow: usize, 163 | outflow: usize, 164 | b: &mut Bencher, 165 | ) { 166 | let mut buf = B::construct(); 167 | let prealloc_cap = buf.chunk_mut().len(); 168 | b.iter(|| { 169 | buf.produce(inflow); 170 | while buf.remaining() >= prealloc_cap { 171 | buf.consume_vectored(outflow); 172 | } 173 | }); 174 | } 175 | 176 | #[bench] 177 | fn pressured_in_50_out_50_percent(b: &mut Bencher) { 178 | pump_pressured::(BUF_SIZE / 2, BUF_SIZE / 2, b); 179 | } 180 | 181 | #[bench] 182 | fn pressured_in_300_out_50_percent(b: &mut Bencher) { 183 | pump_pressured::(BUF_SIZE * 3, BUF_SIZE / 2, b); 184 | } 185 | 186 | #[bench] 187 | fn pressured_in_310_out_50_percent(b: &mut Bencher) { 188 | pump_pressured::(BUF_SIZE * 3 + BUF_SIZE / 10, BUF_SIZE / 2, b); 189 | } 190 | 191 | #[bench] 192 | fn pressured_in_350_out_50_percent(b: &mut Bencher) { 193 | pump_pressured::(BUF_SIZE * 3 + BUF_SIZE / 2, BUF_SIZE / 2, b); 194 | } 195 | 196 | #[bench] 197 | fn pressured_in_900_out_50_percent(b: &mut Bencher) { 198 | pump_pressured::(BUF_SIZE * 9, BUF_SIZE / 2, b); 199 | } 200 | 201 | #[bench] 202 | fn pressured_in_150_out_100_percent(b: &mut Bencher) { 203 | pump_pressured::(BUF_SIZE + BUF_SIZE / 2, BUF_SIZE, b); 204 | } 205 | 206 | #[bench] 207 | fn pressured_in_200_out_100_percent(b: &mut Bencher) { 208 | pump_pressured::(BUF_SIZE * 2, BUF_SIZE, b); 209 | } 210 | 211 | #[bench] 212 | fn pressured_in_210_out_100_percent(b: &mut Bencher) { 213 | pump_pressured::(BUF_SIZE * 2 + BUF_SIZE / 10, BUF_SIZE, b); 214 | } 215 | 216 | #[bench] 217 | fn pressured_in_300_out_100_percent(b: &mut Bencher) { 218 | pump_pressured::(BUF_SIZE * 3, BUF_SIZE, b); 219 | } 220 | 221 | #[bench] 222 | fn pressured_in_900_out_100_percent(b: &mut Bencher) { 223 | pump_pressured::(BUF_SIZE * 9, BUF_SIZE, b); 224 | } 225 | 226 | fn pass_bytes_through( 227 | b: &mut Bencher, 228 | chunk_size: usize, 229 | cnt: usize, 230 | ) { 231 | let mut buf = B::construct_with_profile(chunk_size, cnt); 232 | b.iter(|| { 233 | let mut salami = Bytes::from(vec![0; chunk_size * cnt]); 234 | for _ in 0..cnt { 235 | buf.put_bytes(salami.split_to(chunk_size)); 236 | } 237 | while buf.has_remaining() { 238 | buf.consume_vectored(BUF_SIZE); 239 | } 240 | }); 241 | } 242 | 243 | #[bench] 244 | fn pass_bytes_through_sliced_by_16(b: &mut Bencher) { 245 | pass_bytes_through::(b, BUF_SIZE / 16, 16); 246 | } 247 | 248 | #[bench] 249 | fn pass_bytes_through_sliced_by_64(b: &mut Bencher) { 250 | pass_bytes_through::(b, BUF_SIZE / 64, 64); 251 | } 252 | 253 | #[bench] 254 | fn pass_bytes_through_sliced_by_256(b: &mut Bencher) { 255 | pass_bytes_through::(b, BUF_SIZE / 256, 256); 256 | } 257 | 258 | #[bench] 259 | fn pass_bytes_through_medium_sized(b: &mut Bencher) { 260 | pass_bytes_through::(b, BUF_SIZE / 4, 16); 261 | } 262 | 263 | #[bench] 264 | fn pass_bytes_through_larger_than_buf(b: &mut Bencher) { 265 | pass_bytes_through::(b, BUF_SIZE * 2, 2); 266 | } 267 | 268 | fn mix_slice_and_bytes( 269 | b: &mut Bencher, 270 | slice_len: usize, 271 | bytes_len: usize, 272 | ) { 273 | let mut buf = B::construct(); 274 | let v = vec![0; slice_len]; 275 | b.iter(|| { 276 | buf.put_slice(&v); 277 | buf.put_bytes(Bytes::from(vec![0; bytes_len])); 278 | while buf.has_remaining() { 279 | buf.consume_vectored(BUF_SIZE); 280 | } 281 | }); 282 | } 283 | 284 | #[bench] 285 | fn mix_slice_and_bytes_32_32(b: &mut Bencher) { 286 | mix_slice_and_bytes::(b, 32, 32) 287 | } 288 | 289 | #[bench] 290 | fn mix_slice_and_bytes_32_4096(b: &mut Bencher) { 291 | mix_slice_and_bytes::(b, 32, 4096) 292 | } 293 | 294 | #[instantiate_tests()] 295 | mod loosely_chunked_bytes {} 296 | 297 | #[instantiate_tests()] 298 | mod strictly_chunked_bytes {} 299 | 300 | #[instantiate_tests()] 301 | mod bytes_mut {} 302 | } 303 | -------------------------------------------------------------------------------- /examples/encoder.rs: -------------------------------------------------------------------------------- 1 | // This example demonstrates the primary use case for `ChunkedBytes`: 2 | // a buffer to serialize data of protocol messages without reallocations 3 | // and write it to output with efficiency of `AsyncWrite` implementations 4 | // that make use of `Buf::chunks_vectored`. 5 | 6 | use bytes::{Buf, BufMut, Bytes, BytesMut}; 7 | use chunked_bytes::ChunkedBytes; 8 | use futures::io::{self, AsyncWrite, Error, IoSlice}; 9 | use futures::prelude::*; 10 | use futures::ready; 11 | use futures::Sink; 12 | use pin_project::pin_project; 13 | 14 | use std::iter::{self, FromIterator}; 15 | use std::pin::Pin; 16 | use std::task::{Context, Poll}; 17 | 18 | const MESSAGE_MAGIC: &[u8] = b"mess"; 19 | 20 | pub struct Message { 21 | int_field: u32, 22 | string_field: String, 23 | blob_field: Bytes, 24 | } 25 | 26 | fn encode_message(buf: &mut ChunkedBytes, msg: Message) { 27 | buf.put_slice(MESSAGE_MAGIC); 28 | 29 | buf.put_u32(1); 30 | buf.put_u32(msg.int_field); 31 | 32 | buf.put_u32(2); 33 | buf.put_u64(msg.string_field.len() as u64); 34 | buf.put(msg.string_field.as_bytes()); 35 | 36 | buf.put_u32(3); 37 | buf.put_u64(msg.blob_field.len() as u64); 38 | // Take the bytes without copying 39 | buf.put_bytes(msg.blob_field); 40 | } 41 | 42 | #[pin_project] 43 | pub struct EncodingWriter { 44 | #[pin] 45 | out: T, 46 | buf: ChunkedBytes, 47 | } 48 | 49 | impl EncodingWriter { 50 | pub fn new(out: T) -> Self { 51 | EncodingWriter { 52 | buf: ChunkedBytes::new(), 53 | out, 54 | } 55 | } 56 | 57 | fn poll_write_buf( 58 | self: Pin<&mut Self>, 59 | cx: &mut Context<'_>, 60 | ) -> Poll> { 61 | let mut io_bufs = [IoSlice::new(&[]); 16]; 62 | let mut this = self.project(); 63 | let io_vec_len = this.buf.chunks_vectored(&mut io_bufs); 64 | let bytes_written = ready!(this 65 | .out 66 | .as_mut() 67 | .poll_write_vectored(cx, &io_bufs[..io_vec_len]))?; 68 | this.buf.advance(bytes_written); 69 | Poll::Ready(Ok(())) 70 | } 71 | 72 | fn poll_flush_buf( 73 | mut self: Pin<&mut Self>, 74 | cx: &mut Context<'_>, 75 | ) -> Poll> { 76 | while self.as_mut().project().buf.has_remaining() { 77 | ready!(self.as_mut().poll_write_buf(cx))?; 78 | } 79 | Poll::Ready(Ok(())) 80 | } 81 | } 82 | 83 | impl Sink for EncodingWriter { 84 | type Error = Error; 85 | 86 | fn poll_ready( 87 | mut self: Pin<&mut Self>, 88 | cx: &mut Context<'_>, 89 | ) -> Poll> { 90 | // Here's a way to provide back-pressure on the sink: 91 | // rather than allowing the buffer to grow, drain it until 92 | // there is only the staging buffer with room to fill. 93 | let chunk_size = self.buf.chunk_size_hint(); 94 | while self.as_mut().buf.remaining() >= chunk_size { 95 | ready!(self.as_mut().poll_write_buf(cx))?; 96 | } 97 | Poll::Ready(Ok(())) 98 | } 99 | 100 | fn start_send(self: Pin<&mut Self>, msg: Message) -> Result<(), Error> { 101 | let this = self.project(); 102 | encode_message(this.buf, msg); 103 | Ok(()) 104 | } 105 | 106 | fn poll_flush( 107 | mut self: Pin<&mut Self>, 108 | cx: &mut Context<'_>, 109 | ) -> Poll> { 110 | ready!(self.as_mut().poll_flush_buf(cx))?; 111 | self.project().out.poll_flush(cx) 112 | } 113 | 114 | fn poll_close( 115 | mut self: Pin<&mut Self>, 116 | cx: &mut Context<'_>, 117 | ) -> Poll> { 118 | ready!(self.as_mut().poll_flush_buf(cx))?; 119 | self.project().out.poll_close(cx) 120 | } 121 | } 122 | 123 | #[tokio::main] 124 | async fn main() -> io::Result<()> { 125 | // Pretend we received the data from input into a Bytes handle 126 | let blob = BytesMut::from_iter(iter::repeat(b'\xa5').take(8000)); 127 | 128 | let msg = Message { 129 | int_field: 42, 130 | string_field: "Hello, world!".into(), 131 | blob_field: blob.freeze(), 132 | }; 133 | 134 | let sink = io::sink(); 135 | let mut writer = EncodingWriter::new(sink); 136 | 137 | writer.send(msg).await?; 138 | Ok(()) 139 | } 140 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | max_width = 80 2 | -------------------------------------------------------------------------------- /src/chunked.rs: -------------------------------------------------------------------------------- 1 | use crate::{DrainChunks, IntoChunks}; 2 | 3 | use bytes::buf::{Buf, BufMut, UninitSlice}; 4 | use bytes::{Bytes, BytesMut}; 5 | 6 | use std::cmp::min; 7 | use std::collections::VecDeque; 8 | use std::io::IoSlice; 9 | 10 | const DEFAULT_CHUNK_SIZE: usize = 4096; 11 | 12 | #[derive(Debug)] 13 | pub(crate) struct Inner { 14 | staging: BytesMut, 15 | chunks: VecDeque, 16 | chunk_size: usize, 17 | } 18 | 19 | impl Default for Inner { 20 | #[inline] 21 | fn default() -> Self { 22 | Inner { 23 | staging: BytesMut::new(), 24 | chunks: VecDeque::new(), 25 | chunk_size: DEFAULT_CHUNK_SIZE, 26 | } 27 | } 28 | } 29 | 30 | pub(crate) enum AdvanceStopped { 31 | InChunk, 32 | InStaging(usize), 33 | } 34 | 35 | impl Inner { 36 | #[inline] 37 | pub fn with_chunk_size(chunk_size: usize) -> Self { 38 | Inner { 39 | chunk_size, 40 | ..Default::default() 41 | } 42 | } 43 | 44 | #[inline] 45 | pub fn with_profile(chunk_size: usize, chunking_capacity: usize) -> Self { 46 | Inner { 47 | staging: BytesMut::new(), 48 | chunks: VecDeque::with_capacity(chunking_capacity), 49 | chunk_size, 50 | } 51 | } 52 | 53 | #[inline] 54 | pub fn chunk_size(&self) -> usize { 55 | self.chunk_size 56 | } 57 | 58 | #[inline] 59 | pub fn is_empty(&self) -> bool { 60 | self.chunks.is_empty() && self.staging.is_empty() 61 | } 62 | 63 | #[inline] 64 | pub fn staging_len(&self) -> usize { 65 | self.staging.len() 66 | } 67 | 68 | #[inline] 69 | pub fn staging_capacity(&self) -> usize { 70 | self.staging.capacity() 71 | } 72 | 73 | #[inline] 74 | pub fn push_chunk(&mut self, chunk: Bytes) { 75 | debug_assert!(!chunk.is_empty()); 76 | self.chunks.push_back(chunk) 77 | } 78 | 79 | #[inline] 80 | pub fn flush(&mut self) { 81 | if !self.staging.is_empty() { 82 | let bytes = self.staging.split().freeze(); 83 | self.push_chunk(bytes) 84 | } 85 | } 86 | 87 | #[inline] 88 | pub fn drain_chunks(&mut self) -> DrainChunks<'_> { 89 | DrainChunks::new(self.chunks.drain(..)) 90 | } 91 | 92 | #[inline] 93 | pub fn into_chunks(mut self) -> IntoChunks { 94 | if !self.staging.is_empty() { 95 | self.chunks.push_back(self.staging.freeze()); 96 | } 97 | IntoChunks::new(self.chunks.into_iter()) 98 | } 99 | 100 | pub fn reserve_staging(&mut self) -> usize { 101 | let cap = self.staging.capacity(); 102 | 103 | // We are here when either: 104 | // a) the buffer has never been used and never allocated; 105 | // b) the producer has filled a previously allocated buffer, 106 | // and the consumer may have read a part or the whole of it. 107 | // Our goal is to reserve space in the staging buffer without 108 | // forcing it to reallocate to a larger capacity. 109 | // 110 | // To reuse the allocation of `BytesMut` in the vector form with 111 | // the offset `off` and remaining capacity `cap` while reserving 112 | // `additional` bytes, the following needs to apply: 113 | // 114 | // off >= additional && off >= cap / 2 115 | // 116 | // We have: 117 | // 118 | // off + cap == allocated_size >= chunk_size 119 | // 120 | // From this, we can derive the following condition check: 121 | let cutoff = cap.saturating_add(cap / 2); 122 | let additional = if cutoff > self.chunk_size { 123 | // Alas, the bytes still in the staging buffer are likely to 124 | // necessitate a new allocation. Split them off to a chunk 125 | // first, so that the new allocation does not have to copy 126 | // them and the total required capacity is `self.chunk_size`. 127 | self.flush(); 128 | self.chunk_size 129 | } else { 130 | // This amount will get BytesMut to reuse the allocation and 131 | // copy back the bytes if there are no chunks left unconsumed. 132 | // Otherwise, it will reallocate to its previous capacity. 133 | // A virgin buffer will be allocated to `self.chunk_size`. 134 | self.chunk_size - cap 135 | }; 136 | self.staging.reserve(additional); 137 | self.staging.capacity() 138 | } 139 | 140 | #[inline] 141 | pub fn remaining_mut(&self) -> usize { 142 | self.staging.remaining_mut() 143 | } 144 | 145 | #[inline] 146 | pub unsafe fn advance_mut(&mut self, cnt: usize) { 147 | self.staging.advance_mut(cnt); 148 | } 149 | 150 | #[inline] 151 | pub fn chunk_mut(&mut self) -> &mut UninitSlice { 152 | self.staging.chunk_mut() 153 | } 154 | 155 | pub fn remaining(&self) -> usize { 156 | self.chunks 157 | .iter() 158 | .fold(self.staging.len(), |sum, chunk| sum + chunk.len()) 159 | } 160 | 161 | #[inline] 162 | pub fn chunk(&self) -> &[u8] { 163 | if let Some(chunk) = self.chunks.front() { 164 | chunk 165 | } else { 166 | self.staging.chunk() 167 | } 168 | } 169 | 170 | pub fn advance(&mut self, mut cnt: usize) -> AdvanceStopped { 171 | loop { 172 | match self.chunks.front_mut() { 173 | None => { 174 | self.staging.advance(cnt); 175 | return AdvanceStopped::InStaging(cnt); 176 | } 177 | Some(chunk) => { 178 | let len = chunk.len(); 179 | if cnt < len { 180 | chunk.advance(cnt); 181 | return AdvanceStopped::InChunk; 182 | } else { 183 | cnt -= len; 184 | self.chunks.pop_front(); 185 | } 186 | } 187 | } 188 | } 189 | } 190 | 191 | pub fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { 192 | let n = { 193 | let zipped = dst.iter_mut().zip(self.chunks.iter()); 194 | let len = zipped.len(); 195 | for (io_slice, chunk) in zipped { 196 | *io_slice = IoSlice::new(chunk); 197 | } 198 | len 199 | }; 200 | 201 | if n < dst.len() && !self.staging.is_empty() { 202 | dst[n] = IoSlice::new(&self.staging); 203 | n + 1 204 | } else { 205 | n 206 | } 207 | } 208 | 209 | pub fn copy_to_bytes(&mut self, len: usize) -> Bytes { 210 | if self.chunks.is_empty() { 211 | return self.staging.copy_to_bytes(len); 212 | } 213 | let mut to_copy = min(len, self.remaining()); 214 | let mut buf = BytesMut::with_capacity(to_copy); 215 | loop { 216 | match self.chunks.front_mut() { 217 | None => { 218 | buf.put((&mut self.staging).take(to_copy)); 219 | break; 220 | } 221 | Some(chunk) => { 222 | if chunk.len() > to_copy { 223 | buf.put(chunk.take(to_copy)); 224 | break; 225 | } else { 226 | buf.extend_from_slice(chunk); 227 | to_copy -= chunk.len(); 228 | } 229 | } 230 | } 231 | self.chunks.pop_front(); 232 | } 233 | buf.freeze() 234 | } 235 | } 236 | -------------------------------------------------------------------------------- /src/iter.rs: -------------------------------------------------------------------------------- 1 | use bytes::Bytes; 2 | 3 | use std::collections::vec_deque; 4 | use std::iter::FusedIterator; 5 | 6 | /// The iterator produced by the `drain_chunks` method of `ChunkedBytes`. 7 | pub struct DrainChunks<'a> { 8 | inner: vec_deque::Drain<'a, Bytes>, 9 | } 10 | 11 | impl<'a> DrainChunks<'a> { 12 | #[inline] 13 | pub(crate) fn new(inner: vec_deque::Drain<'a, Bytes>) -> Self { 14 | DrainChunks { inner } 15 | } 16 | } 17 | 18 | impl<'a> Iterator for DrainChunks<'a> { 19 | type Item = Bytes; 20 | 21 | #[inline] 22 | fn next(&mut self) -> Option { 23 | self.inner.next() 24 | } 25 | 26 | #[inline] 27 | fn size_hint(&self) -> (usize, Option) { 28 | self.inner.size_hint() 29 | } 30 | } 31 | 32 | impl<'a> ExactSizeIterator for DrainChunks<'a> {} 33 | impl<'a> FusedIterator for DrainChunks<'a> {} 34 | 35 | /// The iterator produced by the `into_chunks` method of `ChunkedBytes`. 36 | pub struct IntoChunks { 37 | inner: vec_deque::IntoIter, 38 | } 39 | 40 | impl IntoChunks { 41 | #[inline] 42 | pub(crate) fn new(inner: vec_deque::IntoIter) -> Self { 43 | IntoChunks { inner } 44 | } 45 | } 46 | 47 | impl Iterator for IntoChunks { 48 | type Item = Bytes; 49 | 50 | #[inline] 51 | fn next(&mut self) -> Option { 52 | self.inner.next() 53 | } 54 | 55 | #[inline] 56 | fn size_hint(&self) -> (usize, Option) { 57 | self.inner.size_hint() 58 | } 59 | } 60 | 61 | impl ExactSizeIterator for IntoChunks {} 62 | impl FusedIterator for IntoChunks {} 63 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A non-contiguous buffer for efficient serialization of data structures 2 | //! and vectored output. 3 | //! 4 | //! This crate provides `ChunkedBytes`, a [rope]-like byte container based on 5 | //! `Bytes` and `BytesMut` from the `bytes` crate. Its primary purpose is to 6 | //! serve as an intermediate buffer for serializing fields of data structures 7 | //! into byte sequences of varying length, without whole-buffer reallocations 8 | //! like those performed to grow a `Vec`, and then consuming the bytes in bulk, 9 | //! split into regularly sized chunks suitable for [vectored output]. 10 | //! 11 | //! [rope]: https://en.wikipedia.org/wiki/Rope_(data_structure) 12 | //! [vectored output]: https://en.wikipedia.org/wiki/Vectored_I/O 13 | //! 14 | //! `ChunkedBytes` implements the traits `Buf` and `BufMut` for read and write 15 | //! access to the buffered data. It also provides the `put_bytes` method 16 | //! for appending a `Bytes` slice to its queue of non-contiguous chunks without 17 | //! copying the data. 18 | //! 19 | //! # Examples 20 | //! 21 | //! ``` 22 | //! use bytes::{Buf, BufMut, Bytes}; 23 | //! use chunked_bytes::ChunkedBytes; 24 | //! use std::io::{self, IoSlice, Read, Write}; 25 | //! use std::net::{SocketAddr, TcpListener, TcpStream, Shutdown}; 26 | //! use std::thread; 27 | //! 28 | //! fn write_vectored( 29 | //! buf: &mut ChunkedBytes, 30 | //! mut out: W, 31 | //! ) -> io::Result { 32 | //! let mut io_bufs = [IoSlice::new(&[]); 32]; 33 | //! let io_vec_len = buf.chunks_vectored(&mut io_bufs); 34 | //! let bytes_written = out.write_vectored(&io_bufs[..io_vec_len])?; 35 | //! buf.advance(bytes_written); 36 | //! Ok(bytes_written) 37 | //! } 38 | //! 39 | //! fn main() -> io::Result<()> { 40 | //! const MESSAGE: &[u8] = b"I \xf0\x9f\x96\xa4 \x00\xc0\xff\xee"; 41 | //! 42 | //! let listen_addr = "127.0.0.1:0".parse::().unwrap(); 43 | //! let server = TcpListener::bind(listen_addr)?; 44 | //! let server_addr = server.local_addr()?; 45 | //! 46 | //! let server_handle: thread::JoinHandle> = 47 | //! thread::spawn(move || { 48 | //! let (mut receiver, _) = server.accept()?; 49 | //! let mut buf = Vec::with_capacity(64); 50 | //! receiver.read_to_end(&mut buf)?; 51 | //! assert_eq!(buf.as_slice(), MESSAGE); 52 | //! Ok(()) 53 | //! }); 54 | //! 55 | //! let mut sender = TcpStream::connect(server_addr)?; 56 | //! 57 | //! let mut buf = ChunkedBytes::with_chunk_size_hint(4096); 58 | //! 59 | //! buf.put("I ".as_bytes()); 60 | //! buf.put_bytes(Bytes::from("🖤 ")); 61 | //! buf.put_u32(0xc0ffee); 62 | //! 63 | //! let bytes_written = write_vectored(&mut buf, &mut sender)?; 64 | //! assert_eq!(bytes_written, MESSAGE.len()); 65 | //! 66 | //! sender.shutdown(Shutdown::Write)?; 67 | //! 68 | //! server_handle.join().expect("server thread panicked")?; 69 | //! Ok(()) 70 | //! } 71 | 72 | #![warn(clippy::all)] 73 | #![warn(future_incompatible)] 74 | #![warn(missing_docs)] 75 | #![warn(rust_2018_idioms)] 76 | #![doc(test(no_crate_inject, attr(deny(warnings, rust_2018_idioms))))] 77 | 78 | pub mod loosely; 79 | pub mod strictly; 80 | 81 | mod chunked; 82 | mod iter; 83 | 84 | pub use self::iter::{DrainChunks, IntoChunks}; 85 | pub use self::loosely::ChunkedBytes; 86 | 87 | #[cfg(test)] 88 | mod tests; 89 | -------------------------------------------------------------------------------- /src/loosely.rs: -------------------------------------------------------------------------------- 1 | //! Buffer with a loose adherence to the preferred chunk size. 2 | 3 | use super::chunked::Inner; 4 | use crate::{DrainChunks, IntoChunks}; 5 | 6 | use bytes::buf::{Buf, BufMut, UninitSlice}; 7 | use bytes::Bytes; 8 | 9 | use std::fmt; 10 | use std::io::IoSlice; 11 | 12 | /// A non-contiguous buffer for efficient serialization of data structures. 13 | /// 14 | /// A `ChunkedBytes` container has a staging buffer to coalesce small byte 15 | /// sequences of source data, and a queue of byte chunks split off the staging 16 | /// buffer that can be incrementally consumed by an output API such as an object 17 | /// implementing `AsyncWrite`. Once the number of bytes in the staging 18 | /// buffer reaches a certain configured chunk size, the buffer content is 19 | /// split off to form a new chunk. 20 | /// 21 | /// This variant of the container does not enforce an upper limit on the size 22 | /// of contiguous chunks, being optimized for performance. If your application 23 | /// needs the sizes of the produced chunks to be capped, 24 | /// use `strictly::ChunkedBytes` instead. 25 | /// 26 | /// Refer to the documentation on the methods available for `ChunkedBytes`, 27 | /// including the methods of traits `Buf` and `BufMut`, for details on working 28 | /// with this container. 29 | #[derive(Debug, Default)] 30 | pub struct ChunkedBytes { 31 | inner: Inner, 32 | } 33 | 34 | impl ChunkedBytes { 35 | /// Creates a new `ChunkedBytes` container with the preferred chunk size 36 | /// set to a default value. 37 | #[inline] 38 | pub fn new() -> Self { 39 | Default::default() 40 | } 41 | 42 | /// Creates a new `ChunkedBytes` container with the given chunk size 43 | /// to prefer. 44 | #[inline] 45 | pub fn with_chunk_size_hint(chunk_size: usize) -> Self { 46 | ChunkedBytes { 47 | inner: Inner::with_chunk_size(chunk_size), 48 | } 49 | } 50 | 51 | /// The fully detailed constructor for `ChunkedBytes`. 52 | /// The preferred chunk size is given in `chunk_size`, and an upper 53 | /// estimate of the number of chunks this container could be expected to 54 | /// have at any moment of time should be given in `chunking_capacity`. 55 | /// More chunks can still be held, but this may cause reallocations of 56 | /// internal data structures. 57 | #[inline] 58 | pub fn with_profile(chunk_size: usize, chunking_capacity: usize) -> Self { 59 | ChunkedBytes { 60 | inner: Inner::with_profile(chunk_size, chunking_capacity), 61 | } 62 | } 63 | 64 | /// Returns the size this `ChunkedBytes` container uses as the threshold 65 | /// for splitting off complete chunks. 66 | /// 67 | /// Note that the size of produced chunks may be larger or smaller than the 68 | /// configured value, due to the allocation strategy used internally by 69 | /// the implementation and also depending on the pattern of usage. 70 | #[inline] 71 | pub fn chunk_size_hint(&self) -> usize { 72 | self.inner.chunk_size() 73 | } 74 | 75 | /// Returns true if the `ChunkedBytes` container has no complete chunks 76 | /// and the staging buffer is empty. 77 | #[inline] 78 | pub fn is_empty(&self) -> bool { 79 | self.inner.is_empty() 80 | } 81 | 82 | #[cfg(test)] 83 | pub fn staging_capacity(&self) -> usize { 84 | self.inner.staging_capacity() 85 | } 86 | 87 | /// Splits any bytes that are currently in the staging buffer into a new 88 | /// complete chunk. 89 | /// If the staging buffer is empty, this method does nothing. 90 | /// 91 | /// Most users should not need to call this method. It is called 92 | /// internally when needed by the methods that advance the writing 93 | /// position. 94 | #[inline] 95 | pub fn flush(&mut self) { 96 | self.inner.flush() 97 | } 98 | 99 | /// Appends a `Bytes` slice to the container without copying the data. 100 | /// 101 | /// If `chunk` is empty, this method does nothing. Otherwise, 102 | /// if there are any bytes currently in the staging buffer, they are split 103 | /// to form a complete chunk. Next, the given slice is appended as the 104 | /// next chunk. 105 | /// 106 | /// # Performance Notes 107 | /// 108 | /// For a small slice originating from a buffer that is not split 109 | /// or shared between other `Bytes` instances, copying the bytes with 110 | /// `BufMut::put_slice` may be faster than the overhead of 111 | /// atomic reference counting induced by use of this method. 112 | #[inline] 113 | pub fn put_bytes(&mut self, chunk: Bytes) { 114 | if !chunk.is_empty() { 115 | self.flush(); 116 | self.inner.push_chunk(chunk); 117 | } 118 | } 119 | 120 | /// Returns an iterator that removes complete chunks from the 121 | /// `ChunkedBytes` container and yields the removed chunks as `Bytes` 122 | /// slice handles. This does not include bytes in the staging buffer. 123 | /// 124 | /// The chunks are removed even if the iterator is dropped without being 125 | /// consumed until the end. It is unspecified how many chunks are removed 126 | /// if the `DrainChunks` value is not dropped, but the borrow it holds 127 | /// expires (e.g. due to `std::mem::forget`). 128 | #[inline] 129 | pub fn drain_chunks(&mut self) -> DrainChunks<'_> { 130 | self.inner.drain_chunks() 131 | } 132 | 133 | /// Consumes the `ChunkedBytes` container to produce an iterator over 134 | /// its chunks. If there are bytes in the staging buffer, they are yielded 135 | /// as the last chunk. 136 | /// 137 | /// The memory allocated for `IntoChunks` may be slightly more than the 138 | /// `ChunkedBytes` container it consumes. This is an infrequent side effect 139 | /// of making the internal state efficient in general for iteration. 140 | #[inline] 141 | pub fn into_chunks(self) -> IntoChunks { 142 | self.inner.into_chunks() 143 | } 144 | } 145 | 146 | unsafe impl BufMut for ChunkedBytes { 147 | #[inline] 148 | fn remaining_mut(&self) -> usize { 149 | self.inner.remaining_mut() 150 | } 151 | 152 | #[inline] 153 | unsafe fn advance_mut(&mut self, cnt: usize) { 154 | self.inner.advance_mut(cnt); 155 | } 156 | 157 | /// Returns a mutable slice of unwritten bytes available in 158 | /// the staging buffer, starting at the current writing position. 159 | /// 160 | /// The length of the slice may be larger than the preferred chunk 161 | /// size due to the allocation strategy used internally by 162 | /// the implementation. 163 | #[inline] 164 | fn chunk_mut(&mut self) -> &mut UninitSlice { 165 | if self.inner.staging_len() == self.inner.staging_capacity() { 166 | self.inner.reserve_staging(); 167 | } 168 | self.inner.chunk_mut() 169 | } 170 | } 171 | 172 | impl Buf for ChunkedBytes { 173 | #[inline] 174 | fn remaining(&self) -> usize { 175 | self.inner.remaining() 176 | } 177 | 178 | #[inline] 179 | fn has_remaining(&self) -> bool { 180 | !self.is_empty() 181 | } 182 | 183 | /// Returns a slice of the bytes in the first extant complete chunk, 184 | /// or the bytes in the staging buffer if there are no unconsumed chunks. 185 | /// 186 | /// It is more efficient to use `chunks_vectored` to gather all the disjoint 187 | /// slices for vectored output. 188 | #[inline] 189 | fn chunk(&self) -> &[u8] { 190 | self.inner.chunk() 191 | } 192 | 193 | /// Advances the reading position by `cnt`, dropping the `Bytes` references 194 | /// to any complete chunks that the position has been advanced past 195 | /// and then advancing the starting position of the first remaining chunk. 196 | /// If there are no complete chunks left, the reading position is advanced 197 | /// in the staging buffer, effectively removing the consumed bytes. 198 | /// 199 | /// # Panics 200 | /// 201 | /// This function may panic when `cnt > self.remaining()`. 202 | /// 203 | #[inline] 204 | fn advance(&mut self, cnt: usize) { 205 | let _ = self.inner.advance(cnt); 206 | } 207 | 208 | /// Fills `dst` sequentially with the slice views of the chunks, then 209 | /// the bytes in the staging buffer if any remain and there is 210 | /// another unfilled entry left in `dst`. Returns the number of `IoSlice` 211 | /// entries filled. 212 | #[inline] 213 | fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { 214 | self.inner.chunks_vectored(dst) 215 | } 216 | 217 | #[inline] 218 | fn copy_to_bytes(&mut self, len: usize) -> Bytes { 219 | self.inner.copy_to_bytes(len) 220 | } 221 | } 222 | 223 | impl fmt::Write for ChunkedBytes { 224 | #[inline] 225 | fn write_str(&mut self, s: &str) -> fmt::Result { 226 | if self.remaining_mut() >= s.len() { 227 | self.put_slice(s.as_bytes()); 228 | Ok(()) 229 | } else { 230 | Err(fmt::Error) 231 | } 232 | } 233 | 234 | // The default implementation delegates to 235 | // fmt::write(&mut self as &mut dyn fmt::Write, args) 236 | #[inline] 237 | fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result { 238 | fmt::write(self, args) 239 | } 240 | } 241 | -------------------------------------------------------------------------------- /src/strictly.rs: -------------------------------------------------------------------------------- 1 | //! Buffer with a strict limit on the chunk sizes. 2 | 3 | use super::chunked::{AdvanceStopped, Inner}; 4 | use crate::{DrainChunks, IntoChunks}; 5 | 6 | use bytes::buf::{Buf, BufMut, UninitSlice}; 7 | use bytes::Bytes; 8 | 9 | use std::cmp::min; 10 | use std::fmt; 11 | use std::io::IoSlice; 12 | 13 | /// A non-contiguous buffer for efficient serialization of data structures. 14 | /// 15 | /// A `ChunkedBytes` container has a staging buffer to coalesce small byte 16 | /// sequences of source data, and a queue of byte chunks split off the staging 17 | /// buffer that can be incrementally consumed by an output API such as an object 18 | /// implementing `AsyncWrite`. Once the number of bytes in the staging 19 | /// buffer reaches a certain configured chunk size, the buffer content is 20 | /// split off to form a new chunk. 21 | /// 22 | /// Unlike `loosely::ChunkedBytes`, this variant of the `ChunkedBytes` container 23 | /// never produces chunks larger than the configured size. This comes at a cost 24 | /// of increased processing overhead and sometimes more allocated memory needed 25 | /// to keep the buffered data, so the applications that don't benefit from 26 | /// the strict limit should prefer `loosely::ChunkedBytes`. 27 | /// 28 | /// Refer to the documentation on the methods available for `ChunkedBytes`, 29 | /// including the methods of traits `Buf` and `BufMut`, for details on working 30 | /// with this container. 31 | #[derive(Debug, Default)] 32 | pub struct ChunkedBytes { 33 | inner: Inner, 34 | // Maintains own capacity counter because `BytesMut` can't guarantee 35 | // the exact requested capacity. 36 | cap: usize, 37 | } 38 | 39 | impl ChunkedBytes { 40 | /// Creates a new `ChunkedBytes` container with the chunk size limit 41 | /// set to a default value. 42 | #[inline] 43 | pub fn new() -> Self { 44 | Default::default() 45 | } 46 | 47 | /// Creates a new `ChunkedBytes` container with the given chunk size limit. 48 | #[inline] 49 | pub fn with_chunk_size_limit(chunk_size: usize) -> Self { 50 | ChunkedBytes { 51 | inner: Inner::with_chunk_size(chunk_size), 52 | cap: 0, 53 | } 54 | } 55 | 56 | /// The fully detailed constructor for `ChunkedBytes`. 57 | /// The chunk size limit is given in `chunk_size`, and an upper 58 | /// estimate of the number of chunks this container could be expected to 59 | /// have at any moment of time should be given in `chunking_capacity`. 60 | /// More chunks can still be held, but this may cause reallocations of 61 | /// internal data structures. 62 | #[inline] 63 | pub fn with_profile(chunk_size: usize, chunking_capacity: usize) -> Self { 64 | ChunkedBytes { 65 | inner: Inner::with_profile(chunk_size, chunking_capacity), 66 | cap: 0, 67 | } 68 | } 69 | 70 | /// Returns the size this `ChunkedBytes` container uses as the limit 71 | /// for splitting off complete chunks. 72 | /// 73 | /// Note that the size of produced chunks may be smaller than the 74 | /// configured value, due to the allocation strategy used internally by 75 | /// the implementation and also depending on the pattern of usage. 76 | #[inline] 77 | pub fn chunk_size_limit(&self) -> usize { 78 | self.inner.chunk_size() 79 | } 80 | 81 | /// Returns true if the `ChunkedBytes` container has no complete chunks 82 | /// and the staging buffer is empty. 83 | #[inline] 84 | pub fn is_empty(&self) -> bool { 85 | self.inner.is_empty() 86 | } 87 | 88 | #[cfg(test)] 89 | pub fn staging_capacity(&self) -> usize { 90 | self.inner.staging_capacity() 91 | } 92 | 93 | /// Splits any bytes that are currently in the staging buffer into a new 94 | /// complete chunk. 95 | /// If the staging buffer is empty, this method does nothing. 96 | /// 97 | /// Most users should not need to call this method. It is called 98 | /// internally when needed by the methods that advance the writing 99 | /// position. 100 | #[inline] 101 | pub fn flush(&mut self) { 102 | debug_assert!(self.inner.staging_len() <= self.inner.chunk_size()); 103 | self.inner.flush() 104 | } 105 | 106 | /// Appends a `Bytes` slice to the container without copying the data. 107 | /// 108 | /// If `src` is empty, this method does nothing. Otherwise, 109 | /// if there are any bytes currently in the staging buffer, they are split 110 | /// to form a complete chunk. Next, `src` is appended as a sequence of 111 | /// chunks, split if necessary so that all chunks except the last are 112 | /// sized to the chunk size limit. 113 | /// 114 | /// # Performance Notes 115 | /// 116 | /// For a small slice originating from a buffer that is not split 117 | /// or shared between other `Bytes` instances, copying the bytes with 118 | /// `BufMut::put_slice` may be faster than the overhead of 119 | /// atomic reference counting induced by use of this method. 120 | pub fn put_bytes(&mut self, mut src: Bytes) { 121 | if !src.is_empty() { 122 | self.flush(); 123 | let chunk_size = self.inner.chunk_size(); 124 | while src.len() > chunk_size { 125 | self.inner.push_chunk(src.split_to(chunk_size)); 126 | } 127 | self.inner.push_chunk(src); 128 | } 129 | } 130 | 131 | /// Returns an iterator that removes complete chunks from the 132 | /// `ChunkedBytes` container and yields the removed chunks as `Bytes` 133 | /// slice handles. This does not include bytes in the staging buffer. 134 | /// 135 | /// The chunks are removed even if the iterator is dropped without being 136 | /// consumed until the end. It is unspecified how many chunks are removed 137 | /// if the `DrainChunks` value is not dropped, but the borrow it holds 138 | /// expires (e.g. due to `std::mem::forget`). 139 | #[inline] 140 | pub fn drain_chunks(&mut self) -> DrainChunks<'_> { 141 | self.inner.drain_chunks() 142 | } 143 | 144 | /// Consumes the `ChunkedBytes` container to produce an iterator over 145 | /// its chunks. If there are bytes in the staging buffer, they are yielded 146 | /// as the last src. 147 | /// 148 | /// The memory allocated for `IntoChunks` may be slightly more than the 149 | /// `ChunkedBytes` container it consumes. This is an infrequent side effect 150 | /// of making the internal state efficient in general for iteration. 151 | #[inline] 152 | pub fn into_chunks(self) -> IntoChunks { 153 | debug_assert!(self.inner.staging_len() <= self.inner.chunk_size()); 154 | self.inner.into_chunks() 155 | } 156 | } 157 | 158 | unsafe impl BufMut for ChunkedBytes { 159 | #[inline] 160 | fn remaining_mut(&self) -> usize { 161 | self.inner.remaining_mut() 162 | } 163 | 164 | #[inline] 165 | unsafe fn advance_mut(&mut self, cnt: usize) { 166 | assert!( 167 | self.inner.staging_len() + cnt <= self.cap, 168 | "new_len = {}; capacity = {}", 169 | self.inner.staging_len() + cnt, 170 | self.cap 171 | ); 172 | self.inner.advance_mut(cnt); 173 | } 174 | 175 | fn chunk_mut(&mut self) -> &mut UninitSlice { 176 | if self.inner.staging_len() == self.cap { 177 | let new_cap = self.inner.reserve_staging(); 178 | self.cap = min(new_cap, self.chunk_size_limit()) 179 | } 180 | let chunk = self.inner.chunk_mut(); 181 | let len = min(chunk.len(), self.cap); 182 | &mut chunk[..len] 183 | } 184 | } 185 | 186 | impl Buf for ChunkedBytes { 187 | #[inline] 188 | fn remaining(&self) -> usize { 189 | self.inner.remaining() 190 | } 191 | 192 | #[inline] 193 | fn has_remaining(&self) -> bool { 194 | !self.is_empty() 195 | } 196 | 197 | /// Returns a slice of the bytes in the first extant complete chunk, 198 | /// or the bytes in the staging buffer if there are no unconsumed chunks. 199 | /// 200 | /// It is more efficient to use `chunks_vectored` to gather all the disjoint 201 | /// slices for vectored output. 202 | #[inline] 203 | fn chunk(&self) -> &[u8] { 204 | self.inner.chunk() 205 | } 206 | 207 | /// Advances the reading position by `cnt`, dropping the `Bytes` references 208 | /// to any complete chunks that the position has been advanced past 209 | /// and then advancing the starting position of the first remaining chunk. 210 | /// If there are no complete chunks left, the reading position is advanced 211 | /// in the staging buffer, effectively removing the consumed bytes. 212 | /// 213 | /// # Panics 214 | /// 215 | /// This function may panic when `cnt > self.remaining()`. 216 | /// 217 | fn advance(&mut self, cnt: usize) { 218 | match self.inner.advance(cnt) { 219 | AdvanceStopped::InChunk => {} 220 | AdvanceStopped::InStaging(adv) => { 221 | self.cap -= adv; 222 | } 223 | } 224 | } 225 | 226 | /// Fills `dst` sequentially with the slice views of the chunks, then 227 | /// the bytes in the staging buffer if any remain and there is 228 | /// another unfilled entry left in `dst`. Returns the number of `IoSlice` 229 | /// entries filled. 230 | #[inline] 231 | fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize { 232 | debug_assert!(self.inner.staging_len() <= self.inner.chunk_size()); 233 | self.inner.chunks_vectored(dst) 234 | } 235 | 236 | #[inline] 237 | fn copy_to_bytes(&mut self, len: usize) -> Bytes { 238 | self.inner.copy_to_bytes(len) 239 | } 240 | } 241 | 242 | impl fmt::Write for ChunkedBytes { 243 | #[inline] 244 | fn write_str(&mut self, s: &str) -> fmt::Result { 245 | if self.remaining_mut() >= s.len() { 246 | self.put_slice(s.as_bytes()); 247 | Ok(()) 248 | } else { 249 | Err(fmt::Error) 250 | } 251 | } 252 | 253 | // The default implementation delegates to 254 | // fmt::write(&mut self as &mut dyn fmt::Write, args) 255 | #[inline] 256 | fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result { 257 | fmt::write(self, args) 258 | } 259 | } 260 | -------------------------------------------------------------------------------- /src/tests.rs: -------------------------------------------------------------------------------- 1 | use crate::{loosely, strictly, DrainChunks}; 2 | use bytes::{Buf, BufMut}; 3 | 4 | trait TestBuf: Buf + BufMut { 5 | fn with_chunk_size(size: usize) -> Self; 6 | fn drain_chunks(&mut self) -> DrainChunks<'_>; 7 | fn staging_capacity(&self) -> usize; 8 | } 9 | 10 | impl TestBuf for loosely::ChunkedBytes { 11 | fn with_chunk_size(size: usize) -> Self { 12 | loosely::ChunkedBytes::with_chunk_size_hint(size) 13 | } 14 | 15 | fn drain_chunks(&mut self) -> DrainChunks<'_> { 16 | self.drain_chunks() 17 | } 18 | 19 | fn staging_capacity(&self) -> usize { 20 | self.staging_capacity() 21 | } 22 | } 23 | 24 | impl TestBuf for strictly::ChunkedBytes { 25 | fn with_chunk_size(size: usize) -> Self { 26 | strictly::ChunkedBytes::with_chunk_size_limit(size) 27 | } 28 | 29 | fn drain_chunks(&mut self) -> DrainChunks<'_> { 30 | self.drain_chunks() 31 | } 32 | 33 | fn staging_capacity(&self) -> usize { 34 | self.staging_capacity() 35 | } 36 | } 37 | 38 | #[generic_tests::define] 39 | mod properties { 40 | use super::*; 41 | 42 | #[test] 43 | fn reserve_does_not_grow_staging_buffer() { 44 | let mut buf = B::with_chunk_size(8); 45 | let cap = buf.chunk_mut().len(); 46 | assert!(cap >= 8); 47 | 48 | buf.put(&vec![0; cap][..]); 49 | assert_eq!(buf.chunk_mut().len(), cap); 50 | { 51 | let mut chunks = buf.drain_chunks(); 52 | let chunk = chunks.next().expect("expected a chunk to be flushed"); 53 | assert_eq!(chunk.len(), cap); 54 | assert!(chunks.next().is_none()); 55 | } 56 | 57 | buf.put(&vec![0; cap - 4][..]); 58 | buf.advance(cap - 6); 59 | buf.put(&[0; 4][..]); 60 | assert_eq!(buf.chunk_mut().len(), cap); 61 | { 62 | let mut chunks = buf.drain_chunks(); 63 | let chunk = chunks.next().expect("expected a chunk to be flushed"); 64 | assert_eq!(chunk.len(), 6); 65 | assert!(chunks.next().is_none()); 66 | } 67 | 68 | buf.put(&vec![0; cap - 5][..]); 69 | buf.advance(cap - 5); 70 | buf.put(&[0; 5][..]); 71 | assert_eq!(buf.chunk_mut().len(), cap - 5); 72 | assert_eq!(buf.staging_capacity(), cap); 73 | assert!( 74 | buf.drain_chunks().next().is_none(), 75 | "expected no chunks to be flushed" 76 | ); 77 | } 78 | 79 | #[instantiate_tests()] 80 | mod loosely_chunked_bytes {} 81 | 82 | #[instantiate_tests()] 83 | mod strictly_chunked_bytes {} 84 | } 85 | --------------------------------------------------------------------------------