1use asupersync::Cx;
26use asupersync::stream::Stream;
27use std::io::{self, Read, Seek, SeekFrom};
28use std::path::Path;
29use std::pin::Pin;
30use std::task::{Context, Poll};
31
32pub const DEFAULT_CHUNK_SIZE: usize = 64 * 1024;
34
35pub const DEFAULT_MAX_BUFFER_SIZE: usize = 4 * 1024 * 1024;
37
38#[derive(Debug, Clone)]
40pub struct StreamConfig {
41 chunk_size: usize,
43 max_buffer_size: usize,
45 checkpoint_enabled: bool,
47}
48
49impl Default for StreamConfig {
50 fn default() -> Self {
51 Self {
52 chunk_size: DEFAULT_CHUNK_SIZE,
53 max_buffer_size: DEFAULT_MAX_BUFFER_SIZE,
54 checkpoint_enabled: true,
55 }
56 }
57}
58
59impl StreamConfig {
60 #[must_use]
62 pub fn new() -> Self {
63 Self::default()
64 }
65
66 #[must_use]
68 pub fn with_chunk_size(mut self, size: usize) -> Self {
69 self.chunk_size = size.max(1024); self
71 }
72
73 #[must_use]
75 pub fn with_max_buffer_size(mut self, size: usize) -> Self {
76 self.max_buffer_size = size;
77 self
78 }
79
80 #[must_use]
82 pub fn with_checkpoint(mut self, enabled: bool) -> Self {
83 self.checkpoint_enabled = enabled;
84 self
85 }
86
87 #[must_use]
89 pub fn chunk_size(&self) -> usize {
90 self.chunk_size
91 }
92
93 #[must_use]
95 pub fn max_buffer_size(&self) -> usize {
96 self.max_buffer_size
97 }
98
99 #[must_use]
101 pub fn checkpoint_enabled(&self) -> bool {
102 self.checkpoint_enabled
103 }
104}
105
106#[derive(Debug)]
108pub enum StreamError {
109 Io(io::Error),
111 Cancelled,
113 BufferFull,
115}
116
117impl std::fmt::Display for StreamError {
118 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119 match self {
120 Self::Io(e) => write!(f, "streaming I/O error: {e}"),
121 Self::Cancelled => write!(f, "stream cancelled"),
122 Self::BufferFull => write!(f, "stream buffer full"),
123 }
124 }
125}
126
127impl std::error::Error for StreamError {
128 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
129 match self {
130 Self::Io(e) => Some(e),
131 _ => None,
132 }
133 }
134}
135
136impl From<io::Error> for StreamError {
137 fn from(e: io::Error) -> Self {
138 Self::Io(e)
139 }
140}
141
142pub struct CancelAwareStream<S> {
152 inner: S,
153 cx: Cx,
154 cancelled: bool,
155}
156
157impl<S> CancelAwareStream<S> {
158 pub fn new(inner: S, cx: Cx) -> Self {
160 Self {
161 inner,
162 cx,
163 cancelled: false,
164 }
165 }
166
167 #[must_use]
169 pub fn is_cancelled(&self) -> bool {
170 self.cancelled
171 }
172}
173
174impl<S> Stream for CancelAwareStream<S>
175where
176 S: Stream + Unpin,
177{
178 type Item = S::Item;
179
180 fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
181 if self.cx.is_cancel_requested() {
183 self.cancelled = true;
184 return Poll::Ready(None);
185 }
186
187 Pin::new(&mut self.inner).poll_next(ctx)
189 }
190}
191
192enum FileStreamState {
194 Active {
196 file: std::fs::File,
197 buffer: Vec<u8>,
198 remaining: u64,
199 },
200 Complete,
202 Error,
204}
205
206pub struct FileStream {
228 state: FileStreamState,
229 cx: Cx,
230 config: StreamConfig,
231}
232
233impl FileStream {
234 pub fn open<P: AsRef<Path>>(path: P, cx: Cx, config: StreamConfig) -> io::Result<Self> {
246 let mut file = std::fs::File::open(path)?;
247 let metadata = file.metadata()?;
248 let file_size = metadata.len();
249
250 file.seek(SeekFrom::Start(0))?;
252
253 let buffer = Vec::with_capacity(config.chunk_size);
254
255 Ok(Self {
256 state: FileStreamState::Active {
257 file,
258 buffer,
259 remaining: file_size,
260 },
261 cx,
262 config,
263 })
264 }
265
266 pub fn open_range<P: AsRef<Path>>(
282 path: P,
283 start: u64,
284 length: u64,
285 cx: Cx,
286 config: StreamConfig,
287 ) -> io::Result<Self> {
288 let mut file = std::fs::File::open(path)?;
289 file.seek(SeekFrom::Start(start))?;
290
291 let buffer = Vec::with_capacity(config.chunk_size);
292
293 Ok(Self {
294 state: FileStreamState::Active {
295 file,
296 buffer,
297 remaining: length,
298 },
299 cx,
300 config,
301 })
302 }
303
304 #[must_use]
306 pub fn remaining(&self) -> u64 {
307 match &self.state {
308 FileStreamState::Active { remaining, .. } => *remaining,
309 _ => 0,
310 }
311 }
312
313 #[must_use]
315 pub fn is_complete(&self) -> bool {
316 matches!(self.state, FileStreamState::Complete)
317 }
318}
319
320impl Stream for FileStream {
321 type Item = Vec<u8>;
322
323 fn poll_next(mut self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
324 if self.cx.is_cancel_requested() {
326 self.state = FileStreamState::Complete;
327 return Poll::Ready(None);
328 }
329
330 let chunk_size = self.config.chunk_size;
332
333 match &mut self.state {
334 FileStreamState::Active {
335 file,
336 buffer,
337 remaining,
338 } => {
339 if *remaining == 0 {
340 self.state = FileStreamState::Complete;
341 return Poll::Ready(None);
342 }
343
344 let to_read = (chunk_size as u64).min(*remaining) as usize;
346
347 buffer.clear();
349 buffer.resize(to_read, 0);
350
351 match file.read(&mut buffer[..to_read]) {
353 Ok(0) => {
354 self.state = FileStreamState::Complete;
356 Poll::Ready(None)
357 }
358 Ok(n) => {
359 *remaining -= n as u64;
360 buffer.truncate(n);
361
362 let chunk = std::mem::take(buffer);
364 *buffer = Vec::with_capacity(chunk_size);
365
366 Poll::Ready(Some(chunk))
367 }
368 Err(e) if e.kind() == io::ErrorKind::Interrupted => {
369 _ctx.waker().wake_by_ref();
371 Poll::Pending
372 }
373 Err(_) => {
374 self.state = FileStreamState::Error;
375 Poll::Ready(None)
376 }
377 }
378 }
379 FileStreamState::Complete | FileStreamState::Error => Poll::Ready(None),
380 }
381 }
382}
383
384#[allow(unsafe_code)]
391unsafe impl Send for FileStream {}
392
393pub struct ChunkedBytes {
397 data: Vec<u8>,
398 position: usize,
399 chunk_size: usize,
400}
401
402impl ChunkedBytes {
403 #[must_use]
405 pub fn new(data: Vec<u8>, chunk_size: usize) -> Self {
406 Self {
407 data,
408 position: 0,
409 chunk_size: chunk_size.max(1),
410 }
411 }
412
413 #[must_use]
415 pub fn with_default_chunks(data: Vec<u8>) -> Self {
416 Self::new(data, DEFAULT_CHUNK_SIZE)
417 }
418
419 #[must_use]
421 pub fn total_size(&self) -> usize {
422 self.data.len()
423 }
424
425 #[must_use]
427 pub fn remaining(&self) -> usize {
428 self.data.len().saturating_sub(self.position)
429 }
430}
431
432impl Stream for ChunkedBytes {
433 type Item = Vec<u8>;
434
435 fn poll_next(mut self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
436 if self.position >= self.data.len() {
437 return Poll::Ready(None);
438 }
439
440 let end = (self.position + self.chunk_size).min(self.data.len());
441 let chunk = self.data[self.position..end].to_vec();
442 self.position = end;
443
444 Poll::Ready(Some(chunk))
445 }
446}
447
448pub trait StreamingResponseExt {
450 fn stream_file<P: AsRef<Path>>(
462 path: P,
463 cx: Cx,
464 content_type: &[u8],
465 ) -> io::Result<fastapi_core::Response>;
466
467 fn stream_file_with_config<P: AsRef<Path>>(
473 path: P,
474 cx: Cx,
475 content_type: &[u8],
476 config: StreamConfig,
477 ) -> io::Result<fastapi_core::Response>;
478
479 fn stream_file_range<P: AsRef<Path>>(
495 path: P,
496 range: crate::range::ByteRange,
497 total_size: u64,
498 cx: Cx,
499 content_type: &[u8],
500 ) -> io::Result<fastapi_core::Response>;
501
502 fn stream_file_range_with_config<P: AsRef<Path>>(
508 path: P,
509 range: crate::range::ByteRange,
510 total_size: u64,
511 cx: Cx,
512 content_type: &[u8],
513 config: StreamConfig,
514 ) -> io::Result<fastapi_core::Response>;
515}
516
517impl StreamingResponseExt for fastapi_core::Response {
518 fn stream_file<P: AsRef<Path>>(
519 path: P,
520 cx: Cx,
521 content_type: &[u8],
522 ) -> io::Result<fastapi_core::Response> {
523 Self::stream_file_with_config(path, cx, content_type, StreamConfig::default())
524 }
525
526 fn stream_file_with_config<P: AsRef<Path>>(
527 path: P,
528 cx: Cx,
529 content_type: &[u8],
530 config: StreamConfig,
531 ) -> io::Result<fastapi_core::Response> {
532 let stream = FileStream::open(path, cx, config)?;
533
534 Ok(fastapi_core::Response::ok()
535 .header("content-type", content_type.to_vec())
536 .header("accept-ranges", b"bytes".to_vec())
537 .body(fastapi_core::ResponseBody::stream(stream)))
538 }
539
540 fn stream_file_range<P: AsRef<Path>>(
541 path: P,
542 range: crate::range::ByteRange,
543 total_size: u64,
544 cx: Cx,
545 content_type: &[u8],
546 ) -> io::Result<fastapi_core::Response> {
547 Self::stream_file_range_with_config(
548 path,
549 range,
550 total_size,
551 cx,
552 content_type,
553 StreamConfig::default(),
554 )
555 }
556
557 fn stream_file_range_with_config<P: AsRef<Path>>(
558 path: P,
559 range: crate::range::ByteRange,
560 total_size: u64,
561 cx: Cx,
562 content_type: &[u8],
563 config: StreamConfig,
564 ) -> io::Result<fastapi_core::Response> {
565 let stream = FileStream::open_range(path, range.start, range.len(), cx, config)?;
566
567 Ok(fastapi_core::Response::partial_content()
568 .header("content-type", content_type.to_vec())
569 .header("accept-ranges", b"bytes".to_vec())
570 .header(
571 "content-range",
572 range.content_range_header(total_size).into_bytes(),
573 )
574 .header("content-length", range.len().to_string().into_bytes())
575 .body(fastapi_core::ResponseBody::stream(stream)))
576 }
577}
578
579#[cfg(test)]
580mod tests {
581 use super::*;
582 use std::task::Waker;
583
584 fn noop_waker() -> Waker {
585 Waker::noop().clone()
586 }
587
588 #[test]
589 fn stream_config_defaults() {
590 let config = StreamConfig::default();
591 assert_eq!(config.chunk_size(), DEFAULT_CHUNK_SIZE);
592 assert_eq!(config.max_buffer_size(), DEFAULT_MAX_BUFFER_SIZE);
593 assert!(config.checkpoint_enabled());
594 }
595
596 #[test]
597 fn stream_config_custom() {
598 let config = StreamConfig::new()
599 .with_chunk_size(1024)
600 .with_max_buffer_size(2048)
601 .with_checkpoint(false);
602
603 assert_eq!(config.chunk_size(), 1024);
604 assert_eq!(config.max_buffer_size(), 2048);
605 assert!(!config.checkpoint_enabled());
606 }
607
608 #[test]
609 fn stream_config_minimum_chunk_size() {
610 let config = StreamConfig::new().with_chunk_size(100);
611 assert_eq!(config.chunk_size(), 1024);
613 }
614
615 #[test]
616 fn chunked_bytes_basic() {
617 let data = b"Hello, World!".to_vec();
618 let mut stream = ChunkedBytes::new(data.clone(), 5);
619
620 assert_eq!(stream.total_size(), 13);
621 assert_eq!(stream.remaining(), 13);
622
623 let waker = noop_waker();
624 let mut ctx = Context::from_waker(&waker);
625
626 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
628 assert_eq!(chunk, Poll::Ready(Some(b"Hello".to_vec())));
629 assert_eq!(stream.remaining(), 8);
630
631 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
633 assert_eq!(chunk, Poll::Ready(Some(b", Wor".to_vec())));
634
635 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
637 assert_eq!(chunk, Poll::Ready(Some(b"ld!".to_vec())));
638
639 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
641 assert_eq!(chunk, Poll::Ready(None));
642 }
643
644 #[test]
645 fn chunked_bytes_empty() {
646 let mut stream = ChunkedBytes::new(Vec::new(), 5);
647 let waker = noop_waker();
648 let mut ctx = Context::from_waker(&waker);
649
650 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
651 assert_eq!(chunk, Poll::Ready(None));
652 }
653
654 #[test]
655 fn chunked_bytes_exact_chunk_size() {
656 let data = b"12345".to_vec();
657 let mut stream = ChunkedBytes::new(data, 5);
658
659 let waker = noop_waker();
660 let mut ctx = Context::from_waker(&waker);
661
662 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
664 assert_eq!(chunk, Poll::Ready(Some(b"12345".to_vec())));
665
666 let chunk = Pin::new(&mut stream).poll_next(&mut ctx);
668 assert_eq!(chunk, Poll::Ready(None));
669 }
670
671 #[test]
672 fn cancel_aware_stream_propagates_items() {
673 let inner = asupersync::stream::iter(vec![1, 2, 3]);
674 let cx = Cx::for_testing();
675 let mut stream = CancelAwareStream::new(inner, cx);
676
677 let waker = noop_waker();
678 let mut ctx = Context::from_waker(&waker);
679
680 assert_eq!(
681 Pin::new(&mut stream).poll_next(&mut ctx),
682 Poll::Ready(Some(1))
683 );
684 assert_eq!(
685 Pin::new(&mut stream).poll_next(&mut ctx),
686 Poll::Ready(Some(2))
687 );
688 assert_eq!(
689 Pin::new(&mut stream).poll_next(&mut ctx),
690 Poll::Ready(Some(3))
691 );
692 assert_eq!(Pin::new(&mut stream).poll_next(&mut ctx), Poll::Ready(None));
693
694 assert!(!stream.is_cancelled());
695 }
696
697 #[test]
698 fn stream_error_display() {
699 let err = StreamError::Cancelled;
700 assert_eq!(format!("{err}"), "stream cancelled");
701
702 let err = StreamError::BufferFull;
703 assert_eq!(format!("{err}"), "stream buffer full");
704
705 let io_err = io::Error::new(io::ErrorKind::NotFound, "file not found");
706 let err = StreamError::Io(io_err);
707 assert!(format!("{err}").contains("streaming I/O error"));
708 }
709
710 #[test]
715 fn stream_file_adds_accept_ranges_header() {
716 let temp_dir = std::env::temp_dir();
718 let test_file = temp_dir.join("test_stream_accept_ranges.txt");
719 std::fs::write(&test_file, b"Hello, streaming world!").unwrap();
720
721 let cx = Cx::for_testing();
722 let response = fastapi_core::Response::stream_file(&test_file, cx, b"text/plain").unwrap();
723
724 let accept_ranges = response
725 .headers()
726 .iter()
727 .find(|(name, _)| name == "accept-ranges")
728 .map(|(_, value)| String::from_utf8_lossy(value).to_string());
729
730 assert_eq!(accept_ranges, Some("bytes".to_string()));
731
732 let _ = std::fs::remove_file(test_file);
734 }
735
736 #[test]
737 fn stream_file_range_returns_206() {
738 use crate::range::ByteRange;
739
740 let temp_dir = std::env::temp_dir();
742 let test_file = temp_dir.join("test_stream_range_206.txt");
743 std::fs::write(&test_file, b"0123456789ABCDEF").unwrap();
744
745 let cx = Cx::for_testing();
746 let range = ByteRange::new(0, 4); let response = fastapi_core::Response::stream_file_range(
748 &test_file,
749 range,
750 16, cx,
752 b"text/plain",
753 )
754 .unwrap();
755
756 assert_eq!(response.status().as_u16(), 206);
758
759 let _ = std::fs::remove_file(test_file);
761 }
762
763 #[test]
764 fn stream_file_range_sets_content_range_header() {
765 use crate::range::ByteRange;
766
767 let temp_dir = std::env::temp_dir();
769 let test_file = temp_dir.join("test_stream_content_range.txt");
770 std::fs::write(&test_file, b"0123456789ABCDEF").unwrap();
771
772 let cx = Cx::for_testing();
773 let range = ByteRange::new(5, 9); let response = fastapi_core::Response::stream_file_range(
775 &test_file,
776 range,
777 16, cx,
779 b"text/plain",
780 )
781 .unwrap();
782
783 let content_range = response
784 .headers()
785 .iter()
786 .find(|(name, _)| name == "content-range")
787 .map(|(_, value)| String::from_utf8_lossy(value).to_string());
788
789 assert_eq!(content_range, Some("bytes 5-9/16".to_string()));
790
791 let _ = std::fs::remove_file(test_file);
793 }
794
795 #[test]
796 fn stream_file_range_sets_content_length_header() {
797 use crate::range::ByteRange;
798
799 let temp_dir = std::env::temp_dir();
801 let test_file = temp_dir.join("test_stream_content_length.txt");
802 std::fs::write(&test_file, b"0123456789ABCDEF").unwrap();
803
804 let cx = Cx::for_testing();
805 let range = ByteRange::new(0, 99); let response = fastapi_core::Response::stream_file_range(
807 &test_file,
808 range,
809 16, cx,
811 b"text/plain",
812 )
813 .unwrap();
814
815 let content_length = response
816 .headers()
817 .iter()
818 .find(|(name, _)| name == "content-length")
819 .map(|(_, value)| String::from_utf8_lossy(value).to_string());
820
821 assert_eq!(content_length, Some("100".to_string()));
823
824 let _ = std::fs::remove_file(test_file);
826 }
827
828 #[test]
833 fn stream_large_response_in_chunks() {
834 const TARGET_SIZE: usize = 10 * 1024 * 1024; const CHUNK_SIZE: usize = 64 * 1024; let data: Vec<u8> = (0..TARGET_SIZE).map(|i| (i % 256) as u8).collect();
840 let mut stream = ChunkedBytes::new(data.clone(), CHUNK_SIZE);
841
842 let waker = noop_waker();
843 let mut ctx = Context::from_waker(&waker);
844
845 let mut total_received = 0usize;
846 let mut chunk_count = 0usize;
847
848 loop {
849 match Pin::new(&mut stream).poll_next(&mut ctx) {
850 Poll::Ready(Some(chunk)) => {
851 if total_received + CHUNK_SIZE <= TARGET_SIZE {
853 assert_eq!(
854 chunk.len(),
855 CHUNK_SIZE,
856 "Non-final chunks should be {CHUNK_SIZE} bytes"
857 );
858 }
859 total_received += chunk.len();
860 chunk_count += 1;
861 }
862 Poll::Ready(None) => break,
863 Poll::Pending => panic!("ChunkedBytes should never return Pending"),
864 }
865 }
866
867 assert_eq!(total_received, TARGET_SIZE, "Should receive all 10MB");
868 let expected_chunks = TARGET_SIZE.div_ceil(CHUNK_SIZE);
869 assert_eq!(
870 chunk_count, expected_chunks,
871 "Should have correct number of chunks"
872 );
873 }
874
875 #[test]
876 fn cancel_aware_stream_stops_on_cancellation() {
877 let data = vec![1, 2, 3, 4, 5];
879 let inner = asupersync::stream::iter(data);
880 let cx = Cx::for_testing();
881
882 cx.set_cancel_requested(true);
884
885 let mut stream = CancelAwareStream::new(inner, cx);
886
887 let waker = noop_waker();
888 let mut ctx = Context::from_waker(&waker);
889
890 assert_eq!(Pin::new(&mut stream).poll_next(&mut ctx), Poll::Ready(None));
892 assert!(
893 stream.is_cancelled(),
894 "Stream should be marked as cancelled"
895 );
896 }
897
898 #[test]
899 fn file_stream_reads_complete_file() {
900 let temp_dir = std::env::temp_dir();
902 let test_file = temp_dir.join("test_file_stream_complete.bin");
903
904 const FILE_SIZE: usize = 256 * 1024;
906 let data: Vec<u8> = (0..FILE_SIZE).map(|i| (i % 256) as u8).collect();
907 std::fs::write(&test_file, &data).unwrap();
908
909 let cx = Cx::for_testing();
910 let config = StreamConfig::new().with_chunk_size(32 * 1024);
911 let mut stream = FileStream::open(&test_file, cx, config).unwrap();
912
913 let waker = noop_waker();
914 let mut ctx = Context::from_waker(&waker);
915
916 let mut total_received = 0usize;
917 let mut received_data = Vec::new();
918
919 loop {
920 match Pin::new(&mut stream).poll_next(&mut ctx) {
921 Poll::Ready(Some(chunk)) => {
922 total_received += chunk.len();
923 received_data.extend(chunk);
924 }
925 Poll::Ready(None) => break,
926 Poll::Pending => {
927 }
929 }
930 }
931
932 assert_eq!(total_received, FILE_SIZE, "Should receive complete file");
933 assert_eq!(received_data, data, "Data should match original");
934
935 let _ = std::fs::remove_file(test_file);
937 }
938
939 #[test]
940 fn chunked_bytes_total_size_is_correct() {
941 const SIZE: usize = 1024 * 100; let data: Vec<u8> = vec![0u8; SIZE];
944 let stream = ChunkedBytes::new(data, 1024);
945
946 assert_eq!(
947 stream.total_size(),
948 SIZE,
949 "Total size should be known upfront"
950 );
951 }
952
953 #[test]
954 fn file_stream_size_is_known_via_remaining() {
955 let temp_dir = std::env::temp_dir();
957 let test_file = temp_dir.join("test_file_size_known.txt");
958
959 const FILE_SIZE: usize = 12345;
960 let data: Vec<u8> = vec![b'X'; FILE_SIZE];
961 std::fs::write(&test_file, &data).unwrap();
962
963 let cx = Cx::for_testing();
964 let config = StreamConfig::default();
965 let stream = FileStream::open(&test_file, cx, config).unwrap();
966
967 assert_eq!(
969 stream.remaining(),
970 FILE_SIZE as u64,
971 "File size should be known via remaining()"
972 );
973
974 let _ = std::fs::remove_file(test_file);
976 }
977}