1use std::{io::SeekFrom, pin::Pin};
2
3use axum::body::Bytes;
4use futures_util::{Stream, TryStreamExt, stream};
5use shardline_protocol::ByteRange;
6use shardline_storage::{LocalObjectStore, ObjectKey, ObjectStore, S3ObjectStore};
7use tokio::{
8 fs::File,
9 io::{AsyncReadExt, AsyncSeekExt},
10};
11
12use crate::{ServerError, error::ObjectStoreError, object_store::ServerObjectStore};
13
14pub const STREAM_READ_BUFFER_BYTES: u64 = 1024 * 1024;
15
16pub type ServerByteStream = Pin<Box<dyn Stream<Item = Result<Bytes, ServerError>> + Send>>;
17
18struct LocalObjectByteStreamState {
19 file: File,
20 remaining: u64,
21}
22
23#[cfg(test)]
24pub(crate) async fn local_object_byte_stream(
25 object_store: LocalObjectStore,
26 object_key: ObjectKey,
27 length: u64,
28) -> Result<ServerByteStream, ServerError> {
29 local_store_byte_stream(object_store, object_key, length).await
30}
31
32#[cfg(test)]
33pub(crate) async fn local_object_byte_range_stream(
34 object_store: LocalObjectStore,
35 object_key: ObjectKey,
36 total_length: u64,
37 range: ByteRange,
38) -> Result<ServerByteStream, ServerError> {
39 local_store_byte_range_stream(object_store, object_key, total_length, range).await
40}
41
42pub(crate) async fn object_byte_range_stream(
43 object_store: ServerObjectStore,
44 object_key: ObjectKey,
45 total_length: u64,
46 range: ByteRange,
47) -> Result<ServerByteStream, ServerError> {
48 match object_store {
49 ServerObjectStore::Local(store) => {
50 local_store_byte_range_stream(store, object_key, total_length, range).await
51 }
52 ServerObjectStore::S3(store) => {
53 s3_store_byte_range_stream(store, object_key, total_length, range).await
54 }
55 ServerObjectStore::Blackhole => Err(ServerError::NotFound),
56 }
57}
58
59pub(crate) async fn object_byte_stream(
60 object_store: ServerObjectStore,
61 object_key: ObjectKey,
62 total_length: u64,
63) -> Result<ServerByteStream, ServerError> {
64 if total_length == 0 {
65 let metadata = object_store.metadata(&object_key)?;
66 let Some(metadata) = metadata else {
67 return Err(ServerError::NotFound);
68 };
69 if metadata.length() != 0 {
70 return Err(ServerError::ObjectStore(
71 ObjectStoreError::StoredLengthMismatch,
72 ));
73 }
74
75 return Ok(Box::pin(stream::empty()));
76 }
77
78 let end_inclusive = total_length.checked_sub(1).ok_or(ServerError::Overflow)?;
79 let range = ByteRange::new(0, end_inclusive).map_err(|_error| ServerError::Overflow)?;
80 object_byte_range_stream(object_store, object_key, total_length, range).await
81}
82
83#[cfg(test)]
84async fn local_store_byte_stream(
85 object_store: LocalObjectStore,
86 object_key: ObjectKey,
87 length: u64,
88) -> Result<ServerByteStream, ServerError> {
89 if length == 0 {
90 let file = object_store.open_object_file(&object_key)?;
91 let metadata = file.metadata()?;
92 if metadata.len() != 0 {
93 return Err(ServerError::ObjectStore(
94 ObjectStoreError::StoredLengthMismatch,
95 ));
96 }
97
98 return Ok(Box::pin(stream::empty()));
99 }
100
101 let end_inclusive = length.checked_sub(1).ok_or(ServerError::Overflow)?;
102 let range = ByteRange::new(0, end_inclusive).map_err(|_error| ServerError::Overflow)?;
103 local_store_byte_range_stream(object_store, object_key, length, range).await
104}
105
106async fn local_store_byte_range_stream(
107 object_store: LocalObjectStore,
108 object_key: ObjectKey,
109 total_length: u64,
110 range: ByteRange,
111) -> Result<ServerByteStream, ServerError> {
112 let file = object_store.open_object_file(&object_key)?;
113 let mut file = File::from_std(file);
114 let metadata = file.metadata().await?;
115 if metadata.len() != total_length {
116 return Err(ServerError::ObjectStore(
117 ObjectStoreError::StoredLengthMismatch,
118 ));
119 }
120 if range.end_inclusive() >= total_length {
121 return Err(ServerError::RangeNotSatisfiable);
122 }
123
124 file.seek(SeekFrom::Start(range.start())).await?;
125 let remaining = range.len().ok_or(ServerError::Overflow)?;
126 let state = LocalObjectByteStreamState { file, remaining };
127 let byte_stream = stream::try_unfold(state, |mut state| async move {
128 if state.remaining == 0 {
129 return Ok::<Option<(Bytes, LocalObjectByteStreamState)>, ServerError>(None);
130 }
131
132 let read_len_u64 = state.remaining.min(STREAM_READ_BUFFER_BYTES);
133 let read_len = usize::try_from(read_len_u64)?;
134 let mut buffer = vec![0_u8; read_len];
135 let read = state.file.read(&mut buffer).await?;
136 if read == 0 {
137 return Err(ServerError::ObjectStore(
138 ObjectStoreError::StoredLengthMismatch,
139 ));
140 }
141
142 buffer.truncate(read);
143 let read_u64 = u64::try_from(read)?;
144 state.remaining = state
145 .remaining
146 .checked_sub(read_u64)
147 .ok_or(ServerError::Overflow)?;
148
149 Ok(Some((Bytes::from(buffer), state)))
150 });
151
152 Ok(Box::pin(byte_stream))
153}
154
155async fn s3_store_byte_range_stream(
156 object_store: S3ObjectStore,
157 object_key: ObjectKey,
158 total_length: u64,
159 range: ByteRange,
160) -> Result<ServerByteStream, ServerError> {
161 let metadata = object_store.metadata(&object_key)?;
162 let Some(metadata) = metadata else {
163 return Err(ServerError::NotFound);
164 };
165 if metadata.length() != total_length {
166 return Err(ServerError::ObjectStore(
167 ObjectStoreError::StoredLengthMismatch,
168 ));
169 }
170 if range.end_inclusive() >= total_length {
171 return Err(ServerError::RangeNotSatisfiable);
172 }
173
174 let byte_stream = object_store.stream_range(&object_key, range).await?;
175
176 Ok(Box::pin(byte_stream.map_err(ServerError::from)))
177}
178
179#[cfg(test)]
180mod tests {
181 use futures_util::StreamExt;
182 use shardline_protocol::ByteRange;
183 use shardline_storage::{LocalObjectStore, ObjectKey};
184 use tokio::fs;
185
186 use super::{local_object_byte_range_stream, local_object_byte_stream};
187 use crate::ServerError;
188 use crate::error::ObjectStoreError;
189 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
190 async fn local_object_byte_stream_reads_object_in_segments() {
191 let storage = shardline_test_support::TempStorage::new();
192 let object_store = LocalObjectStore::new(storage.path_buf());
193 assert!(object_store.is_ok());
194 let Ok(object_store) = object_store else {
195 return;
196 };
197 let object_key = ObjectKey::parse("ab/object");
198 assert!(object_key.is_ok());
199 let Ok(object_key) = object_key else {
200 return;
201 };
202 let bytes = vec![7_u8; 64 * 1024 + 3];
203 let path = object_store.path_for_key(&object_key);
204 if let Some(parent) = path.parent() {
205 let created = fs::create_dir_all(parent).await;
206 assert!(created.is_ok());
207 }
208 let written = fs::write(&path, &bytes).await;
209 assert!(written.is_ok());
210
211 let byte_stream = local_object_byte_stream(
212 object_store,
213 object_key,
214 u64::try_from(bytes.len()).unwrap_or(0),
215 )
216 .await;
217 assert!(byte_stream.is_ok());
218 let Ok(mut byte_stream) = byte_stream else {
219 return;
220 };
221 let mut observed = Vec::with_capacity(bytes.len());
222 while let Some(item) = byte_stream.next().await {
223 assert!(item.is_ok());
224 let Ok(chunk) = item else {
225 return;
226 };
227 observed.extend_from_slice(&chunk);
228 }
229
230 assert_eq!(observed, bytes);
231 }
232
233 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
234 async fn local_object_byte_stream_rejects_index_length_mismatch() {
235 let storage = shardline_test_support::TempStorage::new();
236 let object_store = LocalObjectStore::new(storage.path_buf());
237 assert!(object_store.is_ok());
238 let Ok(object_store) = object_store else {
239 return;
240 };
241 let object_key = ObjectKey::parse("ab/object");
242 assert!(object_key.is_ok());
243 let Ok(object_key) = object_key else {
244 return;
245 };
246 let path = object_store.path_for_key(&object_key);
247 if let Some(parent) = path.parent() {
248 let created = fs::create_dir_all(parent).await;
249 assert!(created.is_ok());
250 }
251 let written = fs::write(&path, b"short").await;
252 assert!(written.is_ok());
253
254 let byte_stream = local_object_byte_stream(object_store, object_key, 100).await;
255
256 assert!(matches!(
257 byte_stream,
258 Err(ServerError::ObjectStore(
259 ObjectStoreError::StoredLengthMismatch
260 ))
261 ));
262 }
263
264 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
265 async fn local_object_byte_range_stream_reads_only_requested_range() {
266 let storage = shardline_test_support::TempStorage::new();
267 let object_store = LocalObjectStore::new(storage.path_buf());
268 assert!(object_store.is_ok());
269 let Ok(object_store) = object_store else {
270 return;
271 };
272 let object_key = ObjectKey::parse("ab/object");
273 assert!(object_key.is_ok());
274 let Ok(object_key) = object_key else {
275 return;
276 };
277 let bytes = b"abcdefghijkl".to_vec();
278 let path = object_store.path_for_key(&object_key);
279 if let Some(parent) = path.parent() {
280 let created = fs::create_dir_all(parent).await;
281 assert!(created.is_ok());
282 }
283 let written = fs::write(&path, &bytes).await;
284 assert!(written.is_ok());
285 let range = ByteRange::new(2, 7);
286 assert!(range.is_ok());
287 let Ok(range) = range else {
288 return;
289 };
290
291 let byte_stream = local_object_byte_range_stream(
292 object_store,
293 object_key,
294 u64::try_from(bytes.len()).unwrap_or(0),
295 range,
296 )
297 .await;
298 assert!(byte_stream.is_ok());
299 let Ok(mut byte_stream) = byte_stream else {
300 return;
301 };
302 let mut observed = Vec::new();
303 while let Some(item) = byte_stream.next().await {
304 assert!(item.is_ok());
305 let Ok(chunk) = item else {
306 return;
307 };
308 observed.extend_from_slice(&chunk);
309 }
310
311 assert_eq!(observed, b"cdefgh");
312 }
313}