use std::io::Error;
use std::mem;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamIoMode {
Auto,
Serial,
Parallel,
}
#[derive(Debug, Clone)]
pub struct StreamOptions {
pub block_size: usize,
pub io_mode: StreamIoMode,
}
impl Default for StreamOptions {
fn default() -> Self {
Self {
block_size: 4 * 1024 * 1024,
io_mode: StreamIoMode::Auto,
}
}
}
impl StreamOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_block_size(mut self, size: usize) -> Self {
const MIN_BLOCK_SIZE_BYTES: usize = 1024;
const MAX_BLOCK_SIZE_BYTES: usize = 16 * 1024 * 1024;
self.block_size = size.clamp(MIN_BLOCK_SIZE_BYTES, MAX_BLOCK_SIZE_BYTES);
self
}
pub fn with_io_mode(mut self, mode: StreamIoMode) -> Self {
self.io_mode = mode;
self
}
}
#[derive(Debug)]
pub struct StreamError {
pub shard_index: usize,
pub kind: StreamErrorKind,
}
#[derive(Debug)]
pub enum StreamErrorKind {
Read(std::io::Error),
Write(std::io::Error),
Codec(crate::Error),
}
impl core::fmt::Display for StreamError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match &self.kind {
StreamErrorKind::Read(e) => {
write!(f, "read error on shard {}: {}", self.shard_index, e)
}
StreamErrorKind::Write(e) => {
write!(f, "write error on shard {}: {}", self.shard_index, e)
}
StreamErrorKind::Codec(e) => {
write!(f, "codec error on shard {}: {}", self.shard_index, e)
}
}
}
}
impl std::error::Error for StreamError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match &self.kind {
StreamErrorKind::Read(e) => Some(e),
StreamErrorKind::Write(e) => Some(e),
StreamErrorKind::Codec(e) => Some(e),
}
}
}
impl StreamError {
fn read(shard_index: usize, e: std::io::Error) -> Self {
Self {
shard_index,
kind: StreamErrorKind::Read(e),
}
}
fn write(shard_index: usize, e: std::io::Error) -> Self {
Self {
shard_index,
kind: StreamErrorKind::Write(e),
}
}
fn codec(shard_index: usize, e: crate::Error) -> Self {
Self {
shard_index,
kind: StreamErrorKind::Codec(e),
}
}
}
fn take_stream_error(
first_error: &std::sync::Mutex<Option<StreamError>>,
fallback_message: &'static str,
) -> StreamError {
match first_error.lock() {
Ok(mut guard) => guard
.take()
.unwrap_or_else(|| StreamError::read(0, Error::other(fallback_message))),
Err(poisoned) => {
let mut guard = poisoned.into_inner();
guard
.take()
.unwrap_or_else(|| StreamError::read(0, Error::other(fallback_message)))
}
}
}
fn use_parallel_stream_io(options: &StreamOptions, stream_count: usize) -> bool {
match options.io_mode {
StreamIoMode::Serial => false,
StreamIoMode::Parallel => true,
StreamIoMode::Auto => use_parallel_stream_io_auto(options.block_size, stream_count),
}
}
fn use_parallel_stream_io_auto(block_size: usize, stream_count: usize) -> bool {
if stream_count < 2 || block_size < 256 * 1024 {
return false;
}
if stream_count <= 6 && block_size <= 1024 * 1024 {
return false;
}
block_size >= 4 * 1024 * 1024 && stream_count >= 10
}
fn read_block_with_mode<R: std::io::Read + Send>(
readers: &mut [R],
buffers: &mut [Vec<u8>],
max_len: usize,
use_parallel_io: bool,
read_lengths: &mut Vec<(usize, usize)>,
) -> Result<(bool, usize), StreamError> {
if use_parallel_io {
read_block_par(readers, buffers, max_len)
} else {
read_block(readers, buffers, max_len, read_lengths)
}
}
fn write_block_with_mode<W: std::io::Write + Send>(
writers: &mut [W],
buffers: &[Vec<u8>],
len: usize,
shard_offset: usize,
use_parallel_io: bool,
) -> Result<(), StreamError> {
if use_parallel_io {
write_block_par(writers, buffers, len, shard_offset)
} else {
write_block(writers, buffers, len, shard_offset)
}
}
fn read_block<R: std::io::Read>(
readers: &mut [R],
buffers: &mut [Vec<u8>],
max_len: usize,
read_lengths: &mut Vec<(usize, usize)>,
) -> Result<(bool, usize), StreamError> {
read_lengths.clear();
for (i, (reader, buf)) in readers.iter_mut().zip(buffers.iter_mut()).enumerate() {
let total = read_one_stream(reader, buf, max_len).map_err(|e| StreamError::read(i, e))?;
read_lengths.push((i, total));
}
let actual_len = read_lengths
.iter()
.map(|(_, total)| *total)
.max()
.unwrap_or(0);
zero_pad_short_buffers(buffers, read_lengths, actual_len);
Ok((actual_len == 0, actual_len))
}
fn prepare_read_buffer(buf: &mut Vec<u8>, max_len: usize) {
match buf.len().cmp(&max_len) {
core::cmp::Ordering::Less => buf.resize(max_len, 0),
core::cmp::Ordering::Greater => buf.truncate(max_len),
core::cmp::Ordering::Equal => {}
}
}
fn read_one_stream<R: std::io::Read>(
reader: &mut R,
buf: &mut Vec<u8>,
max_len: usize,
) -> Result<usize, std::io::Error> {
prepare_read_buffer(buf, max_len);
let mut total = 0;
while total < max_len {
match reader.read(&mut buf[total..]) {
Ok(0) => break,
Ok(n) => total += n,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => return Err(e),
}
}
Ok(total)
}
fn zero_pad_short_buffers(
buffers: &mut [Vec<u8>],
read_lengths: &[(usize, usize)],
actual_len: usize,
) {
for &(i, total) in read_lengths {
if total < actual_len {
buffers[i][total..actual_len].fill(0);
}
}
}
fn read_present_cursors_with_mode(
shards: &mut [std::io::Cursor<Vec<u8>>],
buffers: &mut [Vec<u8>],
present: &[bool],
present_indices: &[usize],
max_len: usize,
use_parallel_io: bool,
read_lengths: &mut Vec<(usize, usize)>,
) -> Result<(bool, usize), StreamError> {
read_lengths.clear();
if use_parallel_io {
use rayon::prelude::*;
*read_lengths = shards
.par_iter_mut()
.zip(buffers.par_iter_mut())
.enumerate()
.filter_map(|(i, item)| present[i].then_some((i, item)))
.map(|(i, (shard, buf))| {
let total =
read_one_stream(shard, buf, max_len).map_err(|e| StreamError::read(i, e))?;
buf.truncate(total);
Ok::<_, StreamError>((i, total))
})
.collect::<Result<Vec<_>, _>>()?;
} else {
for &i in present_indices {
let total = read_one_stream(&mut shards[i], &mut buffers[i], max_len)
.map_err(|e| StreamError::read(i, e))?;
buffers[i].truncate(total);
read_lengths.push((i, total));
}
}
let actual_len = read_lengths
.iter()
.map(|(_, total)| *total)
.max()
.unwrap_or(0);
if actual_len != 0 {
for &idx in present_indices {
buffers[idx].resize(actual_len, 0);
}
}
Ok((actual_len == 0, actual_len))
}
fn write_block<W: std::io::Write>(
writers: &mut [W],
buffers: &[Vec<u8>],
len: usize,
shard_offset: usize,
) -> Result<(), StreamError> {
for (i, (writer, buf)) in writers.iter_mut().zip(buffers.iter()).enumerate() {
let mut written = 0;
while written < len {
match writer.write(&buf[written..len]) {
Ok(0) => {
return Err(StreamError::write(
shard_offset + i,
std::io::Error::new(std::io::ErrorKind::WriteZero, "write returned 0"),
));
}
Ok(n) => written += n,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => return Err(StreamError::write(shard_offset + i, e)),
}
}
}
Ok(())
}
fn read_block_par<R: std::io::Read + Send>(
readers: &mut [R],
buffers: &mut [Vec<u8>],
max_len: usize,
) -> Result<(bool, usize), StreamError> {
use rayon::prelude::*;
let read_lengths: Vec<(usize, usize)> = readers
.par_iter_mut()
.zip(buffers.par_iter_mut())
.enumerate()
.map(|(i, (reader, buf))| {
read_one_stream(reader, buf, max_len)
.map(|total| (i, total))
.map_err(|e| StreamError::read(i, e))
})
.collect::<Result<Vec<_>, _>>()?;
let actual_len = read_lengths
.iter()
.map(|(_, total)| *total)
.max()
.unwrap_or(0);
zero_pad_short_buffers(buffers, &read_lengths, actual_len);
Ok((actual_len == 0, actual_len))
}
fn write_block_par<W: std::io::Write + Send>(
writers: &mut [W],
buffers: &[Vec<u8>],
len: usize,
shard_offset: usize,
) -> Result<(), StreamError> {
use rayon::prelude::*;
let first_error: std::sync::Mutex<Option<StreamError>> = std::sync::Mutex::new(None);
writers
.par_iter_mut()
.zip(buffers.par_iter())
.enumerate()
.try_for_each(|(i, (writer, buf))| {
let mut written = 0;
while written < len {
match writer.write(&buf[written..len]) {
Ok(0) => {
if let Ok(mut guard) = first_error.lock()
&& guard.is_none()
{
*guard = Some(StreamError::write(
shard_offset + i,
std::io::Error::new(
std::io::ErrorKind::WriteZero,
"write returned 0",
),
));
}
return Err(());
}
Ok(n) => written += n,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => {
if let Ok(mut guard) = first_error.lock()
&& guard.is_none()
{
*guard = Some(StreamError::write(shard_offset + i, e));
}
return Err(());
}
}
}
Ok(())
})
.map_err(|()| {
take_stream_error(
&first_error,
"parallel stream writer error was not reported",
)
})
}
impl super::ReedSolomon<crate::galois_8::Field> {
pub fn encode_stream(
&self,
data: &mut [impl std::io::Read + Send],
parity: &mut [impl std::io::Write + Send],
options: &StreamOptions,
) -> Result<(), StreamError> {
let block_size = options.block_size;
let data_count = self.data_shard_count;
let parity_count = self.parity_shard_count;
let use_parallel_read = use_parallel_stream_io(options, data_count);
let use_parallel_write = use_parallel_stream_io(options, parity_count);
debug_assert_eq!(data.len(), data_count);
debug_assert_eq!(parity.len(), parity_count);
let mut data_bufs: Vec<Vec<u8>> = (0..data_count)
.map(|_| Vec::with_capacity(block_size))
.collect();
let mut parity_bufs: Vec<Vec<u8>> = (0..parity_count)
.map(|_| Vec::with_capacity(block_size))
.collect();
let mut read_lengths = Vec::with_capacity(data_count);
loop {
let (all_eof, actual_len) = read_block_with_mode(
data,
&mut data_bufs,
block_size,
use_parallel_read,
&mut read_lengths,
)?;
if all_eof {
break;
}
for buf in parity_bufs.iter_mut() {
buf.resize(actual_len, 0);
}
let data_refs: Vec<&[u8]> = data_bufs.iter().map(|b| &b[..actual_len]).collect();
let mut parity_refs: Vec<&mut [u8]> = parity_bufs
.iter_mut()
.map(|b| &mut b[..actual_len])
.collect();
self.encode_sep_par(&data_refs, &mut parity_refs)
.map_err(|e| StreamError::codec(0, e))?;
write_block_with_mode(
parity,
&parity_bufs,
actual_len,
data_count,
use_parallel_write,
)?;
}
Ok(())
}
pub fn verify_stream(
&self,
shards: &mut [impl std::io::Read + Send],
options: &StreamOptions,
) -> Result<bool, StreamError> {
let block_size = options.block_size;
let total = self.total_shard_count;
let use_parallel_read = use_parallel_stream_io(options, total);
debug_assert_eq!(shards.len(), total);
let mut bufs: Vec<Vec<u8>> = (0..total).map(|_| Vec::with_capacity(block_size)).collect();
let mut read_lengths = Vec::with_capacity(total);
loop {
let (all_eof, actual_len) = read_block_with_mode(
shards,
&mut bufs,
block_size,
use_parallel_read,
&mut read_lengths,
)?;
if all_eof {
break;
}
let refs: Vec<&[u8]> = bufs.iter().map(|b| &b[..actual_len]).collect();
let valid = self
.verify_par(&refs)
.map_err(|e| StreamError::codec(0, e))?;
if !valid {
return Ok(false);
}
}
Ok(true)
}
pub fn reconstruct_stream(
&self,
shards: &mut [std::io::Cursor<Vec<u8>>],
options: &StreamOptions,
) -> Result<(), StreamError> {
let block_size = options.block_size;
let total = self.total_shard_count;
debug_assert_eq!(shards.len(), total);
let mut present = vec![false; total];
for (i, shard) in shards.iter().enumerate() {
present[i] = !shard.get_ref().is_empty();
}
let missing_count = present.iter().filter(|&&p| !p).count();
if missing_count > self.parity_shard_count {
return Err(StreamError::codec(0, crate::Error::TooFewShardsPresent));
}
let present_count = total - missing_count;
let use_parallel_read = use_parallel_stream_io(options, present_count);
let present_indices: Vec<usize> = present
.iter()
.enumerate()
.filter_map(|(i, is_present)| is_present.then_some(i))
.collect();
let missing_indices: Vec<usize> = present
.iter()
.enumerate()
.filter_map(|(i, is_present)| (!is_present).then_some(i))
.collect();
let mut bufs: Vec<Vec<u8>> = (0..total)
.map(|i| {
if present[i] {
Vec::with_capacity(block_size)
} else {
Vec::new()
}
})
.collect();
let mut reconstruct_bufs: Vec<Option<Vec<u8>>> = (0..total).map(|_| None).collect();
let mut read_lengths = Vec::with_capacity(total);
loop {
let (all_eof, actual_len) = read_present_cursors_with_mode(
shards,
&mut bufs,
&present,
&present_indices,
block_size,
use_parallel_read,
&mut read_lengths,
)?;
if all_eof {
break;
}
for &idx in &present_indices {
bufs[idx].resize(actual_len, 0);
reconstruct_bufs[idx] = Some(mem::take(&mut bufs[idx]));
}
for &idx in &missing_indices {
reconstruct_bufs[idx] = None;
}
self.reconstruct(&mut reconstruct_bufs)
.map_err(|e| StreamError::codec(0, e))?;
for &idx in &missing_indices {
let recovered = reconstruct_bufs[idx]
.as_ref()
.expect("missing shard buffer");
shards[idx]
.get_mut()
.extend_from_slice(&recovered[..actual_len]);
}
for idx in 0..total {
if let Some(buf) = reconstruct_bufs[idx].take() {
bufs[idx] = buf;
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::galois_8::ReedSolomon;
fn make_codec(data: usize, parity: usize) -> ReedSolomon {
ReedSolomon::new(data, parity).unwrap()
}
fn random_data(len: usize) -> Vec<u8> {
(0..len)
.map(|i| (i.wrapping_mul(73).wrapping_add(17)) as u8)
.collect()
}
#[test]
fn test_encode_stream_basic() {
let rs = make_codec(3, 2);
let shard_size = 4096;
let d0 = random_data(shard_size);
let d1 = random_data(shard_size);
let d2 = random_data(shard_size);
let mut readers: Vec<&[u8]> = vec![d0.as_slice(), d1.as_slice(), d2.as_slice()];
let mut writers: Vec<Vec<u8>> = vec![Vec::new(), Vec::new()];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
assert_eq!(writers[0].len(), shard_size);
assert_eq!(writers[1].len(), shard_size);
let all: Vec<&[u8]> = vec![
d0.as_slice(),
d1.as_slice(),
d2.as_slice(),
writers[0].as_slice(),
writers[1].as_slice(),
];
assert!(rs.verify(&all).unwrap());
}
#[test]
fn test_encode_stream_multi_block() {
let rs = make_codec(3, 2);
let total_size = 10 * 1024; let block_size = 4096;
let d0 = random_data(total_size);
let d1 = random_data(total_size);
let d2 = random_data(total_size);
let mut readers: Vec<&[u8]> = vec![d0.as_slice(), d1.as_slice(), d2.as_slice()];
let mut writers: Vec<Vec<u8>> = vec![Vec::new(), Vec::new()];
let opts = StreamOptions::new().with_block_size(block_size);
rs.encode_stream(&mut readers, &mut writers, &opts).unwrap();
assert_eq!(writers[0].len(), total_size);
assert_eq!(writers[1].len(), total_size);
}
#[test]
fn test_encode_stream_empty() {
let rs = make_codec(3, 2);
let empty: Vec<&[u8]> = vec![&[], &[], &[]];
let mut readers = empty.clone();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(), Vec::new()];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
assert!(writers[0].is_empty());
assert!(writers[1].is_empty());
}
#[test]
fn test_encode_stream_unequal_lengths() {
let rs = make_codec(3, 2);
let d0 = random_data(1000);
let d1 = random_data(500);
let d2 = random_data(800);
let mut readers: Vec<&[u8]> = vec![d0.as_slice(), d1.as_slice(), d2.as_slice()];
let mut writers: Vec<Vec<u8>> = vec![Vec::new(), Vec::new()];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
assert_eq!(writers[0].len(), 1000);
assert_eq!(writers[1].len(), 1000);
}
#[test]
fn test_encode_stream_10x4() {
let rs = make_codec(10, 4);
let shard_size = 64 * 1024;
let data: Vec<Vec<u8>> = (0..10).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 4];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
let mut all: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
for w in &writers {
all.push(w.as_slice());
}
assert!(rs.verify(&all).unwrap());
}
#[test]
fn test_verify_stream_valid() {
let rs = make_codec(3, 2);
let shard_size = 4096;
let d = vec![random_data(shard_size); 3];
let mut readers: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
let mut all_data: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
for w in &writers {
all_data.push(w.as_slice());
}
let mut all_readers: Vec<&[u8]> = all_data;
assert!(
rs.verify_stream(&mut all_readers, &StreamOptions::default())
.unwrap()
);
}
#[test]
fn test_verify_stream_corrupted() {
let rs = make_codec(3, 2);
let shard_size = 4096;
let d = vec![random_data(shard_size); 3];
let mut readers: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
writers[0][0] ^= 0xFF;
let mut all_data: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
for w in &writers {
all_data.push(w.as_slice());
}
let mut all_readers: Vec<&[u8]> = all_data;
assert!(
!rs.verify_stream(&mut all_readers, &StreamOptions::default())
.unwrap()
);
}
#[test]
fn test_reconstruct_stream_single_missing() {
let rs = make_codec(3, 2);
let shard_size = 4096;
let d: Vec<Vec<u8>> = (0..3).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
let mut parity_writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut parity_writers, &StreamOptions::default())
.unwrap();
let mut shards: Vec<std::io::Cursor<Vec<u8>>> = vec![
std::io::Cursor::new(Vec::new()), std::io::Cursor::new(d[1].clone()),
std::io::Cursor::new(d[2].clone()),
std::io::Cursor::new(parity_writers[0].clone()),
std::io::Cursor::new(parity_writers[1].clone()),
];
rs.reconstruct_stream(&mut shards, &StreamOptions::default())
.unwrap();
assert_eq!(shards[0].get_ref(), &d[0]);
}
#[test]
fn test_reconstruct_non_streaming() {
let rs = make_codec(3, 2);
let shard_size = 4096;
let d: Vec<Vec<u8>> = (0..3).map(|_| random_data(shard_size)).collect();
let data_refs: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
let mut p0 = vec![0u8; shard_size];
let mut p1 = vec![0u8; shard_size];
let mut parity_refs: Vec<&mut [u8]> = vec![&mut p0, &mut p1];
rs.encode_sep(&data_refs, &mut parity_refs).unwrap();
let mut shards: Vec<Option<Vec<u8>>> = vec![
None,
Some(d[1].clone()),
Some(d[2].clone()),
Some(p0.clone()),
Some(p1.clone()),
];
rs.reconstruct(&mut shards).unwrap();
assert_eq!(shards[0].as_ref().unwrap(), &d[0]);
}
#[test]
fn test_reconstruct_stream_basic() {
let rs = make_codec(3, 2);
let shard_size = 64;
let d: Vec<Vec<u8>> = (0..3).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = d.iter().map(|s| s.as_slice()).collect();
let mut parity_writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut parity_writers, &StreamOptions::default())
.unwrap();
let all: Vec<&[u8]> = vec![
d[0].as_slice(),
d[1].as_slice(),
d[2].as_slice(),
parity_writers[0].as_slice(),
parity_writers[1].as_slice(),
];
assert!(rs.verify(&all).unwrap());
let mut shards: Vec<std::io::Cursor<Vec<u8>>> = vec![
std::io::Cursor::new(Vec::new()), std::io::Cursor::new(d[1].clone()),
std::io::Cursor::new(d[2].clone()),
std::io::Cursor::new(parity_writers[0].clone()),
std::io::Cursor::new(parity_writers[1].clone()),
];
rs.reconstruct_stream(&mut shards, &StreamOptions::default())
.unwrap();
let recovered = shards[0].get_ref();
assert_eq!(
recovered.len(),
d[0].len(),
"recovered len {} != expected len {}",
recovered.len(),
d[0].len()
);
assert_eq!(
recovered,
&d[0],
"recovered: {:?}, expected: {:?}",
&recovered[..8],
&d[0][..8]
);
}
#[test]
fn test_stream_options_builder() {
let opts = StreamOptions::new().with_block_size(1024 * 1024);
assert_eq!(opts.block_size, 1024 * 1024);
assert_eq!(opts.io_mode, StreamIoMode::Auto);
let opts = StreamOptions::new().with_block_size(100);
assert_eq!(opts.block_size, 1024);
let opts = StreamOptions::new().with_io_mode(StreamIoMode::Serial);
assert_eq!(opts.io_mode, StreamIoMode::Serial);
}
#[test]
fn test_stream_io_auto_decision_thresholds() {
let opts = StreamOptions::new()
.with_block_size(64 * 1024)
.with_io_mode(StreamIoMode::Auto);
assert!(!use_parallel_stream_io(&opts, 14));
let opts = StreamOptions::new()
.with_block_size(1024 * 1024)
.with_io_mode(StreamIoMode::Auto);
assert!(!use_parallel_stream_io(&opts, 6));
let opts = StreamOptions::new()
.with_block_size(4 * 1024 * 1024)
.with_io_mode(StreamIoMode::Auto);
assert!(use_parallel_stream_io(&opts, 10));
let opts = StreamOptions::new()
.with_block_size(64 * 1024)
.with_io_mode(StreamIoMode::Parallel);
assert!(use_parallel_stream_io(&opts, 1));
}
#[test]
fn test_stream_error_display() {
let e = StreamError::read(
3,
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "eof"),
);
let s = format!("{e}");
assert!(s.contains("shard 3"));
assert!(s.contains("eof"));
let e = StreamError::codec(0, crate::Error::TooFewShardsPresent);
let s = format!("{e}");
assert!(s.contains("codec"));
}
#[test]
fn test_encode_stream_concurrent() {
let rs = make_codec(4, 2);
let shard_size = 8 * 1024;
let data: Vec<Vec<u8>> = (0..4).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
let mut all: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
for w in &writers {
all.push(w.as_slice());
}
assert!(rs.verify(&all).unwrap());
}
#[test]
fn test_verify_stream_concurrent() {
let rs = make_codec(4, 2);
let shard_size = 8 * 1024;
let data: Vec<Vec<u8>> = (0..4).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
let mut all: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
for w in &writers {
all.push(w.as_slice());
}
let mut all_readers: Vec<&[u8]> = all;
assert!(
rs.verify_stream(&mut all_readers, &StreamOptions::default())
.unwrap()
);
let mut corrupted = writers[0].clone();
corrupted[0] ^= 0xFF;
let mut all_corrupt: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
all_corrupt.push(corrupted.as_slice());
all_corrupt.push(writers[1].as_slice());
assert!(
!rs.verify_stream(&mut all_corrupt, &StreamOptions::default())
.unwrap()
);
}
#[test]
fn test_reconstruct_stream_concurrent() {
let rs = make_codec(4, 2);
let shard_size = 8 * 1024;
let data: Vec<Vec<u8>> = (0..4).map(|_| random_data(shard_size)).collect();
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut parity_writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut parity_writers, &StreamOptions::default())
.unwrap();
let mut shards: Vec<std::io::Cursor<Vec<u8>>> = vec![
std::io::Cursor::new(data[0].clone()),
std::io::Cursor::new(Vec::new()), std::io::Cursor::new(data[2].clone()),
std::io::Cursor::new(data[3].clone()),
std::io::Cursor::new(Vec::new()), std::io::Cursor::new(parity_writers[1].clone()),
];
rs.reconstruct_stream(&mut shards, &StreamOptions::default())
.unwrap();
assert_eq!(shards[1].get_ref(), &data[1]);
}
#[test]
fn test_concurrent_stream_large_blocks() {
let rs = make_codec(10, 4);
let total_size = 1024 * 1024; let block_size = 256 * 1024;
let data: Vec<Vec<u8>> = (0..10).map(|_| random_data(total_size)).collect();
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 4];
let opts = StreamOptions::new().with_block_size(block_size);
rs.encode_stream(&mut readers, &mut writers, &opts).unwrap();
let mut all: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
for w in &writers {
all.push(w.as_slice());
}
assert!(rs.verify(&all).unwrap());
let mut shards: Vec<std::io::Cursor<Vec<u8>>> = Vec::new();
for d in &data {
shards.push(std::io::Cursor::new(d.clone()));
}
shards[0] = std::io::Cursor::new(Vec::new());
shards[5] = std::io::Cursor::new(Vec::new());
for w in &writers {
shards.push(std::io::Cursor::new(w.clone()));
}
rs.reconstruct_stream(&mut shards, &StreamOptions::default())
.unwrap();
assert_eq!(shards[0].get_ref(), &data[0]);
assert_eq!(shards[5].get_ref(), &data[5]);
}
#[test]
fn test_encode_stream_zero_length() {
let rs = ReedSolomon::new(3, 2).unwrap();
let mut readers: Vec<&[u8]> = vec![b""; 3];
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &StreamOptions::default())
.unwrap();
assert!(writers.iter().all(|w| w.is_empty()));
}
#[test]
fn test_encode_stream_single_byte_block() {
let rs = ReedSolomon::new(2, 1).unwrap();
let d0 = vec![0xABu8; 4];
let d1 = vec![0xCDu8; 4];
let mut readers: Vec<&[u8]> = vec![d0.as_slice(), d1.as_slice()];
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 1];
let opts = StreamOptions::new().with_block_size(1);
rs.encode_stream(&mut readers, &mut writers, &opts).unwrap();
let mut all: Vec<&[u8]> = vec![d0.as_slice(), d1.as_slice()];
for w in &writers {
all.push(w.as_slice());
}
assert!(rs.verify(&all).unwrap());
}
#[test]
fn test_stream_io_modes_encode_verify_match() {
let rs = ReedSolomon::new(4, 2).unwrap();
let shard_size = 32 * 1024;
let data: Vec<Vec<u8>> = (0..4).map(|_| random_data(shard_size)).collect();
let mut expected_parity = Vec::new();
for mode in [
StreamIoMode::Auto,
StreamIoMode::Serial,
StreamIoMode::Parallel,
] {
let opts = StreamOptions::new()
.with_block_size(8 * 1024)
.with_io_mode(mode);
let mut readers: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut writers: Vec<Vec<u8>> = vec![Vec::new(); 2];
rs.encode_stream(&mut readers, &mut writers, &opts).unwrap();
if expected_parity.is_empty() {
expected_parity = writers.clone();
} else {
assert_eq!(writers, expected_parity);
}
let mut all: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
for parity in &writers {
all.push(parity.as_slice());
}
let mut verify_readers = all;
assert!(rs.verify_stream(&mut verify_readers, &opts).unwrap());
}
}
#[test]
fn test_reconstruct_stream_minimum_present() {
let rs = ReedSolomon::new(3, 2).unwrap();
let shard_len = 16usize;
let data: Vec<Vec<u8>> = (0..3).map(|i| vec![i as u8 + 1; shard_len]).collect();
let refs: Vec<&[u8]> = data.iter().map(|d| d.as_slice()).collect();
let mut parity = vec![vec![0u8; shard_len]; 2];
let mut par_refs: Vec<&mut [u8]> = parity.iter_mut().map(|p| &mut p[..]).collect();
rs.encode_sep(&refs, &mut par_refs).unwrap();
let mut shards: Vec<std::io::Cursor<Vec<u8>>> = vec![
std::io::Cursor::new(data[0].clone()),
std::io::Cursor::new(Vec::new()),
std::io::Cursor::new(Vec::new()),
std::io::Cursor::new(parity[0].clone()),
std::io::Cursor::new(parity[1].clone()),
];
rs.reconstruct_stream(&mut shards, &StreamOptions::default())
.unwrap();
assert_eq!(shards[1].get_ref(), &data[1], "shard 1 mismatch");
assert_eq!(shards[2].get_ref(), &data[2], "shard 2 mismatch");
}
}