use crate::{ProgressEntry, ProgressListener, Pusher};
use bytes::Bytes;
use memmap2::MmapMut;
#[cfg(unix)]
use std::os::unix::io::AsRawFd;
#[cfg(windows)]
use std::os::windows::io::AsRawHandle;
pub struct MmapFilePusher {
mmap: MmapMut,
sync_all: bool,
listener: Option<ProgressListener>,
}
impl std::fmt::Debug for MmapFilePusher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MmapFilePusher")
.field("sync_all", &self.sync_all)
.finish_non_exhaustive()
}
}
impl MmapFilePusher {
pub async fn new(file: &tokio::fs::File, size: u64, sync_all: bool) -> std::io::Result<Self> {
file.set_len(size).await?;
#[cfg(unix)]
let raw = file.as_raw_fd();
#[cfg(windows)]
let raw = file.as_raw_handle() as usize;
let mmap = tokio::task::spawn_blocking(move || unsafe {
#[cfg(unix)]
let desc = raw;
#[cfg(windows)]
let desc = raw as std::os::windows::io::RawHandle;
MmapMut::map_mut(desc)
})
.await
.map_err(std::io::Error::other)??;
Ok(Self {
mmap,
sync_all,
listener: None,
})
}
}
impl Pusher for MmapFilePusher {
type Error = std::io::Error;
fn set_listener(&mut self, cb: ProgressListener) {
self.listener = Some(cb);
}
fn push(&mut self, range: &ProgressEntry, bytes: Bytes) -> Result<(), (Self::Error, Bytes)> {
#[allow(clippy::cast_possible_truncation)]
self.mmap[range.start as usize..range.end as usize].copy_from_slice(&bytes);
if let Some(l) = &mut self.listener {
l(range.clone());
}
Ok(())
}
fn flush(&mut self) -> Result<(), Self::Error> {
if self.sync_all {
self.mmap.flush()
} else {
self.mmap.flush_async()
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
use std::sync::{Arc, Mutex};
use std::{fs::File, io::Read, vec::Vec};
use tempfile::NamedTempFile;
#[tokio::test]
async fn test_rand_file_pusher() {
let temp_file = NamedTempFile::new().unwrap();
let file_path = temp_file.path();
let mut pusher = MmapFilePusher::new(&temp_file.reopen().unwrap().into(), 10, false)
.await
.unwrap();
let data = b"234";
let range = 2..5;
pusher.push(&range, data[..].into()).unwrap();
pusher.flush().unwrap();
let mut file_content = Vec::new();
File::open(file_path)
.unwrap()
.read_to_end(&mut file_content)
.unwrap();
assert_eq!(file_content, b"\0\x00234\0\0\0\0\0");
}
#[tokio::test]
async fn test_debug_and_sync_all_flush() {
let temp_file = NamedTempFile::new().unwrap();
let file_path = temp_file.path().to_path_buf();
let mut pusher = MmapFilePusher::new(&temp_file.reopen().unwrap().into(), 10, true)
.await
.unwrap();
let _ = format!("{pusher:?}");
pusher.push(&(2..5), b"234"[..].into()).unwrap();
pusher.flush().unwrap();
let mut file_content = Vec::new();
File::open(&file_path)
.unwrap()
.read_to_end(&mut file_content)
.unwrap();
assert_eq!(file_content, b"\0\x00234\0\0\0\0\0");
}
#[tokio::test]
async fn test_listener_invoked_on_push() {
let temp_file = NamedTempFile::new().unwrap();
let mut pusher = MmapFilePusher::new(&temp_file.reopen().unwrap().into(), 10, false)
.await
.unwrap();
let seen = Arc::new(Mutex::new(None::<ProgressEntry>));
let seen2 = seen.clone();
pusher.set_listener(Box::new(move |r| {
*seen2.lock().unwrap() = Some(r);
}));
pusher.push(&(2..5), b"234"[..].into()).unwrap();
assert_eq!(*seen.lock().unwrap(), Some(2..5));
}
#[tokio::test]
async fn test_mmap_pusher_zero_length_file() {
let temp_file = NamedTempFile::new().unwrap();
let result = MmapFilePusher::new(&temp_file.reopen().unwrap().into(), 0, false).await;
assert!(
result.is_ok(),
"zero-length download should construct an mmap pusher, got: {result:?}"
);
}
}