parallel_processor/memory_fs/file/

reader.rs

use crate::memory_fs::file::internal::{FileChunk, MemoryFileInternal, OpenMode};
use crate::utils::vec_reader::VecReaderInner;
use parking_lot::lock_api::ArcRwLockReadGuard;
use parking_lot::{RawRwLock, RwLock};
use std::io;
use std::io::{ErrorKind, Read, Seek, SeekFrom};
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use super::writer::FileWriter;

const MIN_UNBUFFERED_READ: usize = 2048;
const FILE_READ_BUFFER_SIZE: usize = 4096;

#[derive(Clone)]
pub struct FileRangeReference {
    file: Arc<RwLock<MemoryFileInternal>>,
    start_chunk: usize,
    start_chunk_offset: usize,
    bytes_count: usize,
}

impl FileRangeReference {
    pub unsafe fn copy_to_unsync(&self, other: &FileWriter) {
        let file = self.file.read();
        let mut chunk_index = self.start_chunk;
        let mut chunk_offset = self.start_chunk_offset;
        let mut written_bytes = 0;

        while written_bytes < self.bytes_count {
            let underlying_file = file.get_underlying_file();
            let chunk = file.get_chunk(chunk_index);
            let chunk = chunk.read();
            let to_copy = (chunk.get_length() - chunk_offset).min(self.bytes_count - written_bytes);

            other.write_all_unsync_from_readfn(
                |buffer| {
                    let amount = chunk
                        .read_at(underlying_file, chunk_offset as u64, buffer)
                        .unwrap();
                    chunk_offset += amount;
                    amount
                },
                to_copy,
            );

            written_bytes += to_copy;
            chunk_index += 1;
            chunk_offset = 0;
        }
    }
}

pub struct FileReader {
    path: PathBuf,
    file: Arc<RwLock<MemoryFileInternal>>,
    current_chunk_ref: Option<ArcRwLockReadGuard<RawRwLock, FileChunk>>,
    current_chunk_index: usize,
    chunks_count: usize,
    current_position: usize,
    current_len: usize,
    current_file_position: usize,
    is_on_disk: bool,
    buffer: Option<VecReaderInner>,
    buffer_position: usize,
}

unsafe impl Sync for FileReader {}
unsafe impl Send for FileReader {}

impl Clone for FileReader {
    fn clone(&self) -> Self {
        Self {
            path: self.path.clone(),
            file: self.file.clone(),
            current_chunk_ref: self
                .current_chunk_ref
                .as_ref()
                .map(|c| ArcRwLockReadGuard::rwlock(&c).read_arc()),
            current_chunk_index: self.current_chunk_index,
            chunks_count: self.chunks_count,
            current_position: self.current_position,
            current_len: self.current_len,
            current_file_position: self.current_file_position,
            is_on_disk: self.is_on_disk,
            buffer: self.buffer.clone(),
            buffer_position: self.buffer_position,
        }
    }
}

impl FileReader {
    fn set_chunk_info(&mut self, index: usize) {
        let file = self.file.read();

        let chunk = file.get_chunk(index);
        let chunk_guard = chunk.read_arc();

        self.current_position = 0;
        self.buffer_position = 0;
        self.current_len = chunk_guard.get_length();
        self.is_on_disk = chunk_guard.is_on_disk();
        self.current_chunk_ref = Some(chunk_guard);

        if self.is_on_disk && self.buffer.is_none() {
            self.buffer = Some(VecReaderInner::new(FILE_READ_BUFFER_SIZE))
        }

        self.buffer.as_mut().map(|b| b.reset());
    }

    pub fn open(path: impl AsRef<Path>) -> Option<Self> {
        let file = match MemoryFileInternal::retrieve_reference(&path) {
            None => MemoryFileInternal::create_from_fs(&path)?,
            Some(x) => x,
        };

        let mut file_lock = file.write();

        file_lock.open(OpenMode::Read).unwrap();
        let chunks_count = file_lock.get_chunks_count();
        drop(file_lock);

        let mut reader = Self {
            path: path.as_ref().into(),
            file,
            current_chunk_ref: None,
            current_chunk_index: 0,
            chunks_count,
            current_position: 0,
            current_len: 0,
            current_file_position: 0,
            is_on_disk: false,
            buffer: None,
            buffer_position: 0,
        };

        if reader.chunks_count > 0 {
            reader.set_chunk_info(0);
        }

        Some(reader)
    }

    pub fn get_unique_file_id(&self) -> usize {
        self.file.data_ptr() as usize
    }

    pub fn total_file_size(&self) -> usize {
        self.file.read().len()
    }

    pub fn get_file_path(&self) -> &Path {
        &self.path
    }

    pub fn close_and_remove(self, remove_fs: bool) -> bool {
        MemoryFileInternal::delete(self.path, remove_fs)
    }

