fast_pull/file/
pusher.rs

1extern crate std;
2use crate::{ProgressEntry, RandPusher, SeqPusher, Total};
3use mmap_io::{MemoryMappedFile, MmapIoError, MmapMode, flush::FlushPolicy};
4use std::{boxed::Box, collections::VecDeque, path::Path};
5use tokio::{
6    fs::{File, OpenOptions},
7    io::{self, AsyncSeekExt, AsyncWriteExt, BufWriter, SeekFrom},
8};
9
10#[derive(thiserror::Error, Debug)]
11pub enum FilePusherError {
12    #[error(transparent)]
13    MmapIo(#[from] MmapIoError),
14    #[error(transparent)]
15    TokioIo(#[from] io::Error),
16}
17
18#[derive(Debug)]
19pub struct SeqFilePusher {
20    buffer: BufWriter<File>,
21}
22impl SeqFilePusher {
23    pub fn new(file: File, buffer_size: usize) -> Self {
24        Self {
25            buffer: BufWriter::with_capacity(buffer_size, file),
26        }
27    }
28}
29impl SeqPusher for SeqFilePusher {
30    type Error = FilePusherError;
31    async fn push(&mut self, content: &[u8]) -> Result<(), Self::Error> {
32        Ok(self.buffer.write_all(content).await?)
33    }
34    async fn flush(&mut self) -> Result<(), Self::Error> {
35        Ok(self.buffer.flush().await?)
36    }
37}
38
39#[derive(Debug)]
40pub struct RandFilePusherMmap {
41    mmap: MemoryMappedFile,
42    downloaded: usize,
43    buffer_size: usize,
44}
45impl RandFilePusherMmap {
46    pub async fn new(
47        path: impl AsRef<Path>,
48        size: u64,
49        buffer_size: usize,
50    ) -> Result<Self, FilePusherError> {
51        let mmap_builder = MemoryMappedFile::builder(&path)
52            .mode(MmapMode::ReadWrite)
53            .flush_policy(FlushPolicy::Manual);
54        Ok(Self {
55            mmap: if path.as_ref().try_exists()? {
56                OpenOptions::new()
57                    .write(true)
58                    .open(path)
59                    .await?
60                    .set_len(size)
61                    .await?;
62                mmap_builder.open()
63            } else {
64                mmap_builder.size(size).create()
65            }?,
66            downloaded: 0,
67            buffer_size,
68        })
69    }
70}
71impl RandPusher for RandFilePusherMmap {
72    type Error = FilePusherError;
73    async fn push(&mut self, range: ProgressEntry, bytes: &[u8]) -> Result<(), Self::Error> {
74        self.mmap
75            .as_slice_mut(range.start, range.total())?
76            .as_mut()
77            .copy_from_slice(bytes);
78        self.downloaded += bytes.len();
79        if self.downloaded >= self.buffer_size {
80            self.mmap.flush_async().await?;
81            self.downloaded = 0;
82        }
83        Ok(())
84    }
85    async fn flush(&mut self) -> Result<(), Self::Error> {
86        self.mmap.flush_async().await?;
87        Ok(())
88    }
89}
90
91#[derive(Debug)]
92pub struct RandFilePusherStd {
93    buffer: BufWriter<File>,
94    cache: VecDeque<(u64, Box<[u8]>)>,
95    p: u64,
96    cache_size: usize,
97    buffer_size: usize,
98}
99impl RandFilePusherStd {
100    pub async fn new(file: File, size: u64, buffer_size: usize) -> Result<Self, FilePusherError> {
101        file.set_len(size).await?;
102        Ok(Self {
103            buffer: BufWriter::with_capacity(buffer_size, file),
104            cache: VecDeque::new(),
105            p: 0,
106            cache_size: 0,
107            buffer_size,
108        })
109    }
110}
111impl RandPusher for RandFilePusherStd {
112    type Error = FilePusherError;
113    async fn push(&mut self, range: ProgressEntry, bytes: &[u8]) -> Result<(), Self::Error> {
114        let pos = self.cache.partition_point(|(i, _)| i < &range.start);
115        self.cache_size += bytes.len();
116        self.cache.insert(pos, (range.start, bytes.into()));
117        if self.cache_size >= self.buffer_size {
118            self.flush().await?;
119        }
120        Ok(())
121    }
122    async fn flush(&mut self) -> Result<(), Self::Error> {
123        while let Some((start, bytes)) = self.cache.front() {
124            let len = bytes.len();
125            self.cache_size -= len;
126            if *start != self.p {
127                self.buffer.seek(SeekFrom::Start(*start)).await?;
128                self.p = *start;
129            }
130            self.buffer.write_all(bytes).await?;
131            self.p += len as u64;
132            self.cache.pop_front();
133        }
134        self.buffer.flush().await?;
135        Ok(())
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use bytes::Bytes;
143    use std::vec::Vec;
144    use tempfile::NamedTempFile;
145    use tokio::io::AsyncReadExt;
146
147    #[tokio::test]
148    async fn test_seq_file_pusher() {
149        // 创建一个临时文件用于测试
150        let temp_file = NamedTempFile::new().unwrap();
151        let file_path = temp_file.path().to_path_buf();
152
153        // 初始化 SeqFilePusher
154        let mut pusher = SeqFilePusher::new(temp_file.reopen().unwrap().into(), 1024);
155
156        // 写入数据
157        let data1 = Bytes::from("Hello, ");
158        let data2 = Bytes::from("world!");
159        pusher.push(&data1).await.unwrap();
160        pusher.push(&data2).await.unwrap();
161        pusher.flush().await.unwrap();
162
163        // 验证文件内容
164        let mut file_content = Vec::new();
165        File::open(&file_path)
166            .await
167            .unwrap()
168            .read_to_end(&mut file_content)
169            .await
170            .unwrap();
171        assert_eq!(file_content, b"Hello, world!");
172    }
173
174    #[tokio::test]
175    async fn test_rand_file_pusher() {
176        // 创建一个临时文件用于测试
177        let temp_file = NamedTempFile::new().unwrap();
178        let file_path = temp_file.path();
179
180        // 初始化 RandFilePusher，假设文件大小为 10 字节
181        let mut pusher = RandFilePusherMmap::new(file_path, 10, 8 * 1024 * 1024)
182            .await
183            .unwrap();
184
185        // 写入数据
186        let data = Bytes::from("234");
187        let range = 2..5;
188        pusher.push(range, &data).await.unwrap();
189        pusher.flush().await.unwrap();
190
191        // 验证文件内容
192        let mut file_content = Vec::new();
193        File::open(&file_path)
194            .await
195            .unwrap()
196            .read_to_end(&mut file_content)
197            .await
198            .unwrap();
199        assert_eq!(file_content, b"\0\x00234\0\0\0\0\0");
200    }
201}
fast_pull/file/pusher.rs

fast_pull/file/
pusher.rs
