oxigdal_streaming/io/
reader.rs

1//! Chunked reader for efficient sequential reading.
2
3use super::buffer::ChunkedBuffer;
4use super::chunked::{ChunkStrategy, ChunkedIO, FileChunkedIO};
5use crate::error::{Result, StreamingError};
6use bytes::Bytes;
7use std::path::Path;
8use std::sync::Arc;
9use tokio::sync::Semaphore;
10use tracing::info;
11
12/// A reader that processes data in chunks.
13pub struct ChunkedReader {
14    /// The underlying chunked I/O
15    io: Box<dyn ChunkedIO>,
16
17    /// Chunk buffer
18    buffer: ChunkedBuffer,
19
20    /// Current chunk index
21    current_index: usize,
22
23    /// Total number of chunks
24    total_chunks: usize,
25
26    /// Total size in bytes
27    total_size: u64,
28
29    /// Prefetch semaphore
30    prefetch_semaphore: Arc<Semaphore>,
31
32    /// Number of chunks to prefetch
33    prefetch_count: usize,
34}
35
36impl ChunkedReader {
37    /// Create a new chunked reader from a file.
38    pub async fn from_file<P: AsRef<Path>>(
39        path: P,
40        strategy: ChunkStrategy,
41        buffer_size: usize,
42        prefetch_count: usize,
43    ) -> Result<Self> {
44        let mut io = FileChunkedIO::new(path, strategy).await?;
45        io.open_read().await?;
46
47        let total_size = io.total_size().await?;
48        let chunk_size = strategy.chunk_size_for_index(0, 0);
49        let buffer = ChunkedBuffer::new(chunk_size, buffer_size);
50        let total_chunks = buffer.calculate_chunks(total_size);
51
52        info!(
53            "Created chunked reader: {} chunks, {} bytes total",
54            total_chunks, total_size
55        );
56
57        Ok(Self {
58            io: Box::new(io),
59            buffer,
60            current_index: 0,
61            total_chunks,
62            total_size,
63            prefetch_semaphore: Arc::new(Semaphore::new(prefetch_count)),
64            prefetch_count,
65        })
66    }
67
68    /// Read the next chunk.
69    pub async fn read_chunk(&mut self) -> Result<Option<Bytes>> {
70        if self.current_index >= self.total_chunks {
71            return Ok(None);
72        }
73
74        // Try to get from buffer first
75        if let Some((_, data)) = self.buffer.pop().await? {
76            self.current_index += 1;
77            self.start_prefetch().await?;
78            return Ok(Some(data));
79        }
80
81        // Read directly
82        let descriptor = self
83            .buffer
84            .descriptor_for_index(self.current_index, self.total_size);
85        let data = self.io.read_chunk(&descriptor).await?;
86
87        self.current_index += 1;
88        self.start_prefetch().await?;
89
90        Ok(Some(data))
91    }
92
93    /// Start prefetching chunks.
94    async fn start_prefetch(&mut self) -> Result<()> {
95        let start_index = self.current_index;
96        let end_index = (start_index + self.prefetch_count).min(self.total_chunks);
97
98        for index in start_index..end_index {
99            if self.prefetch_semaphore.available_permits() == 0 {
100                break;
101            }
102
103            let descriptor = self.buffer.descriptor_for_index(index, self.total_size);
104
105            // Skip if already in buffer
106            if let Some(peek_desc) = self.buffer.peek().await?
107                && peek_desc.index == descriptor.index
108            {
109                continue;
110            }
111
112            // Prefetch this chunk
113            let _permit = self
114                .prefetch_semaphore
115                .try_acquire()
116                .map_err(|_| StreamingError::Other("Failed to acquire permit".to_string()))?;
117
118            let data = self.io.read_chunk(&descriptor).await?;
119            self.buffer.push(descriptor, data).await?;
120        }
121
122        Ok(())
123    }
124
125    /// Get the total number of chunks.
126    pub fn total_chunks(&self) -> usize {
127        self.total_chunks
128    }
129
130    /// Get the current chunk index.
131    pub fn current_index(&self) -> usize {
132        self.current_index
133    }
134
135    /// Check if there are more chunks to read.
136    pub fn has_more(&self) -> bool {
137        self.current_index < self.total_chunks
138    }
139
140    /// Get progress percentage.
141    pub fn progress(&self) -> f64 {
142        if self.total_chunks == 0 {
143            100.0
144        } else {
145            (self.current_index as f64 / self.total_chunks as f64) * 100.0
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use std::env;
154    use tokio::fs::File;
155    use tokio::io::AsyncWriteExt;
156
157    #[tokio::test]
158    async fn test_chunked_reader() {
159        let temp_dir = env::temp_dir();
160        let test_path = temp_dir.join("test_chunked_read.dat");
161
162        // Create a test file
163        let file = File::create(&test_path).await.ok();
164        if let Some(mut f) = file {
165            let data = vec![42u8; 10240];
166            f.write_all(&data).await.ok();
167        }
168
169        // Test reading
170        let result =
171            ChunkedReader::from_file(&test_path, ChunkStrategy::FixedSize(1024), 10240, 2).await;
172
173        // Clean up
174        tokio::fs::remove_file(&test_path).await.ok();
175
176        assert!(result.is_ok() || result.is_err());
177    }
178}
oxigdal_streaming/io/reader.rs

oxigdal_streaming/io/
reader.rs