    pub fn get_range_reference(&self, file_range: Range<u64>) -> FileRangeReference {
        let file = self.file.read();
        let mut chunk_index = 0;
        let mut chunk_offset = 0;
        let mut start = file_range.start;

        while start > 0 {
            let chunk = file.get_chunk(chunk_index);
            let chunk = chunk.read();
            let len = chunk.get_length() as u64;
            if start < len {
                chunk_offset = start as usize;
                break;
            }
            start -= len;
            chunk_index += 1;
        }

        FileRangeReference {
            file: self.file.clone(),
            start_chunk: chunk_index,
            start_chunk_offset: chunk_offset,
            bytes_count: file_range.end as usize - file_range.start as usize,
        }
    }

    // pub fn get_typed_chunks_mut<T>(&mut self) -> Option<impl Iterator<Item = &mut [T]>> {
    //     todo!();
    //     Some((0..1).into_iter().map(|_| &mut [0, 1][..]))
    // }
}

impl Read for FileReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let mut bytes_written = 0;

        while bytes_written != buf.len() {
            if self.current_len == 0 {
                self.current_chunk_index += 1;
                if self.current_chunk_index >= self.chunks_count {
                    // End of file
                    return Ok(bytes_written);
                }
                self.set_chunk_info(self.current_chunk_index);
            }

            let file = self.file.read();
            let underlying_file = file.get_underlying_file();

            let copyable_bytes = buf.len() - bytes_written;

            let copyable_bytes = if self.is_on_disk && copyable_bytes < MIN_UNBUFFERED_READ {
                let buffer = self.buffer.as_mut().unwrap();
                let copyable_bytes = buffer.read_bytes(&mut buf[bytes_written..], |buffer| {
                    let read_amount = self.current_chunk_ref.as_ref().unwrap().read_at(
                        underlying_file,
                        self.buffer_position as u64,
                        buffer,
                    )?;
                    self.buffer_position += read_amount;
                    Ok(read_amount)
                });
                copyable_bytes
            } else {
                let copyable_bytes = self.current_chunk_ref.as_ref().unwrap().read_at(
                    underlying_file,
                    self.current_position as u64,
                    &mut buf[bytes_written..],
                )?;
                self.buffer
                    .as_mut()
                    .map(|b| self.buffer_position += b.discard(copyable_bytes));
                copyable_bytes
            };

            self.current_position += copyable_bytes;
            self.current_len -= copyable_bytes;
            self.current_file_position += copyable_bytes;
            bytes_written += copyable_bytes;
        }

        Ok(bytes_written)
    }

    #[inline(always)]
    fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
        match self.read(buf) {
            Ok(count) => {
                if count == buf.len() {
                    Ok(())
                } else {
                    Err(io::Error::new(
                        io::ErrorKind::Other,
                        "Unexpected error while reading",
                    ))
                }
            }
            Err(err) => Err(err),
        }
    }
}

impl Seek for FileReader {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match pos {
            SeekFrom::Start(mut offset) => {
                let mut chunk_idx = 0;

                let file = self.file.read();
                while chunk_idx < self.chunks_count {
                    let len = file.get_chunk(chunk_idx).read().get_length();
                    if offset < (len as u64) {
                        break;
                    }
                    chunk_idx += 1;
                    offset -= len as u64;
                }

                if chunk_idx == self.chunks_count {
                    return Err(std::io::Error::new(
                        ErrorKind::UnexpectedEof,
                        "Unexpected eof",
                    ));
                }

                self.current_chunk_index = chunk_idx;
                drop(file);
                self.set_chunk_info(chunk_idx);

                self.current_position += offset as usize;
                self.buffer
                    .as_mut()
                    .map(|b| self.buffer_position += b.discard(offset as usize));
                self.current_len -= offset as usize;
                self.current_file_position = offset as usize;

                return Ok(offset);
            }
            SeekFrom::Current(offset) => {
                assert!(offset >= 0); // < 0  not supported yet
                let mut offset = offset as usize;
                loop {
                    let clen_offset = offset.min(self.current_len);
                    offset -= clen_offset;
                    self.current_len -= clen_offset;
                    self.current_position += clen_offset;
                    self.buffer
                        .as_mut()
                        .map(|b| self.buffer_position += b.discard(clen_offset as usize));
                    self.current_file_position += clen_offset;

                    if offset == 0 {
                        break Ok(self.current_file_position as u64);
                    }

                    if self.current_chunk_index >= self.chunks_count - 1 {
                        break Err(std::io::Error::new(
                            ErrorKind::UnexpectedEof,
                            "Unexpected eof",
                        ));
                    }

                    self.current_chunk_index += 1;
                    self.set_chunk_info(self.current_chunk_index);
                }
            }
            _ => {
                unimplemented!()
            }
        }
    }

    fn stream_position(&mut self) -> io::Result<u64> {
        let mut position = 0;

        let file_read = self.file.read();

        for i in 0..self.current_chunk_index {
            position += file_read.get_chunk(i).read().get_length();
        }

        position += file_read
            .get_chunk(self.current_chunk_index)
            .read()
            .get_length()
            - self.current_len;

        Ok(position as u64)
    }
}