914 lines
38 KiB
Rust
914 lines
38 KiB
Rust
//! SharedBuffer is given a Reader and SharedBufferReader implements the Reader trait.
|
|
//! SharedBuffer reads ahead in the underlying file and saves the data.
|
|
//! SharedBufferReaders can be created for the buffer and independently keep track of each reader's read location.
|
|
//! The background reader keeps track of the progress of each client. After data has been read by all readers,
|
|
//! the buffer is recycled and reading ahead continues.
|
|
//! A primary use case is the underlying reader being decompressing a file, which can be computationally expensive.
|
|
//! The clients of SharedBufferReaders could be parallel instances which need access to the decompressed data.
|
|
use {
|
|
crate::waitable_condvar::WaitableCondvar,
|
|
log::*,
|
|
solana_measure::measure::Measure,
|
|
std::{
|
|
io::*,
|
|
sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
Arc, Mutex, RwLock,
|
|
},
|
|
thread::{Builder, JoinHandle},
|
|
time::Duration,
|
|
},
|
|
};
|
|
|
|
// tunable parameters:
|
|
// # bytes allocated and populated by reading ahead
|
|
const TOTAL_BUFFER_BUDGET_DEFAULT: usize = 2_000_000_000;
|
|
// data is read-ahead and saved in chunks of this many bytes
|
|
const CHUNK_SIZE_DEFAULT: usize = 100_000_000;
|
|
|
|
type OneSharedBuffer = Arc<Vec<u8>>;
|
|
|
|
struct SharedBufferInternal {
|
|
bg_reader_data: Arc<SharedBufferBgReader>,
|
|
|
|
bg_reader_join_handle: Mutex<Option<JoinHandle<()>>>,
|
|
|
|
// Keep track of the next read location per outstanding client.
|
|
// index is client's my_client_index.
|
|
// Value at index is index into buffers where that client is currently reading.
|
|
// Any buffer at index < min(clients) can be recycled or destroyed.
|
|
clients: RwLock<Vec<usize>>,
|
|
|
|
// unpacking callers read from 'data'. newly_read_data is transferred to 'data when 'data' is exhausted.
|
|
// This minimizes lock contention since bg file reader has to have almost constant write access.
|
|
data: RwLock<Vec<OneSharedBuffer>>,
|
|
|
|
// it is convenient to have one of these around
|
|
empty_buffer: OneSharedBuffer,
|
|
}
|
|
|
|
pub struct SharedBuffer {
|
|
instance: Arc<SharedBufferInternal>,
|
|
}
|
|
|
|
impl SharedBuffer {
|
|
pub fn new<T: 'static + Read + std::marker::Send>(reader: T) -> Self {
|
|
Self::new_with_sizes(TOTAL_BUFFER_BUDGET_DEFAULT, CHUNK_SIZE_DEFAULT, reader)
|
|
}
|
|
fn new_with_sizes<T: 'static + Read + std::marker::Send>(
|
|
total_buffer_budget: usize,
|
|
chunk_size: usize,
|
|
reader: T,
|
|
) -> Self {
|
|
assert!(total_buffer_budget > 0);
|
|
assert!(chunk_size > 0);
|
|
let instance = SharedBufferInternal {
|
|
bg_reader_data: Arc::new(SharedBufferBgReader::new()),
|
|
data: RwLock::new(vec![OneSharedBuffer::default()]), // initialize with 1 vector of empty data at data[0]
|
|
|
|
// default values
|
|
bg_reader_join_handle: Mutex::default(),
|
|
clients: RwLock::default(),
|
|
empty_buffer: OneSharedBuffer::default(),
|
|
};
|
|
let instance = Arc::new(instance);
|
|
let bg_reader_data = instance.bg_reader_data.clone();
|
|
|
|
let handle = Builder::new()
|
|
.name("solCompFileRead".to_string())
|
|
.spawn(move || {
|
|
// importantly, this thread does NOT hold a refcount on the arc of 'instance'
|
|
bg_reader_data.read_entire_file_in_bg(reader, total_buffer_budget, chunk_size);
|
|
});
|
|
*instance.bg_reader_join_handle.lock().unwrap() = Some(handle.unwrap());
|
|
Self { instance }
|
|
}
|
|
}
|
|
|
|
pub struct SharedBufferReader {
|
|
instance: Arc<SharedBufferInternal>,
|
|
my_client_index: usize,
|
|
// index in 'instance' of the current buffer this reader is reading from.
|
|
// The current buffer is referenced from 'current_data'.
|
|
// Until we exhaust this buffer, we don't need to get a lock to read from this.
|
|
current_buffer_index: usize,
|
|
// the index within current_data where we will next read
|
|
index_in_current_data: usize,
|
|
current_data: OneSharedBuffer,
|
|
|
|
// convenient to have access to
|
|
empty_buffer: OneSharedBuffer,
|
|
}
|
|
|
|
impl Drop for SharedBufferInternal {
|
|
fn drop(&mut self) {
|
|
if let Some(handle) = self.bg_reader_join_handle.lock().unwrap().take() {
|
|
self.bg_reader_data.stop.store(true, Ordering::Relaxed);
|
|
handle.join().unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for SharedBufferReader {
|
|
fn drop(&mut self) {
|
|
self.client_done_reading();
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct SharedBufferBgReader {
|
|
stop: AtomicBool,
|
|
// error encountered during read
|
|
error: RwLock<std::io::Result<usize>>,
|
|
// bg thread reads to 'newly_read_data' and signals
|
|
newly_read_data: RwLock<Vec<OneSharedBuffer>>,
|
|
// set when newly_read_data gets new data written to it and can be transferred
|
|
newly_read_data_signal: WaitableCondvar,
|
|
|
|
// currently available set of buffers for bg to read into
|
|
// during operation, this is exhausted as the bg reads ahead
|
|
// As all clients are done with an earlier buffer, it is recycled by being put back into this vec for the bg thread to pull out.
|
|
buffers: RwLock<Vec<OneSharedBuffer>>,
|
|
// signaled when a new buffer is added to buffers. This throttles the bg reading.
|
|
new_buffer_signal: WaitableCondvar,
|
|
|
|
bg_eof_reached: AtomicBool,
|
|
}
|
|
|
|
impl SharedBufferBgReader {
|
|
fn new() -> Self {
|
|
SharedBufferBgReader {
|
|
buffers: RwLock::new(vec![]),
|
|
error: RwLock::new(Ok(0)),
|
|
|
|
// easy defaults
|
|
stop: AtomicBool::new(false),
|
|
newly_read_data: RwLock::default(),
|
|
newly_read_data_signal: WaitableCondvar::default(),
|
|
new_buffer_signal: WaitableCondvar::default(),
|
|
bg_eof_reached: AtomicBool::default(),
|
|
}
|
|
}
|
|
|
|
fn default_wait_timeout() -> Duration {
|
|
Duration::from_millis(100) // short enough to be unnoticable in case of trouble, long enough for efficient waiting
|
|
}
|
|
fn wait_for_new_buffer(&self) -> bool {
|
|
self.new_buffer_signal
|
|
.wait_timeout(Self::default_wait_timeout())
|
|
}
|
|
fn num_buffers(total_buffer_budget: usize, chunk_size: usize) -> usize {
|
|
std::cmp::max(1, total_buffer_budget / chunk_size) // at least 1 buffer
|
|
}
|
|
fn set_error(&self, error: std::io::Error) {
|
|
*self.error.write().unwrap() = Err(error);
|
|
self.newly_read_data_signal.notify_all(); // any client waiting for new data needs to wake up and check for errors
|
|
}
|
|
|
|
// read ahead the entire file.
|
|
// This is governed by the supply of buffers.
|
|
// Buffers are likely limited to cap memory usage.
|
|
// A buffer is recycled after the last client finishes reading from it.
|
|
// When a buffer is available (initially or recycled), this code wakes up and reads into that buffer.
|
|
fn read_entire_file_in_bg<T: 'static + Read + std::marker::Send>(
|
|
&self,
|
|
mut reader: T,
|
|
total_buffer_budget: usize,
|
|
chunk_size: usize,
|
|
) {
|
|
let now = std::time::Instant::now();
|
|
let mut read_us = 0;
|
|
|
|
let mut max_bytes_read = 0;
|
|
let mut wait_us = 0;
|
|
let mut total_bytes = 0;
|
|
let mut error = SharedBufferReader::default_error();
|
|
let mut remaining_buffers_to_allocate = Self::num_buffers(total_buffer_budget, chunk_size);
|
|
loop {
|
|
if self.stop.load(Ordering::Relaxed) {
|
|
// unsure what error is most appropriate here.
|
|
// bg reader was told to stop. All clients need to see that as an error if they try to read.
|
|
self.set_error(std::io::Error::from(std::io::ErrorKind::TimedOut));
|
|
break;
|
|
}
|
|
let mut buffers = self.buffers.write().unwrap();
|
|
let buffer = buffers.pop();
|
|
drop(buffers);
|
|
let mut dest_data = if let Some(dest_data) = buffer {
|
|
// assert that this should not result in a vector copy
|
|
// These are internal buffers and should not be held by anyone else.
|
|
assert_eq!(Arc::strong_count(&dest_data), 1);
|
|
dest_data
|
|
} else if remaining_buffers_to_allocate > 0 {
|
|
// we still haven't allocated all the buffers we are allowed to allocate
|
|
remaining_buffers_to_allocate -= 1;
|
|
Arc::new(vec![0; chunk_size])
|
|
} else {
|
|
// nowhere to write, so wait for a buffer to become available
|
|
let mut wait_for_new_buffer = Measure::start("wait_for_new_buffer");
|
|
self.wait_for_new_buffer();
|
|
wait_for_new_buffer.stop();
|
|
wait_us += wait_for_new_buffer.as_us();
|
|
continue; // check stop, try to get a buffer again
|
|
};
|
|
let target = Arc::make_mut(&mut dest_data);
|
|
let dest_size = target.len();
|
|
|
|
let mut bytes_read = 0;
|
|
let mut eof = false;
|
|
let mut error_received = false;
|
|
|
|
while bytes_read < dest_size {
|
|
let mut time_read = Measure::start("read");
|
|
// Read from underlying reader into the remaining range in dest_data
|
|
// Note that this read takes less time (up to 2x) if we read into the same static buffer location each call.
|
|
// But, we have to copy the data out later, so we choose to pay the price at read time to put the data where it is useful.
|
|
let result = reader.read(&mut target[bytes_read..]);
|
|
time_read.stop();
|
|
read_us += time_read.as_us();
|
|
match result {
|
|
Ok(size) => {
|
|
if size == 0 {
|
|
eof = true;
|
|
break;
|
|
}
|
|
total_bytes += size;
|
|
max_bytes_read = std::cmp::max(max_bytes_read, size);
|
|
bytes_read += size;
|
|
// loop to read some more. Underlying reader does not usually read all we ask for.
|
|
}
|
|
Err(err) => {
|
|
error_received = true;
|
|
error = err;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if bytes_read > 0 {
|
|
// store this buffer in the bg data list
|
|
target.truncate(bytes_read);
|
|
let mut data = self.newly_read_data.write().unwrap();
|
|
data.push(dest_data);
|
|
drop(data);
|
|
self.newly_read_data_signal.notify_all();
|
|
}
|
|
|
|
if eof {
|
|
self.bg_eof_reached.store(true, Ordering::Relaxed);
|
|
self.newly_read_data_signal.notify_all(); // anyone waiting for new data needs to know that we reached eof
|
|
break;
|
|
}
|
|
|
|
if error_received {
|
|
// do not ask for more data from 'reader'. We got an error and saved all the data we got before the error.
|
|
// but, wait to set error until we have added our buffer to newly_read_data
|
|
self.set_error(error);
|
|
break;
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"reading entire decompressed file took: {} us, bytes: {}, read_us: {}, waiting_for_buffer_us: {}, largest fetch: {}, error: {:?}",
|
|
now.elapsed().as_micros(),
|
|
total_bytes,
|
|
read_us,
|
|
wait_us,
|
|
max_bytes_read,
|
|
self.error.read().unwrap()
|
|
);
|
|
}
|
|
}
|
|
|
|
impl SharedBufferInternal {
|
|
fn wait_for_newly_read_data(&self) -> bool {
|
|
self.bg_reader_data
|
|
.newly_read_data_signal
|
|
.wait_timeout(SharedBufferBgReader::default_wait_timeout())
|
|
}
|
|
// bg reader uses write lock on 'newly_read_data' each time a buffer is read or recycled
|
|
// client readers read from 'data' using read locks
|
|
// when all of 'data' has been exhausted by clients, 1 client needs to transfer from 'newly_read_data' to 'data' one time.
|
|
// returns true if any data was added to 'data'
|
|
fn transfer_data_from_bg(&self) -> bool {
|
|
let mut from_lock = self.bg_reader_data.newly_read_data.write().unwrap();
|
|
if from_lock.is_empty() {
|
|
// no data available from bg
|
|
return false;
|
|
}
|
|
// grab all data from bg
|
|
let mut newly_read_data: Vec<OneSharedBuffer> = std::mem::take(&mut *from_lock);
|
|
// append all data to fg
|
|
let mut to_lock = self.data.write().unwrap();
|
|
// from_lock has to be held until we have the to_lock lock. Otherwise, we can race with another reader and append to to_lock out of order.
|
|
drop(from_lock);
|
|
to_lock.append(&mut newly_read_data);
|
|
true // data was transferred
|
|
}
|
|
fn has_reached_eof(&self) -> bool {
|
|
self.bg_reader_data.bg_eof_reached.load(Ordering::Relaxed)
|
|
}
|
|
}
|
|
|
|
// only public methods are new and from trait Read
|
|
impl SharedBufferReader {
|
|
pub fn new(original_instance: &SharedBuffer) -> Self {
|
|
let original_instance = &original_instance.instance;
|
|
let current_buffer_index = 0;
|
|
let mut list = original_instance.clients.write().unwrap();
|
|
let my_client_index = list.len();
|
|
if my_client_index > 0 {
|
|
let current_min = list.iter().min().unwrap();
|
|
if current_min > &0 {
|
|
drop(list);
|
|
panic!("SharedBufferReaders must all be created before the first one reads");
|
|
}
|
|
}
|
|
list.push(current_buffer_index);
|
|
drop(list);
|
|
|
|
Self {
|
|
instance: Arc::clone(original_instance),
|
|
my_client_index,
|
|
current_buffer_index,
|
|
index_in_current_data: 0,
|
|
// startup condition for our local reference to the buffer we want to read from.
|
|
// data[0] will always exist. It will be empty, But that is ok. Corresponds to current_buffer_index initial value of 0.
|
|
current_data: original_instance.data.read().unwrap()[0].clone(),
|
|
empty_buffer: original_instance.empty_buffer.clone(),
|
|
}
|
|
}
|
|
fn default_error() -> std::io::Error {
|
|
// AN error
|
|
std::io::Error::from(std::io::ErrorKind::TimedOut)
|
|
}
|
|
fn client_done_reading(&mut self) {
|
|
// has the effect of causing nobody to ever again wait on this reader's progress
|
|
self.update_client_index(usize::MAX);
|
|
}
|
|
|
|
// this client will now be reading from current_buffer_index
|
|
// We may be able to recycle the buffer(s) this client may have been previously potentially using.
|
|
fn update_client_index(&mut self, new_buffer_index: usize) {
|
|
let previous_buffer_index = self.current_buffer_index;
|
|
self.current_buffer_index = new_buffer_index;
|
|
let client_index = self.my_client_index;
|
|
let mut indexes = self.instance.clients.write().unwrap();
|
|
indexes[client_index] = new_buffer_index;
|
|
drop(indexes);
|
|
let mut new_min = *self.instance.clients.read().unwrap().iter().min().unwrap();
|
|
// if new_min == usize::MAX, then every caller is done reading. We could shut down the bg reader and effectively drop everything.
|
|
new_min = std::cmp::min(new_min, self.instance.data.read().unwrap().len());
|
|
|
|
// if any buffer indexes are now no longer used by any readers, then this reader was the last reader holding onto some indexes.
|
|
if new_min > previous_buffer_index {
|
|
// if bg reader reached eof, there is no need to recycle any buffers and they can all be dropped
|
|
let eof = self.instance.has_reached_eof();
|
|
|
|
for recycle in previous_buffer_index..new_min {
|
|
let remove = {
|
|
let mut data = self.instance.data.write().unwrap();
|
|
std::mem::replace(&mut data[recycle], self.empty_buffer.clone())
|
|
};
|
|
if remove.is_empty() {
|
|
continue; // another thread beat us swapping out this buffer, so nothing to recycle here
|
|
}
|
|
|
|
if !eof {
|
|
// if !eof, recycle this buffer and notify waiting reader(s)
|
|
// if eof, just drop buffer this buffer since it isn't needed for reading anymore
|
|
self.instance
|
|
.bg_reader_data
|
|
.buffers
|
|
.write()
|
|
.unwrap()
|
|
.push(remove);
|
|
self.instance.bg_reader_data.new_buffer_signal.notify_all();
|
|
// new buffer available for bg reader
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Read for SharedBufferReader {
|
|
// called many times by client to read small buffer lengths
|
|
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
|
|
let dest_len = buf.len();
|
|
let mut offset_in_dest = 0;
|
|
|
|
let mut eof_seen = false;
|
|
'outer: while offset_in_dest < dest_len {
|
|
// this code is optimized for the common case where we can satisfy this entire read request from current_data without locks
|
|
let source = &*self.current_data;
|
|
|
|
let remaining_source_len = source.len() - self.index_in_current_data;
|
|
let bytes_to_transfer = std::cmp::min(dest_len - offset_in_dest, remaining_source_len);
|
|
// copy what we can
|
|
buf[offset_in_dest..(offset_in_dest + bytes_to_transfer)].copy_from_slice(
|
|
&source
|
|
[self.index_in_current_data..(self.index_in_current_data + bytes_to_transfer)],
|
|
);
|
|
self.index_in_current_data += bytes_to_transfer;
|
|
offset_in_dest += bytes_to_transfer;
|
|
|
|
if offset_in_dest >= dest_len {
|
|
break;
|
|
}
|
|
|
|
// we exhausted the current buffer
|
|
// increment current_buffer_index get the next buffer to continue reading
|
|
self.current_data = self.empty_buffer.clone(); // unref it so it can be recycled without copy
|
|
self.index_in_current_data = 0;
|
|
self.update_client_index(self.current_buffer_index + 1);
|
|
|
|
let instance = &*self.instance;
|
|
let mut lock;
|
|
// hang out in this loop until the buffer we need is available
|
|
loop {
|
|
lock = instance.data.read().unwrap();
|
|
if self.current_buffer_index < lock.len() {
|
|
break;
|
|
}
|
|
drop(lock);
|
|
|
|
if self.instance.transfer_data_from_bg() {
|
|
continue;
|
|
}
|
|
|
|
// another thread may have transferred data, so check again to see if we have data now
|
|
lock = instance.data.read().unwrap();
|
|
if self.current_buffer_index < lock.len() {
|
|
break;
|
|
}
|
|
drop(lock);
|
|
|
|
if eof_seen {
|
|
// eof detected on previous iteration, we have had a chance to read all data that was buffered, and there is not enough for us
|
|
break 'outer;
|
|
}
|
|
|
|
// no data, we could not transfer, and still no data, so check for eof.
|
|
// If we got an eof, then we have to check again for data to make sure there isn't data now that we may be able to transfer or read. Our reading can lag behind the bg read ahead.
|
|
if instance.has_reached_eof() {
|
|
eof_seen = true;
|
|
continue;
|
|
}
|
|
|
|
{
|
|
// Since the bg reader could not satisfy our read, now is a good time to check to see if the bg reader encountered an error.
|
|
// Note this is a write lock because we want to get the actual error detected and return it here and avoid races with other readers if we tried a read and then subsequent write lock.
|
|
// This would be simpler if I could clone an io error.
|
|
let mut error = instance.bg_reader_data.error.write().unwrap();
|
|
if error.is_err() {
|
|
// replace the current error (with AN error instead of ok)
|
|
// return the original error
|
|
return std::mem::replace(&mut *error, Err(Self::default_error()));
|
|
}
|
|
}
|
|
|
|
// no data to transfer, and file not finished, but no error, so wait for bg reader to read some more data
|
|
instance.wait_for_newly_read_data();
|
|
}
|
|
|
|
// refresh current_data inside the lock
|
|
self.current_data = Arc::clone(&lock[self.current_buffer_index]);
|
|
}
|
|
Ok(offset_in_dest)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub mod tests {
|
|
use {
|
|
super::*,
|
|
crossbeam_channel::{unbounded, Receiver},
|
|
rayon::prelude::*,
|
|
};
|
|
|
|
type SimpleReaderReceiverType = Receiver<(Vec<u8>, Option<std::io::Error>)>;
|
|
struct SimpleReader {
|
|
pub receiver: SimpleReaderReceiverType,
|
|
pub data: Vec<u8>,
|
|
pub done: bool,
|
|
pub err: Option<std::io::Error>,
|
|
}
|
|
impl SimpleReader {
|
|
fn new(receiver: SimpleReaderReceiverType) -> Self {
|
|
Self {
|
|
receiver,
|
|
data: Vec::default(),
|
|
done: false,
|
|
err: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Read for SimpleReader {
|
|
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
|
|
if !self.done && self.data.is_empty() {
|
|
let (mut data, err) = self.receiver.recv().unwrap();
|
|
if err.is_some() {
|
|
self.err = err;
|
|
}
|
|
if data.is_empty() {
|
|
self.done = true;
|
|
} else {
|
|
self.data.append(&mut data);
|
|
}
|
|
}
|
|
if self.err.is_some() {
|
|
return Err(self.err.take().unwrap());
|
|
}
|
|
let len_request = buf.len();
|
|
let len_data = self.data.len();
|
|
let to_read = std::cmp::min(len_request, len_data);
|
|
buf[0..to_read].copy_from_slice(&self.data[0..to_read]);
|
|
self.data.drain(0..to_read);
|
|
Ok(to_read)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "total_buffer_budget > 0")]
|
|
fn test_shared_buffer_buffers_invalid() {
|
|
solana_logger::setup();
|
|
let (_sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
SharedBuffer::new_with_sizes(0, 1, file);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "chunk_size > 0")]
|
|
fn test_shared_buffer_buffers_invalid2() {
|
|
solana_logger::setup();
|
|
let (_sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
SharedBuffer::new_with_sizes(1, 0, file);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "SharedBufferReaders must all be created before the first one reads")]
|
|
fn test_shared_buffer_start_too_late() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
let sent = vec![1, 2, 3];
|
|
let _ = sender.send((sent, None));
|
|
let _ = sender.send((done_signal, None));
|
|
assert!(reader.read_to_end(&mut data).is_ok());
|
|
SharedBufferReader::new(&shared_buffer); // created after reader already read
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_simple_read_to_end() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
let sent = vec![1, 2, 3];
|
|
let _ = sender.send((sent.clone(), None));
|
|
let _ = sender.send((done_signal, None));
|
|
assert!(reader.read_to_end(&mut data).is_ok());
|
|
assert_eq!(sent, data);
|
|
}
|
|
|
|
fn get_error() -> std::io::Error {
|
|
std::io::Error::from(std::io::ErrorKind::WriteZero)
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_simple_read() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let done_signal = vec![];
|
|
|
|
let sent = vec![1, 2, 3];
|
|
let mut data = vec![0; sent.len()];
|
|
let _ = sender.send((sent.clone(), None));
|
|
let _ = sender.send((done_signal, None));
|
|
assert_eq!(reader.read(&mut data[..]).unwrap(), sent.len());
|
|
assert_eq!(sent, data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_error() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
let _ = sender.send((done_signal, Some(get_error())));
|
|
assert_eq!(
|
|
reader.read_to_end(&mut data).unwrap_err().kind(),
|
|
get_error().kind()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_2_errors() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut reader2 = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
let _ = sender.send((done_signal, Some(get_error())));
|
|
assert_eq!(
|
|
reader.read_to_end(&mut data).unwrap_err().kind(),
|
|
get_error().kind()
|
|
);
|
|
// #2 will read 2nd, so should get default error, but still an error
|
|
assert_eq!(
|
|
reader2.read_to_end(&mut data).unwrap_err().kind(),
|
|
SharedBufferReader::default_error().kind()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_2_errors_after_read() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut reader2 = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
// send some data
|
|
let sent = vec![1, 2, 3];
|
|
let _ = sender.send((sent.clone(), None));
|
|
// send an error
|
|
let _ = sender.send((done_signal, Some(get_error())));
|
|
assert_eq!(
|
|
reader.read_to_end(&mut data).unwrap_err().kind(),
|
|
get_error().kind()
|
|
);
|
|
// #2 will read valid bytes first and succeed, then get error
|
|
let mut data = vec![0; sent.len()];
|
|
// this read should succeed because it was prior to error being received by bg reader
|
|
assert_eq!(reader2.read(&mut data[..]).unwrap(), sent.len(),);
|
|
assert_eq!(sent, data);
|
|
assert_eq!(
|
|
reader2.read_to_end(&mut data).unwrap_err().kind(),
|
|
SharedBufferReader::default_error().kind()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_2_errors_after_read2() {
|
|
solana_logger::setup();
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer = SharedBuffer::new(file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
let mut reader2 = SharedBufferReader::new(&shared_buffer);
|
|
let mut data = Vec::new();
|
|
let done_signal = vec![];
|
|
|
|
// send some data
|
|
let sent = vec![1, 2, 3];
|
|
let _ = sender.send((sent.clone(), None));
|
|
// send an error
|
|
let _ = sender.send((done_signal, Some(get_error())));
|
|
assert_eq!(
|
|
reader.read_to_end(&mut data).unwrap_err().kind(),
|
|
get_error().kind()
|
|
);
|
|
// #2 will read valid bytes first and succeed, then get error
|
|
let mut data = vec![0; sent.len()];
|
|
// this read should succeed because it is reading data prior to error being received by bg reader
|
|
let expected_len = 1;
|
|
for i in 0..sent.len() {
|
|
let len = reader2.read(&mut data[i..=i]);
|
|
assert!(len.is_ok(), "{len:?}, progress: {i}");
|
|
assert_eq!(len.unwrap(), expected_len, "progress: {i}");
|
|
}
|
|
assert_eq!(sent, data);
|
|
assert_eq!(
|
|
reader2.read(&mut data[0..=0]).unwrap_err().kind(),
|
|
SharedBufferReader::default_error().kind()
|
|
);
|
|
}
|
|
|
|
// read either all or in specified block sizes
|
|
fn test_read_all(
|
|
reader: &mut SharedBufferReader,
|
|
individual_read_size: Option<usize>,
|
|
) -> Vec<u8> {
|
|
let mut data = Vec::new();
|
|
match individual_read_size {
|
|
Some(size) => {
|
|
loop {
|
|
let mut buffer = vec![0; size];
|
|
let result = reader.read(&mut buffer[..]);
|
|
assert!(result.is_ok());
|
|
let len = result.unwrap();
|
|
if len == 0 {
|
|
break; // done reading
|
|
}
|
|
buffer.truncate(len);
|
|
data.append(&mut buffer);
|
|
}
|
|
}
|
|
None => {
|
|
let result = reader.read_to_end(&mut data);
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), data.len());
|
|
}
|
|
}
|
|
data
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_drop_reader2() {
|
|
let done_signal = vec![];
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let budget_sz = 100;
|
|
let chunk_sz = 10;
|
|
let shared_buffer = SharedBuffer::new_with_sizes(budget_sz, chunk_sz, file);
|
|
let size = budget_sz * 2;
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
// with the Read trait, we don't know we are eof until we get Ok(0) from the underlying reader.
|
|
// This can't happen until we have enough space to store another chunk, thus we try to read another chunk and see the Ok(0) returned.
|
|
// Thus, we have to use size < budget_sz here instead of <=
|
|
let reader2 = SharedBufferReader::new(&shared_buffer);
|
|
|
|
let sent = (0..size)
|
|
.map(|i| ((i + size) % 256) as u8)
|
|
.collect::<Vec<_>>();
|
|
|
|
let _ = sender.send((sent.clone(), None));
|
|
let _ = sender.send((done_signal, None));
|
|
|
|
// can't read all data because it is 2x the buffer budget
|
|
let mut data = vec![0; budget_sz];
|
|
assert!(reader.read(&mut data[0..budget_sz]).is_ok());
|
|
drop(reader2);
|
|
let mut rest = test_read_all(&mut reader, None);
|
|
data.append(&mut rest);
|
|
assert_eq!(sent, data);
|
|
}
|
|
|
|
fn adjusted_buffer_size(total_buffer_budget: usize, chunk_size: usize) -> usize {
|
|
let num_buffers = SharedBufferBgReader::num_buffers(total_buffer_budget, chunk_size);
|
|
num_buffers * chunk_size
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_buffer_sweep() {
|
|
solana_logger::setup();
|
|
// try the inflection points with 1 to 3 readers, including a parallel reader
|
|
// a few different chunk sizes
|
|
for chunk_sz in [1, 2, 10] {
|
|
// same # of buffers as default
|
|
let equivalent_buffer_sz =
|
|
chunk_sz * (TOTAL_BUFFER_BUDGET_DEFAULT / CHUNK_SIZE_DEFAULT);
|
|
// 1 buffer, 2 buffers,
|
|
for budget_sz in [
|
|
1,
|
|
chunk_sz,
|
|
chunk_sz * 2,
|
|
equivalent_buffer_sz - 1,
|
|
equivalent_buffer_sz,
|
|
equivalent_buffer_sz * 2,
|
|
] {
|
|
for read_sz in [0, 1, chunk_sz - 1, chunk_sz, chunk_sz + 1] {
|
|
let read_sz = if read_sz > 0 { Some(read_sz) } else { None };
|
|
for reader_ct in 1..=3 {
|
|
for data_size in [
|
|
0,
|
|
1,
|
|
chunk_sz - 1,
|
|
chunk_sz,
|
|
chunk_sz + 1,
|
|
chunk_sz * 2 - 1,
|
|
chunk_sz * 2,
|
|
chunk_sz * 2 + 1,
|
|
budget_sz - 1,
|
|
budget_sz,
|
|
budget_sz + 1,
|
|
budget_sz * 2,
|
|
budget_sz * 2 - 1,
|
|
budget_sz * 2 + 1,
|
|
] {
|
|
let adjusted_budget_sz = adjusted_buffer_size(budget_sz, chunk_sz);
|
|
let done_signal = vec![];
|
|
let (sender, receiver) = unbounded();
|
|
let file = SimpleReader::new(receiver);
|
|
let shared_buffer =
|
|
SharedBuffer::new_with_sizes(budget_sz, chunk_sz, file);
|
|
let mut reader = SharedBufferReader::new(&shared_buffer);
|
|
// with the Read trait, we don't know we are eof until we get Ok(0) from the underlying reader.
|
|
// This can't happen until we have enough space to store another chunk, thus we try to read another chunk and see the Ok(0) returned.
|
|
// Thus, we have to use data_size < adjusted_budget_sz here instead of <=
|
|
let second_reader = reader_ct > 1
|
|
&& data_size < adjusted_budget_sz
|
|
&& read_sz
|
|
.as_ref()
|
|
.map(|sz| sz < &adjusted_budget_sz)
|
|
.unwrap_or(true);
|
|
let reader2 = if second_reader {
|
|
Some(SharedBufferReader::new(&shared_buffer))
|
|
} else {
|
|
None
|
|
};
|
|
let sent = (0..data_size)
|
|
.map(|i| ((i + data_size) % 256) as u8)
|
|
.collect::<Vec<_>>();
|
|
|
|
let parallel_reader = reader_ct > 2;
|
|
let handle = if parallel_reader {
|
|
// Avoid to create more than the number of threads available in the
|
|
// current rayon threadpool. Deadlock could happen otherwise.
|
|
let threads = std::cmp::min(8, rayon::current_num_threads());
|
|
Some({
|
|
let parallel = (0..threads)
|
|
.map(|_| {
|
|
// create before any reading starts
|
|
let reader_ = SharedBufferReader::new(&shared_buffer);
|
|
let sent_ = sent.clone();
|
|
(reader_, sent_)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
Builder::new()
|
|
.spawn(move || {
|
|
parallel.into_par_iter().for_each(
|
|
|(mut reader, sent)| {
|
|
let data = test_read_all(&mut reader, read_sz);
|
|
assert_eq!(
|
|
sent,
|
|
data,
|
|
"{:?}",
|
|
(
|
|
chunk_sz,
|
|
budget_sz,
|
|
read_sz,
|
|
reader_ct,
|
|
data_size,
|
|
adjusted_budget_sz
|
|
)
|
|
);
|
|
},
|
|
)
|
|
})
|
|
.unwrap()
|
|
})
|
|
} else {
|
|
None
|
|
};
|
|
drop(shared_buffer); // readers should work fine even if shared buffer is dropped
|
|
let _ = sender.send((sent.clone(), None));
|
|
let _ = sender.send((done_signal, None));
|
|
let data = test_read_all(&mut reader, read_sz);
|
|
assert_eq!(
|
|
sent,
|
|
data,
|
|
"{:?}",
|
|
(
|
|
chunk_sz,
|
|
budget_sz,
|
|
read_sz,
|
|
reader_ct,
|
|
data_size,
|
|
adjusted_budget_sz
|
|
)
|
|
);
|
|
// a 2nd reader would stall us if we exceed the total buffer size
|
|
if second_reader {
|
|
// #2 will read valid bytes first and succeed, then get error
|
|
let data = test_read_all(&mut reader2.unwrap(), read_sz);
|
|
assert_eq!(sent, data);
|
|
}
|
|
if parallel_reader {
|
|
assert!(handle.unwrap().join().is_ok());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|