use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::JoinHandle;
use rudb_common::Result;
use crate::File;
use crate::submit::{Completion, Filler, Request, Response};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Config {
pub threads: usize,
pub coalesce_gap: u64,
pub coalesce_span: u64,
pub coalesce: bool,
}
impl Config {
#[must_use]
pub fn local_disk() -> Self {
Self {
threads: cores().clamp(2, 8),
coalesce_gap: 16 << 10,
coalesce_span: 1 << 20,
coalesce: true,
}
}
#[must_use]
pub fn object_store() -> Self {
Self {
threads: (cores() * 8).clamp(32, 128),
coalesce_gap: 512 << 10,
coalesce_span: 8 << 20,
coalesce: true,
}
}
#[must_use]
pub fn with_threads(mut self, threads: usize) -> Self {
self.threads = threads.max(1);
self
}
#[must_use]
pub fn coalescing(mut self, gap: u64) -> Self {
self.coalesce = true;
self.coalesce_gap = gap;
self
}
#[must_use]
pub fn not_coalescing(mut self) -> Self {
self.coalesce = false;
self
}
}
impl Default for Config {
fn default() -> Self {
Self::local_disk()
}
}
fn cores() -> usize {
std::thread::available_parallelism().map_or(4, std::num::NonZero::get)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Stats {
pub requests: u64,
pub reads: u64,
pub wanted: u64,
pub read: u64,
}
#[derive(Debug, Default)]
struct Counters {
requests: AtomicU64,
reads: AtomicU64,
wanted: AtomicU64,
read: AtomicU64,
}
impl Counters {
fn snapshot(&self) -> Stats {
Stats {
requests: self.requests.load(Ordering::Relaxed),
reads: self.reads.load(Ordering::Relaxed),
wanted: self.wanted.load(Ordering::Relaxed),
read: self.read.load(Ordering::Relaxed),
}
}
}
#[derive(Debug)]
struct Part {
index: usize,
offset: u64,
buf: Vec<u8>,
}
#[derive(Debug)]
struct Job {
file: Arc<dyn File>,
filler: Filler,
offset: u64,
span: usize,
parts: Vec<Part>,
}
impl Job {
fn run(self, counters: &Counters) {
let Self { file, filler, offset, span, mut parts } = self;
counters.reads.fetch_add(1, Ordering::Relaxed);
if parts.len() == 1 {
let part = parts.pop().unwrap_or_else(|| unreachable!("checked one part"));
let Part { index, offset, mut buf } = part;
let outcome = file.read_at(offset, &mut buf).map(|read| {
counters.read.fetch_add(read as u64, Ordering::Relaxed);
Response::new(index, offset, read, buf)
});
filler.finish(index, outcome);
return;
}
let mut scratch = vec![0u8; span];
match file.read_at(offset, &mut scratch) {
Ok(read) => {
counters.read.fetch_add(read as u64, Ordering::Relaxed);
for part in parts {
let Part { index, offset: at, mut buf } = part;
let start = (at - offset) as usize;
let got = read.saturating_sub(start).min(buf.len());
buf[..got].copy_from_slice(&scratch[start..start + got]);
filler.finish(index, Ok(Response::new(index, at, got, buf)));
}
}
Err(error) => {
for part in parts {
filler.finish(part.index, Err(error.clone()));
}
}
}
}
}
#[derive(Debug)]
struct Queue {
jobs: VecDeque<Job>,
closed: bool,
}
#[derive(Debug)]
struct Shared {
queue: Mutex<Queue>,
work: Condvar,
counters: Counters,
}
impl Shared {
fn lock(&self) -> std::sync::MutexGuard<'_, Queue> {
self.queue.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
#[derive(Debug, Clone)]
pub struct Pool {
shared: Arc<Shared>,
#[allow(dead_code, reason = "holding this is the point of it, reading it is not")]
threads: Arc<Threads>,
config: Config,
}
#[derive(Debug)]
struct Threads {
shared: Arc<Shared>,
handles: Mutex<Vec<JoinHandle<()>>>,
}
impl Drop for Threads {
fn drop(&mut self) {
self.shared.lock().closed = true;
self.shared.work.notify_all();
let handles = std::mem::take(
&mut *self.handles.lock().unwrap_or_else(std::sync::PoisonError::into_inner),
);
for handle in handles {
let _ = handle.join();
}
}
}
impl Pool {
#[must_use]
pub fn new(config: Config) -> Self {
let shared = Arc::new(Shared {
queue: Mutex::new(Queue { jobs: VecDeque::new(), closed: false }),
work: Condvar::new(),
counters: Counters::default(),
});
let mut handles = Vec::with_capacity(config.threads);
for n in 0..config.threads.max(1) {
let shared = Arc::clone(&shared);
handles.push(
std::thread::Builder::new()
.name(format!("rudb-io-{n}"))
.spawn(move || worker(&shared))
.expect("could not spawn an I/O thread"),
);
}
let threads =
Arc::new(Threads { shared: Arc::clone(&shared), handles: Mutex::new(handles) });
Self { shared, threads, config }
}
#[must_use]
pub fn config(&self) -> Config {
self.config
}
#[must_use]
pub fn stats(&self) -> Stats {
self.shared.counters.snapshot()
}
#[must_use]
pub fn submit(&self, file: &Arc<dyn File>, requests: Vec<Request>) -> Completion {
let count = requests.len();
self.shared.counters.requests.fetch_add(count as u64, Ordering::Relaxed);
let wanted: u64 = requests.iter().map(|r| r.len() as u64).sum();
self.shared.counters.wanted.fetch_add(wanted, Ordering::Relaxed);
let (completion, filler) = Completion::pending(count);
let mut wanted_nothing = Vec::new();
let mut real = Vec::with_capacity(count);
for (index, request) in requests.into_iter().enumerate() {
if request.is_empty() {
wanted_nothing.push((index, request));
} else {
real.push((index, request));
}
}
for (index, request) in wanted_nothing {
let offset = request.offset();
filler.finish(index, Ok(Response::new(index, offset, 0, request.into_buffer())));
}
let jobs = plan(real, self.config, file, &filler);
let mut queue = self.shared.lock();
if queue.closed {
drop(queue);
drop(filler);
return completion;
}
queue.jobs.extend(jobs);
drop(queue);
self.shared.work.notify_all();
completion
}
#[must_use]
pub fn queued(&self) -> usize {
self.shared.lock().jobs.len()
}
}
fn plan(
requests: Vec<(usize, Request)>,
config: Config,
file: &Arc<dyn File>,
filler: &Filler,
) -> Vec<Job> {
let mut parts: Vec<Part> = requests
.into_iter()
.map(|(index, request)| {
let offset = request.offset();
Part { index, offset, buf: request.into_buffer() }
})
.collect();
parts.sort_by_key(|part| part.offset);
let mut jobs: Vec<Job> = Vec::with_capacity(parts.len());
for part in parts {
let end = part.offset + part.buf.len() as u64;
if config.coalesce {
if let Some(last) = jobs.last_mut() {
let last_end = last.offset + last.span as u64;
let gap = part.offset.saturating_sub(last_end);
let span = end.saturating_sub(last.offset);
if gap <= config.coalesce_gap && span <= config.coalesce_span {
last.span = span as usize;
last.parts.push(part);
continue;
}
}
}
jobs.push(Job {
file: Arc::clone(file),
filler: filler.clone(),
offset: part.offset,
span: part.buf.len(),
parts: vec![part],
});
}
jobs
}
fn worker(shared: &Arc<Shared>) {
loop {
let mut queue = shared.lock();
let job = loop {
if let Some(job) = queue.jobs.pop_front() {
break job;
}
if queue.closed {
return;
}
queue = shared.work.wait(queue).unwrap_or_else(std::sync::PoisonError::into_inner);
};
drop(queue);
job.run(&shared.counters);
}
}
#[derive(Debug)]
pub struct Pooled {
file: Arc<dyn File>,
pool: Pool,
}
impl Pooled {
#[must_use]
pub fn new(file: Box<dyn File>, pool: Pool) -> Self {
Self { file: Arc::from(file), pool }
}
#[must_use]
pub fn pool(&self) -> &Pool {
&self.pool
}
}
impl File for Pooled {
fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
self.file.read_at(offset, buf)
}
fn submit(&self, requests: Vec<Request>) -> Completion {
self.pool.submit(&self.file, requests)
}
fn write_at(&self, offset: u64, data: &[u8]) -> Result<()> {
self.file.write_at(offset, data)
}
fn sync(&self) -> Result<()> {
self.file.sync()
}
fn truncate(&self, len: u64) -> Result<()> {
self.file.truncate(len)
}
fn len(&self) -> Result<u64> {
self.file.len()
}
}
#[cfg(test)]
mod tests {
use std::path::Path;
use super::{Config, Pool, Pooled};
use crate::submit::{Request, Response};
use crate::{File, Filesystem, OpenMode, SimFilesystem};
fn ramp(pool: Pool) -> Pooled {
let fs = SimFilesystem::new();
let file = fs.open(Path::new("/data"), OpenMode::Create).unwrap();
let bytes: Vec<u8> = (0..=255u8).collect();
file.write_at(0, &bytes).unwrap();
file.sync().unwrap();
Pooled::new(file, pool)
}
#[test]
fn a_batch_comes_back_answering_the_requests_it_was_given() {
let file = ramp(Pool::new(Config::local_disk()));
let responses = file
.submit(vec![Request::new(100, 4), Request::new(0, 4), Request::new(200, 4)])
.wait()
.unwrap();
assert_eq!(responses[0].bytes(), &[100, 101, 102, 103]);
assert_eq!(responses[1].bytes(), &[0, 1, 2, 3]);
assert_eq!(responses[2].bytes(), &[200, 201, 202, 203]);
}
#[test]
fn one_thread_answers_everything_just_as_well_as_eight() {
for threads in [1, 2, 8] {
let file = ramp(Pool::new(Config::local_disk().with_threads(threads)));
let requests = (0..32).map(|i| Request::new(i * 8, 8)).collect::<Vec<_>>();
let responses = file.submit(requests).wait().unwrap();
assert_eq!(responses.len(), 32, "at {threads} threads");
for (i, response) in responses.iter().enumerate() {
assert_eq!(response.bytes()[0], (i * 8) as u8, "at {threads} threads");
}
}
}
#[test]
fn adjacent_ranges_become_one_read_and_the_bytes_do_not_change() {
let pool = Pool::new(Config::local_disk().with_threads(1).coalescing(0));
let file = ramp(pool.clone());
let responses = file
.submit(vec![Request::new(0, 8), Request::new(8, 8), Request::new(16, 8)])
.wait()
.unwrap();
assert_eq!(pool.stats().reads, 1, "three adjacent ranges are one read");
assert_eq!(pool.stats().requests, 3);
assert_eq!(pool.stats().wanted, 24);
assert_eq!(pool.stats().read, 24, "and no byte was read that nobody asked for");
for (i, response) in responses.iter().enumerate() {
assert_eq!(response.bytes()[0], (i * 8) as u8);
}
}
#[test]
fn a_gap_is_read_and_thrown_away_and_the_byte_count_says_so() {
let pool = Pool::new(Config::local_disk().with_threads(1).coalescing(16));
let file = ramp(pool.clone());
let responses = file.submit(vec![Request::new(0, 4), Request::new(20, 4)]).wait().unwrap();
assert_eq!(pool.stats().reads, 1);
assert_eq!(pool.stats().wanted, 8);
assert_eq!(pool.stats().read, 24, "the sixteen byte gap was read too");
assert_eq!(responses[0].bytes(), &[0, 1, 2, 3]);
assert_eq!(responses[1].bytes(), &[20, 21, 22, 23]);
}
#[test]
fn a_gap_wider_than_the_policy_stays_two_reads() {
let pool = Pool::new(Config::local_disk().with_threads(1).coalescing(4));
let file = ramp(pool.clone());
file.submit(vec![Request::new(0, 4), Request::new(20, 4)]).wait().unwrap();
assert_eq!(pool.stats().reads, 2);
assert_eq!(pool.stats().read, 8);
}
#[test]
fn the_span_limit_stops_a_chain_of_small_gaps_becoming_one_huge_read() {
let mut config = Config::local_disk().with_threads(1).coalescing(12);
config.coalesce_span = 32;
let pool = Pool::new(config);
let file = ramp(pool.clone());
let requests = (0..8).map(|i| Request::new(i * 16, 4)).collect::<Vec<_>>();
file.submit(requests).wait().unwrap();
assert_eq!(pool.stats().reads, 4, "eight ranges over 128 bytes, capped at 32 a read");
}
#[test]
fn both_defaults_coalesce_and_the_local_one_does_it_far_more_narrowly() {
assert!(Config::local_disk().coalesce);
assert!(Config::object_store().coalesce);
assert!(Config::object_store().coalesce_gap >= Config::local_disk().coalesce_gap * 8);
assert!(Config::object_store().threads >= Config::local_disk().threads * 4);
}
#[test]
fn the_local_gap_is_too_small_to_swallow_the_space_between_two_columns() {
let pool = Pool::new(Config::local_disk().with_threads(1));
let fs = SimFilesystem::new();
let handle = fs.open(Path::new("/columns"), OpenMode::Create).unwrap();
handle.write_at(0, &vec![7u8; 2 << 20]).unwrap();
handle.sync().unwrap();
let file = Pooled::new(handle, pool.clone());
let requests = (0..4).map(|i| Request::new(i * (512 << 10), 64 << 10)).collect::<Vec<_>>();
file.submit(requests).wait().unwrap();
assert_eq!(pool.stats().reads, 4, "four columns 448KiB apart are four reads and not one");
assert_eq!(pool.stats().read, pool.stats().wanted, "and nothing else was read");
}
#[test]
fn a_short_read_stays_short_through_a_merge() {
let pool = Pool::new(Config::local_disk().with_threads(1).coalescing(0));
let file = ramp(pool.clone());
let responses =
file.submit(vec![Request::new(248, 8), Request::new(256, 8)]).wait().unwrap();
assert_eq!(pool.stats().reads, 1);
assert!(!responses[0].is_short());
assert!(responses[1].is_short());
assert_eq!(responses[1].read(), 0);
}
#[test]
fn a_failed_read_fails_every_request_it_was_merged_with_and_no_others() {
let fs = SimFilesystem::new();
let file = fs.open(Path::new("/data"), OpenMode::Create).unwrap();
file.write_at(0, &(0..=255u8).collect::<Vec<u8>>()).unwrap();
file.sync().unwrap();
let pool = Pool::new(Config::local_disk().with_threads(1).coalescing(0));
let pooled = Pooled::new(file, pool);
fs.fail_read_at(0);
let mut completion =
pooled.submit(vec![Request::new(0, 8), Request::new(8, 8), Request::new(128, 8)]);
let mut failed = 0;
let mut answered = 0;
while let Some(outcome) = completion.take() {
match outcome {
Ok(_) => answered += 1,
Err(_) => failed += 1,
}
}
assert_eq!((answered, failed), (1, 2));
}
#[test]
fn an_empty_request_is_answered_rather_than_queued() {
let pool = Pool::new(Config::local_disk().with_threads(1));
let file = ramp(pool.clone());
let responses = file.submit(vec![Request::new(0, 0), Request::new(4, 4)]).wait().unwrap();
assert_eq!(pool.stats().reads, 1, "nobody reads nothing");
assert_eq!(responses.len(), 2);
assert_eq!(responses[0].read(), 0);
assert_eq!(responses[1].bytes(), &[4, 5, 6, 7]);
}
#[test]
fn an_empty_batch_is_done_before_it_is_submitted() {
let pool = Pool::new(Config::local_disk());
let file = ramp(pool.clone());
let completion = file.submit(Vec::new());
assert!(completion.is_done());
assert!(completion.wait().unwrap().is_empty());
}
#[test]
fn a_submission_returns_before_the_reads_do() {
let pool = Pool::new(Config::local_disk().with_threads(1));
let file = ramp(pool.clone());
let requests = (0..64).map(|i| Request::new(i * 4, 4)).collect::<Vec<_>>();
let completion = file.submit(requests);
let responses = completion.wait().unwrap();
assert_eq!(responses.len(), 64);
assert_eq!(pool.stats().requests, 64);
}
#[test]
fn everything_submitted_is_answered_once_the_pool_is_stopping() {
let pool = Pool::new(Config::local_disk().with_threads(2));
let file = ramp(pool.clone());
let completion = file.submit(vec![Request::new(0, 4)]);
drop(pool);
let responses = completion.wait();
assert!(responses.is_ok() || responses.is_err());
}
#[test]
fn read_at_on_a_pooled_file_is_the_read_underneath_it() {
let file = ramp(Pool::new(Config::local_disk()));
let mut buf = [0u8; 4];
file.read_exact_at(64, &mut buf).unwrap();
assert_eq!(buf, [64, 65, 66, 67]);
assert_eq!(file.len().unwrap(), 256);
}
#[test]
fn responses_are_in_submission_order_whatever_order_the_threads_finished_in() {
let file = ramp(Pool::new(Config::local_disk().with_threads(8)));
let requests = (0..32).rev().map(|i| Request::new(i * 8, 8)).collect::<Vec<_>>();
let responses = file.submit(requests).wait().unwrap();
let indices: Vec<usize> = responses.iter().map(Response::index).collect();
assert_eq!(indices, (0..32).collect::<Vec<_>>());
for (i, response) in responses.iter().enumerate() {
assert_eq!(response.bytes()[0], ((31 - i) * 8) as u8);
}
}
}