A fork of attic a self-hostable Nix Binary Cache server
server: Use the same read_chunk_async implementation
+15
-37
+7
-19
attic/src/stream.rs
+7
-19
attic/src/stream.rs
···
8
8
use std::task::{Context, Poll};
9
9
10
10
use async_stream::try_stream;
11
-
use bytes::Bytes;
11
+
use bytes::{Bytes, BytesMut};
12
12
use digest::{Digest, Output as DigestOutput};
13
13
use futures::stream::{BoxStream, Stream, StreamExt};
14
14
use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
···
152
152
}
153
153
}
154
154
155
-
/// Greedily reads from a stream for some number of bytes.
156
-
///
157
-
/// This was originally from rust-s3 but completely rewritten to resolve
158
-
/// performance problems.
155
+
/// Greedily reads from a stream to fill a buffer.
159
156
pub async fn read_chunk_async<S: AsyncRead + Unpin + Send>(
160
157
stream: &mut S,
161
-
max_size: usize,
162
-
) -> std::io::Result<Vec<u8>> {
163
-
let mut chunk: Box<[u8]> = vec![0u8; max_size].into_boxed_slice();
164
-
let mut cursor = 0;
165
-
166
-
while cursor < max_size {
167
-
let buf = &mut chunk[cursor..];
168
-
let read = stream.read(buf).await?;
158
+
mut chunk: BytesMut,
159
+
) -> std::io::Result<Bytes> {
160
+
while chunk.len() < chunk.capacity() {
161
+
let read = stream.read_buf(&mut chunk).await?;
169
162
170
163
if read == 0 {
171
164
break;
172
-
} else {
173
-
cursor += read;
174
165
}
175
166
}
176
167
177
-
let mut vec = chunk.into_vec();
178
-
vec.truncate(cursor);
179
-
180
-
Ok(vec)
168
+
Ok(chunk.freeze())
181
169
}
182
170
183
171
#[cfg(test)]
+3
-16
server/src/chunking/mod.rs
+3
-16
server/src/chunking/mod.rs
···
7
7
use bytes::{BufMut, Bytes, BytesMut};
8
8
use fastcdc::FastCDC;
9
9
use futures::stream::Stream;
10
-
use tokio::io::{AsyncRead, AsyncReadExt};
10
+
use tokio::io::AsyncRead;
11
+
12
+
use attic::stream::read_chunk_async;
11
13
12
14
/// Splits a streams into content-defined chunks.
13
15
///
···
61
63
};
62
64
63
65
Box::pin(s)
64
-
}
65
-
66
-
async fn read_chunk_async<S: AsyncRead + Unpin + Send>(
67
-
stream: &mut S,
68
-
mut chunk: BytesMut,
69
-
) -> std::io::Result<Bytes> {
70
-
while chunk.len() < chunk.capacity() {
71
-
let read = stream.read_buf(&mut chunk).await?;
72
-
73
-
if read == 0 {
74
-
break;
75
-
}
76
-
}
77
-
78
-
Ok(chunk.freeze())
79
66
}
80
67
81
68
#[cfg(test)]
+5
-2
server/src/storage/s3.rs
+5
-2
server/src/storage/s3.rs
···
11
11
presigning::config::PresigningConfig,
12
12
Client, Credentials, Endpoint, Region,
13
13
};
14
+
use bytes::BytesMut;
14
15
use futures::future::join_all;
15
16
use futures::stream::StreamExt;
16
17
use serde::{Deserialize, Serialize};
···
171
172
name: String,
172
173
mut stream: &mut (dyn AsyncRead + Unpin + Send),
173
174
) -> ServerResult<RemoteFile> {
174
-
let first_chunk = read_chunk_async(&mut stream, CHUNK_SIZE)
175
+
let buf = BytesMut::with_capacity(CHUNK_SIZE);
176
+
let first_chunk = read_chunk_async(&mut stream, buf)
175
177
.await
176
178
.map_err(ServerError::storage_error)?;
177
179
···
238
240
let chunk = if part_number == 1 {
239
241
first_chunk.take().unwrap()
240
242
} else {
241
-
read_chunk_async(&mut stream, CHUNK_SIZE)
243
+
let buf = BytesMut::with_capacity(CHUNK_SIZE);
244
+
read_chunk_async(&mut stream, buf)
242
245
.await
243
246
.map_err(ServerError::storage_error)?
244
247
};