use crate::error::Result;
use crate::event::{Event, EventKind};
use crate::halt::Halt;
use crate::sector::SectorSource;
use crossbeam_channel::{Receiver, Sender, bounded};
use std::thread::JoinHandle;
pub type EventFn = Box<dyn Fn(Event) + Send + 'static>;
const PREFETCH_CHANNEL_DEPTH: usize = 2;
const SECTOR_ALIGNMENT: u16 = 3;
pub type Batch = std::result::Result<Vec<u8>, std::io::Error>;
pub struct PrefetchedSectorSource {
rx: Receiver<Batch>,
recycle_tx: Sender<Vec<u8>>,
producer: Option<JoinHandle<()>>,
total_sectors: u32,
}
impl PrefetchedSectorSource {
pub fn new<S>(
reader: S,
extents: Vec<crate::disc::Extent>,
batch_sectors: u16,
halt: Option<Halt>,
) -> Result<Self>
where
S: SectorSource + Send + 'static,
{
Self::new_with_events(reader, extents, batch_sectors, halt, None)
}
pub fn new_with_events<S>(
mut reader: S,
extents: Vec<crate::disc::Extent>,
batch_sectors: u16,
halt: Option<Halt>,
event_fn: Option<EventFn>,
) -> Result<Self>
where
S: SectorSource + Send + 'static,
{
if batch_sectors == 0 {
return Err(crate::error::Error::IoError {
source: std::io::Error::from(std::io::ErrorKind::InvalidInput),
});
}
let total_sectors: u32 = extents
.iter()
.map(|e| e.sector_count as u64)
.sum::<u64>()
.min(u32::MAX as u64) as u32;
let bytes_total_extents: u64 = extents.iter().map(|e| e.sector_count as u64 * 2048).sum();
let (tx, rx) = bounded::<Batch>(PREFETCH_CHANNEL_DEPTH);
let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(PREFETCH_CHANNEL_DEPTH + 1);
let batch_bytes = batch_sectors as usize * 2048;
for _ in 0..(PREFETCH_CHANNEL_DEPTH + 1) {
let _ = recycle_tx.send(vec![0u8; batch_bytes]);
}
let producer = std::thread::Builder::new()
.name("freemkv-prefetch".into())
.spawn(move || {
let mut ext_idx = 0usize;
let mut offset: u32 = 0;
let mut bytes_read_total: u64 = 0;
while ext_idx < extents.len() {
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
return;
}
let extent = &extents[ext_idx];
let remaining = extent.sector_count.saturating_sub(offset);
if remaining == 0 {
ext_idx += 1;
offset = 0;
continue;
}
if remaining % SECTOR_ALIGNMENT as u32 != 0
&& remaining < SECTOR_ALIGNMENT as u32
{
let _ = tx.send(Err(crate::error::Error::ExtentNotUnitAligned.into()));
return;
}
let mut sectors = remaining.min(batch_sectors as u32) as u16;
if sectors >= SECTOR_ALIGNMENT {
sectors -= sectors % SECTOR_ALIGNMENT;
} else {
sectors = SECTOR_ALIGNMENT;
}
let bytes = sectors as usize * 2048;
let mut buf = match recycle_rx.recv() {
Ok(b) => b,
Err(_) => return, };
if bytes <= buf.capacity() {
debug_assert!(bytes <= buf.capacity(), "set_len exceeds capacity");
unsafe { buf.set_len(bytes) };
} else {
buf.resize(bytes, 0);
}
let lba = extent.start_lba.saturating_add(offset);
match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) {
Ok(n) => {
if n % 2048 != 0 {
let _ =
tx.send(Err(crate::error::Error::ExtentNotUnitAligned.into()));
return;
}
let sectors_read = (n / 2048) as u32;
buf.truncate(n);
bytes_read_total = bytes_read_total.saturating_add(n as u64);
if let Some(ref f) = event_fn {
f(Event {
kind: EventKind::BytesRead {
bytes: bytes_read_total,
total: bytes_total_extents,
},
});
}
if tx.send(Ok(buf)).is_err() {
return; }
if sectors_read == 0 {
return;
}
offset = offset.saturating_add(sectors_read);
}
Err(e) => {
let _ = tx.send(Err(e.into()));
return;
}
}
}
})
.map_err(|e| crate::error::Error::IoError { source: e })?;
Ok(Self {
rx,
recycle_tx,
producer: Some(producer),
total_sectors,
})
}
pub fn into_channels(self) -> (Receiver<Batch>, Sender<Vec<u8>>, PrefetchShell) {
let total = self.total_sectors;
let me = std::mem::ManuallyDrop::new(self);
let producer = unsafe { std::ptr::read(&me.producer) };
let rx = unsafe { std::ptr::read(&me.rx) };
let recycle = unsafe { std::ptr::read(&me.recycle_tx) };
(rx, recycle, PrefetchShell { producer, total })
}
}
pub struct PrefetchShell {
producer: Option<JoinHandle<()>>,
#[allow(dead_code)]
total: u32,
}
impl Drop for PrefetchShell {
fn drop(&mut self) {
if let Some(h) = self.producer.take() {
let _ = h.join();
}
}
}
impl Drop for PrefetchedSectorSource {
fn drop(&mut self) {
if let Some(h) = self.producer.take() {
let _ = h.join();
}
}
}
impl SectorSource for PrefetchedSectorSource {
fn capacity_sectors(&self) -> u32 {
self.total_sectors
}
fn read_sectors(
&mut self,
_lba: u32,
_count: u16,
buf: &mut [u8],
_recovery: bool,
) -> Result<usize> {
match self.rx.recv() {
Ok(Ok(filled)) => {
if filled.len() > buf.len() {
return Err(crate::error::Error::IoError {
source: std::io::Error::from(std::io::ErrorKind::InvalidInput),
});
}
let n = filled.len();
buf[..n].copy_from_slice(&filled[..n]);
let _ = self.recycle_tx.send(filled);
Ok(n)
}
Ok(Err(e)) => Err(crate::error::Error::IoError { source: e }),
Err(_) => Ok(0),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::disc::Extent;
use crate::error::Result;
use std::sync::mpsc;
use std::time::Duration;
struct EndlessZeroSource;
impl SectorSource for EndlessZeroSource {
fn read_sectors(
&mut self,
_lba: u32,
count: u16,
buf: &mut [u8],
_recovery: bool,
) -> Result<usize> {
let bytes = count as usize * 2048;
buf[..bytes].fill(0);
Ok(bytes)
}
}
struct PatternSource {
capacity: u32,
}
impl SectorSource for PatternSource {
fn capacity_sectors(&self) -> u32 {
self.capacity
}
fn read_sectors(
&mut self,
lba: u32,
count: u16,
buf: &mut [u8],
_recovery: bool,
) -> Result<usize> {
let bytes = count as usize * 2048;
for s in 0..count as usize {
let base = s * 2048;
let tag = (lba.wrapping_add(s as u32) & 0xff) as u8;
for b in &mut buf[base..base + 2048] {
*b = tag;
}
}
Ok(bytes)
}
}
struct ShortFirstSource {
capacity: u32,
first: bool,
}
impl SectorSource for ShortFirstSource {
fn capacity_sectors(&self) -> u32 {
self.capacity
}
fn read_sectors(
&mut self,
lba: u32,
count: u16,
buf: &mut [u8],
_recovery: bool,
) -> Result<usize> {
let give = if self.first {
self.first = false;
SECTOR_ALIGNMENT.min(count)
} else {
count
};
let bytes = give as usize * 2048;
for s in 0..give as usize {
let base = s * 2048;
let tag = (lba.wrapping_add(s as u32) & 0xff) as u8;
for b in &mut buf[base..base + 2048] {
*b = tag;
}
}
Ok(bytes)
}
}
fn big_extent() -> Vec<Extent> {
vec![Extent {
start_lba: 0,
sector_count: u32::MAX,
}]
}
fn within<F: FnOnce() + Send + 'static>(secs: u64, f: F) {
let (done_tx, done_rx) = bounded::<()>(1);
std::thread::spawn(move || {
f();
let _ = done_tx.send(());
});
assert!(
done_rx
.recv_timeout(std::time::Duration::from_secs(secs))
.is_ok(),
"operation did not complete within {secs}s (deadlock)"
);
}
fn with_watchdog<F>(timeout: Duration, f: F)
where
F: FnOnce() + Send + 'static,
{
let (done_tx, done_rx) = mpsc::channel::<()>();
let h = std::thread::spawn(move || {
f();
let _ = done_tx.send(());
});
match done_rx.recv_timeout(timeout) {
Ok(()) => {
let _ = h.join();
}
Err(_) => panic!("watchdog timeout — likely deadlock/hang in prefetch read path"),
}
}
#[test]
fn into_channels_drop_releases_producer() {
within(10, || {
let src = PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 3, None)
.expect("spawn");
let (rx, recycle_tx, shell) = src.into_channels();
drop(rx);
drop(recycle_tx);
drop(shell);
});
}
#[test]
fn halt_releases_producer() {
within(10, || {
let halt = Halt::new();
let src =
PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 3, Some(halt.clone()))
.expect("spawn");
let (rx, recycle_tx, shell) = src.into_channels();
let drainer = std::thread::spawn(move || {
while let Ok(item) = rx.recv() {
if let Ok(buf) = item {
let _ = recycle_tx.send(buf);
}
}
});
halt.cancel();
drop(shell);
let _ = drainer.join();
});
}
#[test]
fn zero_batch_rejected() {
let err = PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 0, None);
assert!(err.is_err(), "zero batch_sectors must be rejected");
}
#[test]
fn direct_reads_past_pool_depth_do_not_deadlock() {
with_watchdog(Duration::from_secs(10), || {
let extents = vec![Extent {
start_lba: 0,
sector_count: 24,
}];
let src = PatternSource { capacity: 24 };
let mut pf = PrefetchedSectorSource::new(src, extents, 3, None).expect("spawn");
let mut buf = vec![0u8; 3 * 2048];
let mut total = 0usize;
for _ in 0..16 {
let n = pf.read_sectors(0, 3, &mut buf, false).unwrap();
if n == 0 {
break; }
total += n;
}
assert_eq!(total, 24 * 2048, "all 24 sectors should be drained");
});
}
#[test]
fn non_multiple_of_three_extent_errors_on_tail() {
with_watchdog(Duration::from_secs(10), || {
let extents = vec![Extent {
start_lba: 100,
sector_count: 8,
}];
let src = PatternSource { capacity: 200 };
let mut pf = PrefetchedSectorSource::new(src, extents, 3, None).expect("spawn");
let mut buf = vec![0u8; 3 * 2048];
let n0 = pf.read_sectors(0, 3, &mut buf, false).unwrap();
assert_eq!(n0, 3 * 2048);
let n1 = pf.read_sectors(0, 3, &mut buf, false).unwrap();
assert_eq!(n1, 3 * 2048);
let err = pf.read_sectors(0, 3, &mut buf, false);
assert!(
err.is_err(),
"non-unit-aligned tail must error, got Ok({:?})",
err
);
});
}
#[test]
fn short_read_does_not_desync_stream() {
with_watchdog(Duration::from_secs(10), || {
let extents = vec![Extent {
start_lba: 0,
sector_count: 9,
}];
let src = ShortFirstSource {
capacity: 9,
first: true,
};
let mut pf = PrefetchedSectorSource::new(src, extents, 9, None).expect("spawn");
let mut buf = vec![0u8; 9 * 2048];
let mut total = 0usize;
for _ in 0..16 {
let n = pf.read_sectors(0, 9, &mut buf, false).unwrap();
if n == 0 {
break;
}
total += n;
}
assert_eq!(
total,
9 * 2048,
"short read must not drop sectors; all 9 must be delivered"
);
});
}
}