1use crate::path::Path;
21use crate::{
22 Attributes, Extensions, ObjectMeta, ObjectStore, ObjectStoreExt, PutMultipartOptions,
23 PutOptions, PutPayloadMut, TagSet, WriteMultipart,
24};
25use bytes::Bytes;
26use futures_util::future::{BoxFuture, FutureExt};
27use futures_util::ready;
28use std::cmp::Ordering;
29use std::io::{Error, ErrorKind, SeekFrom};
30use std::pin::Pin;
31use std::sync::Arc;
32use std::task::{Context, Poll};
33use tokio::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
34
35pub const DEFAULT_BUFFER_SIZE: usize = 1024 * 1024;
37
38pub struct BufReader {
57 store: Arc<dyn ObjectStore>,
59 size: u64,
61 path: Path,
63 cursor: u64,
65 capacity: usize,
67 buffer: Buffer,
69}
70
71impl std::fmt::Debug for BufReader {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 f.debug_struct("BufReader")
74 .field("path", &self.path)
75 .field("size", &self.size)
76 .field("capacity", &self.capacity)
77 .finish()
78 }
79}
80
81enum Buffer {
82 Empty,
83 Pending(BoxFuture<'static, std::io::Result<Bytes>>),
84 Ready(Bytes),
85}
86
87impl BufReader {
88 pub fn new(store: Arc<dyn ObjectStore>, meta: &ObjectMeta) -> Self {
90 Self::with_capacity(store, meta, DEFAULT_BUFFER_SIZE)
91 }
92
93 pub fn with_capacity(store: Arc<dyn ObjectStore>, meta: &ObjectMeta, capacity: usize) -> Self {
95 Self {
96 path: meta.location.clone(),
97 size: meta.size as _,
98 store,
99 capacity,
100 cursor: 0,
101 buffer: Buffer::Empty,
102 }
103 }
104
105 fn poll_fill_buf_impl(
106 &mut self,
107 cx: &mut Context<'_>,
108 amnt: usize,
109 ) -> Poll<std::io::Result<&[u8]>> {
110 let buf = &mut self.buffer;
111 loop {
112 match buf {
113 Buffer::Empty => {
114 let store = Arc::clone(&self.store);
115 let path = self.path.clone();
116 let start = self.cursor.min(self.size) as _;
117 let end = self.cursor.saturating_add(amnt as u64).min(self.size) as _;
118
119 if start == end {
120 return Poll::Ready(Ok(&[]));
121 }
122
123 *buf = Buffer::Pending(Box::pin(async move {
124 Ok(store.get_range(&path, start..end).await?)
125 }))
126 }
127 Buffer::Pending(fut) => match ready!(fut.poll_unpin(cx)) {
128 Ok(b) => *buf = Buffer::Ready(b),
129 Err(e) => return Poll::Ready(Err(e)),
130 },
131 Buffer::Ready(r) => return Poll::Ready(Ok(r)),
132 }
133 }
134 }
135}
136
137impl AsyncSeek for BufReader {
138 fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> std::io::Result<()> {
139 self.cursor = match position {
140 SeekFrom::Start(offset) => offset,
141 SeekFrom::End(offset) => checked_add_signed(self.size, offset).ok_or_else(|| {
142 Error::new(
143 ErrorKind::InvalidInput,
144 format!(
145 "Seeking {offset} from end of {} byte file would result in overflow",
146 self.size
147 ),
148 )
149 })?,
150 SeekFrom::Current(offset) => {
151 checked_add_signed(self.cursor, offset).ok_or_else(|| {
152 Error::new(
153 ErrorKind::InvalidInput,
154 format!(
155 "Seeking {offset} from current offset of {} would result in overflow",
156 self.cursor
157 ),
158 )
159 })?
160 }
161 };
162 self.buffer = Buffer::Empty;
163 Ok(())
164 }
165
166 fn poll_complete(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<u64>> {
167 Poll::Ready(Ok(self.cursor))
168 }
169}
170
171impl AsyncRead for BufReader {
172 fn poll_read(
173 mut self: Pin<&mut Self>,
174 cx: &mut Context<'_>,
175 out: &mut ReadBuf<'_>,
176 ) -> Poll<std::io::Result<()>> {
177 let to_read = out.remaining().max(self.capacity);
179 let r = match ready!(self.poll_fill_buf_impl(cx, to_read)) {
180 Ok(buf) => {
181 let to_consume = out.remaining().min(buf.len());
182 out.put_slice(&buf[..to_consume]);
183 self.consume(to_consume);
184 Ok(())
185 }
186 Err(e) => Err(e),
187 };
188 Poll::Ready(r)
189 }
190}
191
192impl AsyncBufRead for BufReader {
193 fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
194 let capacity = self.capacity;
195 self.get_mut().poll_fill_buf_impl(cx, capacity)
196 }
197
198 fn consume(mut self: Pin<&mut Self>, amt: usize) {
199 match &mut self.buffer {
200 Buffer::Empty => assert_eq!(amt, 0, "cannot consume from empty buffer"),
201 Buffer::Ready(b) => match b.len().cmp(&amt) {
202 Ordering::Less => panic!("{amt} exceeds buffer sized of {}", b.len()),
203 Ordering::Greater => *b = b.slice(amt..),
204 Ordering::Equal => self.buffer = Buffer::Empty,
205 },
206 Buffer::Pending(_) => panic!("cannot consume from pending buffer"),
207 }
208 self.cursor += amt as u64;
209 }
210}
211
212pub struct BufWriter {
222 capacity: usize,
223 max_concurrency: usize,
224 attributes: Option<Attributes>,
225 tags: Option<TagSet>,
226 extensions: Option<Extensions>,
227 state: BufWriterState,
228 store: Arc<dyn ObjectStore>,
229}
230
231impl std::fmt::Debug for BufWriter {
232 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
233 f.debug_struct("BufWriter")
234 .field("capacity", &self.capacity)
235 .finish()
236 }
237}
238
239enum BufWriterState {
240 Buffer(Path, PutPayloadMut),
242 Prepare(BoxFuture<'static, crate::Result<WriteMultipart>>),
244 Write(Option<WriteMultipart>),
246 Flush(BoxFuture<'static, crate::Result<()>>),
248}
249
250impl BufWriter {
251 pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
253 Self::with_capacity(store, path, 10 * 1024 * 1024)
254 }
255
256 pub fn with_capacity(store: Arc<dyn ObjectStore>, path: Path, capacity: usize) -> Self {
258 Self {
259 capacity,
260 store,
261 max_concurrency: 8,
262 attributes: None,
263 tags: None,
264 extensions: None,
265 state: BufWriterState::Buffer(path, PutPayloadMut::new()),
266 }
267 }
268
269 pub fn with_max_concurrency(self, max_concurrency: usize) -> Self {
273 Self {
274 max_concurrency,
275 ..self
276 }
277 }
278
279 pub fn with_attributes(self, attributes: Attributes) -> Self {
281 Self {
282 attributes: Some(attributes),
283 ..self
284 }
285 }
286
287 pub fn with_tags(self, tags: TagSet) -> Self {
289 Self {
290 tags: Some(tags),
291 ..self
292 }
293 }
294
295 pub fn with_extensions(self, extensions: Extensions) -> Self {
302 Self {
303 extensions: Some(extensions),
304 ..self
305 }
306 }
307
308 pub async fn put(&mut self, bytes: Bytes) -> crate::Result<()> {
314 loop {
315 return match &mut self.state {
316 BufWriterState::Write(Some(write)) => {
317 write.wait_for_capacity(self.max_concurrency).await?;
318 write.put(bytes);
319 Ok(())
320 }
321 BufWriterState::Write(None) | BufWriterState::Flush(_) => {
322 panic!("Already shut down")
323 }
324 BufWriterState::Prepare(f) => {
331 self.state = BufWriterState::Write(f.await?.into());
332 continue;
333 }
334 BufWriterState::Buffer(path, b) => {
335 if b.content_length().saturating_add(bytes.len()) < self.capacity {
336 b.push(bytes);
337 Ok(())
338 } else {
339 let buffer = std::mem::take(b);
340 let path = std::mem::take(path);
341 let opts = PutMultipartOptions {
342 attributes: self.attributes.take().unwrap_or_default(),
343 tags: self.tags.take().unwrap_or_default(),
344 extensions: self.extensions.take().unwrap_or_default(),
345 };
346 let upload = self.store.put_multipart_opts(&path, opts).await?;
347 let mut chunked =
348 WriteMultipart::new_with_chunk_size(upload, self.capacity);
349 for chunk in buffer.freeze() {
350 chunked.put(chunk);
351 }
352 chunked.put(bytes);
353 self.state = BufWriterState::Write(Some(chunked));
354 Ok(())
355 }
356 }
357 };
358 }
359 }
360
361 pub async fn abort(&mut self) -> crate::Result<()> {
367 match &mut self.state {
368 BufWriterState::Buffer(_, _) | BufWriterState::Prepare(_) => Ok(()),
369 BufWriterState::Flush(_) => panic!("Already shut down"),
370 BufWriterState::Write(x) => x.take().unwrap().abort().await,
371 }
372 }
373}
374
375impl AsyncWrite for BufWriter {
376 fn poll_write(
377 mut self: Pin<&mut Self>,
378 cx: &mut Context<'_>,
379 buf: &[u8],
380 ) -> Poll<Result<usize, Error>> {
381 let cap = self.capacity;
382 let max_concurrency = self.max_concurrency;
383 loop {
384 return match &mut self.state {
385 BufWriterState::Write(Some(write)) => {
386 ready!(write.poll_for_capacity(cx, max_concurrency))?;
387 write.write(buf);
388 Poll::Ready(Ok(buf.len()))
389 }
390 BufWriterState::Write(None) | BufWriterState::Flush(_) => {
391 panic!("Already shut down")
392 }
393 BufWriterState::Prepare(f) => {
394 self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
395 continue;
396 }
397 BufWriterState::Buffer(path, b) => {
398 if b.content_length().saturating_add(buf.len()) >= cap {
399 let buffer = std::mem::take(b);
400 let path = std::mem::take(path);
401 let opts = PutMultipartOptions {
402 attributes: self.attributes.take().unwrap_or_default(),
403 tags: self.tags.take().unwrap_or_default(),
404 extensions: self.extensions.take().unwrap_or_default(),
405 };
406 let store = Arc::clone(&self.store);
407 self.state = BufWriterState::Prepare(Box::pin(async move {
408 let upload = store.put_multipart_opts(&path, opts).await?;
409 let mut chunked = WriteMultipart::new_with_chunk_size(upload, cap);
410 for chunk in buffer.freeze() {
411 chunked.put(chunk);
412 }
413 Ok(chunked)
414 }));
415 continue;
416 }
417 b.extend_from_slice(buf);
418 Poll::Ready(Ok(buf.len()))
419 }
420 };
421 }
422 }
423
424 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
425 loop {
426 return match &mut self.state {
427 BufWriterState::Write(_) | BufWriterState::Buffer(_, _) => Poll::Ready(Ok(())),
428 BufWriterState::Flush(_) => panic!("Already shut down"),
429 BufWriterState::Prepare(f) => {
430 self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
431 continue;
432 }
433 };
434 }
435 }
436
437 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
438 loop {
439 match &mut self.state {
440 BufWriterState::Prepare(f) => {
441 self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
442 }
443 BufWriterState::Buffer(p, b) => {
444 let buf = std::mem::take(b);
445 let path = std::mem::take(p);
446 let opts = PutOptions {
447 attributes: self.attributes.take().unwrap_or_default(),
448 tags: self.tags.take().unwrap_or_default(),
449 extensions: self.extensions.take().unwrap_or_default(),
450 ..Default::default()
451 };
452 let store = Arc::clone(&self.store);
453 self.state = BufWriterState::Flush(Box::pin(async move {
454 store.put_opts(&path, buf.into(), opts).await?;
455 Ok(())
456 }));
457 }
458 BufWriterState::Flush(f) => return f.poll_unpin(cx).map_err(std::io::Error::from),
459 BufWriterState::Write(x) => {
460 let upload = x.take().ok_or_else(|| {
461 std::io::Error::new(
462 ErrorKind::InvalidInput,
463 "Cannot shutdown a writer that has already been shut down",
464 )
465 })?;
466 self.state = BufWriterState::Flush(
467 async move {
468 upload.finish().await?;
469 Ok(())
470 }
471 .boxed(),
472 )
473 }
474 }
475 }
476 }
477}
478
479#[inline]
483fn checked_add_signed(a: u64, rhs: i64) -> Option<u64> {
484 let (res, overflowed) = a.overflowing_add(rhs as _);
485 let overflow = overflowed ^ (rhs < 0);
486 (!overflow).then_some(res)
487}
488
489#[cfg(test)]
490mod tests {
491 use super::*;
492 use crate::memory::InMemory;
493 use crate::path::Path;
494 use crate::{Attribute, GetOptions, ObjectStoreExt};
495 use itertools::Itertools;
496 use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
497
498 #[tokio::test]
499 async fn test_buf_reader() {
500 let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
501
502 let existent = Path::from("exists.txt");
503 const BYTES: usize = 4096;
504
505 let data: Bytes = b"12345678".iter().cycle().copied().take(BYTES).collect();
506 store.put(&existent, data.clone().into()).await.unwrap();
507
508 let meta = store.head(&existent).await.unwrap();
509
510 let mut reader = BufReader::new(Arc::clone(&store), &meta);
511 let mut out = Vec::with_capacity(BYTES);
512 let read = reader.read_to_end(&mut out).await.unwrap();
513
514 assert_eq!(read, BYTES);
515 assert_eq!(&out, &data);
516
517 let err = reader.seek(SeekFrom::Current(i64::MIN)).await.unwrap_err();
518 assert_eq!(
519 err.to_string(),
520 "Seeking -9223372036854775808 from current offset of 4096 would result in overflow"
521 );
522
523 reader.rewind().await.unwrap();
524
525 let err = reader.seek(SeekFrom::Current(-1)).await.unwrap_err();
526 assert_eq!(
527 err.to_string(),
528 "Seeking -1 from current offset of 0 would result in overflow"
529 );
530
531 reader.seek(SeekFrom::Start(u64::MAX)).await.unwrap();
533 let buf = reader.fill_buf().await.unwrap();
534 assert!(buf.is_empty());
535
536 let err = reader.seek(SeekFrom::Current(1)).await.unwrap_err();
537 assert_eq!(
538 err.to_string(),
539 "Seeking 1 from current offset of 18446744073709551615 would result in overflow"
540 );
541
542 for capacity in [200, 1024, 4096, DEFAULT_BUFFER_SIZE] {
543 let store = Arc::clone(&store);
544 let mut reader = BufReader::with_capacity(store, &meta, capacity);
545
546 let mut bytes_read = 0;
547 loop {
548 let buf = reader.fill_buf().await.unwrap();
549 if buf.is_empty() {
550 assert_eq!(bytes_read, BYTES);
551 break;
552 }
553 assert!(buf.starts_with(b"12345678"));
554 bytes_read += 8;
555 reader.consume(8);
556 }
557
558 let mut buf = Vec::with_capacity(76);
559 reader.seek(SeekFrom::Current(-76)).await.unwrap();
560 reader.read_to_end(&mut buf).await.unwrap();
561 assert_eq!(&buf, &data[BYTES - 76..]);
562
563 reader.rewind().await.unwrap();
564 let buffer = reader.fill_buf().await.unwrap();
565 assert_eq!(buffer, &data[..capacity.min(BYTES)]);
566
567 reader.seek(SeekFrom::Start(325)).await.unwrap();
568 let buffer = reader.fill_buf().await.unwrap();
569 assert_eq!(buffer, &data[325..(325 + capacity).min(BYTES)]);
570
571 reader.seek(SeekFrom::End(0)).await.unwrap();
572 let buffer = reader.fill_buf().await.unwrap();
573 assert!(buffer.is_empty());
574 }
575 }
576
577 #[tokio::test]
579 async fn test_buf_writer() {
580 let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
581 let path = Path::from("file.txt");
582 let attributes = Attributes::from_iter([
583 (Attribute::ContentType, "text/html"),
584 (Attribute::CacheControl, "max-age=604800"),
585 ]);
586
587 let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30)
589 .with_attributes(attributes.clone());
590 writer.write_all(&[0; 20]).await.unwrap();
591 writer.flush().await.unwrap();
592 writer.write_all(&[0; 5]).await.unwrap();
593 writer.shutdown().await.unwrap();
594 let response = store
595 .get_opts(&path, GetOptions::new().with_head(true))
596 .await
597 .unwrap();
598 assert_eq!(response.meta.size, 25);
599 assert_eq!(response.attributes, attributes);
600
601 let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30)
603 .with_attributes(attributes.clone());
604 writer.write_all(&[0; 20]).await.unwrap();
605 writer.flush().await.unwrap();
606 writer.write_all(&[0; 20]).await.unwrap();
607 writer.shutdown().await.unwrap();
608 let response = store
609 .get_opts(&path, GetOptions::new().with_head(true))
610 .await
611 .unwrap();
612 assert_eq!(response.meta.size, 40);
613 assert_eq!(response.attributes, attributes);
614 }
615
616 #[tokio::test]
617 async fn test_buf_writer_with_put() {
618 let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
619 let path = Path::from("file.txt");
620
621 let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30);
623 writer
624 .put(Bytes::from((0..20).collect_vec()))
625 .await
626 .unwrap();
627 writer
628 .put(Bytes::from((20..25).collect_vec()))
629 .await
630 .unwrap();
631 writer.shutdown().await.unwrap();
632 let response = store
633 .get_opts(&path, GetOptions::new().with_head(true))
634 .await
635 .unwrap();
636 assert_eq!(response.meta.size, 25);
637 assert_eq!(response.bytes().await.unwrap(), (0..25).collect_vec());
638
639 let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30);
641 writer
642 .put(Bytes::from((0..20).collect_vec()))
643 .await
644 .unwrap();
645 writer
646 .put(Bytes::from((20..40).collect_vec()))
647 .await
648 .unwrap();
649 writer.shutdown().await.unwrap();
650 let response = store
651 .get_opts(&path, GetOptions::new().with_head(true))
652 .await
653 .unwrap();
654 assert_eq!(response.meta.size, 40);
655 assert_eq!(response.bytes().await.unwrap(), (0..40).collect_vec());
656 }
657}