use std::future::Future;
use std::panic::{self, AssertUnwindSafe};
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};
use std::time::Duration;
use std::{mem, thread};
use flume::SendTimeoutError;
use futures_core::Stream;
use crate::Position;
use crate::event::Event;
use crate::query::Query;
use crate::read::{DEFAULT_MAX_BATCH_EVENTS, ReadError, ReadHandle};
type ReadItem = Result<Vec<(Position, Event)>, ReadError>;
const SHUTDOWN_POLL: Duration = Duration::from_millis(50);
#[derive(Debug, Clone, Copy)]
pub struct ReadPoolConfig {
pub workers: usize,
pub batch_events: usize,
pub batch_bytes: usize,
pub channel_depth: usize,
pub queue_capacity: usize,
}
impl Default for ReadPoolConfig {
fn default() -> ReadPoolConfig {
ReadPoolConfig {
workers: 1,
batch_events: DEFAULT_MAX_BATCH_EVENTS,
batch_bytes: 512 * 1024,
channel_depth: 8,
queue_capacity: 1024,
}
}
}
struct ReadJob {
query: Query,
cursor: Position,
limit: Option<u64>,
reverse: bool,
out: flume::Sender<ReadItem>,
batch_events: usize,
batch_bytes: usize,
#[cfg(test)]
inject_panic: bool,
}
pub struct ReadPool {
tx: Option<flume::Sender<ReadJob>>,
workers: Vec<thread::JoinHandle<()>>,
shutdown: Arc<AtomicBool>,
batch_events: usize,
batch_bytes: usize,
channel_depth: usize,
}
impl ReadPool {
pub fn new(handle: ReadHandle, workers: usize) -> ReadPool {
ReadPool::with_config(
handle,
ReadPoolConfig {
workers,
..ReadPoolConfig::default()
},
)
}
pub fn with_config(handle: ReadHandle, config: ReadPoolConfig) -> ReadPool {
let worker_count = config.workers.max(1);
let batch_events = config.batch_events.max(1);
let batch_bytes = config.batch_bytes.max(1);
let channel_depth = config.channel_depth.max(1);
let queue_capacity = config.queue_capacity.max(1);
let (tx, rx) = flume::bounded::<ReadJob>(queue_capacity);
let shutdown = Arc::new(AtomicBool::new(false));
let mut workers = Vec::with_capacity(worker_count);
for _ in 0..worker_count {
let rx = rx.clone();
let handle = handle.clone();
let shutdown = Arc::clone(&shutdown);
let worker = thread::Builder::new()
.name("tephra-read-pool".to_string())
.spawn(move || worker_loop(handle, rx, shutdown))
.expect("spawn tephra read pool worker");
workers.push(worker);
}
ReadPool {
tx: Some(tx),
workers,
shutdown,
batch_events,
batch_bytes,
channel_depth,
}
}
pub fn read(&self, query: Query, after: Position, limit: Option<u64>) -> ReadStream {
self.stream(query, after, false, limit)
}
pub fn read_back(&self, query: Query, before: Position, limit: Option<u64>) -> ReadStream {
self.stream(query, before, true, limit)
}
fn stream(
&self,
query: Query,
cursor: Position,
reverse: bool,
limit: Option<u64>,
) -> ReadStream {
let (out_tx, out_rx) = flume::bounded::<ReadItem>(self.channel_depth);
let job = ReadJob {
query,
cursor,
limit,
reverse,
out: out_tx,
batch_events: self.batch_events,
batch_bytes: self.batch_bytes,
#[cfg(test)]
inject_panic: false,
};
match &self.tx {
Some(tx) => {
let tx = tx.clone();
let submit: Pin<Box<dyn Future<Output = ()> + Send>> = Box::pin(async move {
let _ = tx.send_async(job).await;
});
ReadStream(StreamState::Submitting {
submit,
out: out_rx,
})
}
None => ReadStream(StreamState::Done),
}
}
pub async fn read_all(
&self,
query: Query,
after: Position,
limit: Option<u64>,
) -> Result<Vec<(Position, Event)>, ReadError> {
self.collect(query, after, false, limit).await
}
pub async fn read_all_back(
&self,
query: Query,
before: Position,
limit: Option<u64>,
) -> Result<Vec<(Position, Event)>, ReadError> {
self.collect(query, before, true, limit).await
}
async fn collect(
&self,
query: Query,
cursor: Position,
reverse: bool,
limit: Option<u64>,
) -> Result<Vec<(Position, Event)>, ReadError> {
let (out_tx, out_rx) = flume::bounded::<ReadItem>(self.channel_depth);
let job = ReadJob {
query,
cursor,
limit,
reverse,
out: out_tx,
batch_events: self.batch_events,
batch_bytes: self.batch_bytes,
#[cfg(test)]
inject_panic: false,
};
if let Some(tx) = &self.tx {
let _ = tx.send_async(job).await;
}
let mut out = Vec::new();
loop {
match out_rx.recv_async().await {
Ok(Ok(mut batch)) => out.append(&mut batch),
Ok(Err(err)) => return Err(err),
Err(_) => return Ok(out),
}
}
}
pub fn shutdown(self) {
}
#[cfg(test)]
fn submit_panicking(&self, query: Query) -> flume::Receiver<ReadItem> {
let (out_tx, out_rx) = flume::bounded::<ReadItem>(self.channel_depth);
let job = ReadJob {
query,
cursor: Position::ZERO,
limit: None,
reverse: false,
out: out_tx,
batch_events: self.batch_events,
batch_bytes: self.batch_bytes,
inject_panic: true,
};
if let Some(tx) = &self.tx {
tx.send(job).expect("submit panicking read");
}
out_rx
}
}
impl Drop for ReadPool {
fn drop(&mut self) {
self.shutdown.store(true, Ordering::Release);
drop(self.tx.take());
for worker in self.workers.drain(..) {
let _ = worker.join();
}
}
}
fn worker_loop(handle: ReadHandle, rx: flume::Receiver<ReadJob>, shutdown: Arc<AtomicBool>) {
while let Ok(job) = rx.recv() {
let out = job.out.clone();
let outcome = panic::catch_unwind(AssertUnwindSafe(|| run_job(&handle, job, &shutdown)));
if outcome.is_err() {
let _ = out.send(ReadItem::Err(ReadError::Aborted));
}
}
}
fn run_job(handle: &ReadHandle, job: ReadJob, shutdown: &AtomicBool) {
let ReadJob {
query,
cursor,
limit,
reverse,
out,
batch_events,
batch_bytes,
#[cfg(test)]
inject_panic,
} = job;
#[cfg(test)]
if inject_panic {
panic!("injected read panic (test hook)");
}
let mut reads = if reverse {
handle.read_back(&query, cursor, limit)
} else {
handle.read(&query, cursor, limit)
};
let mut batch: Vec<(Position, Event)> = Vec::new();
let mut batch_bytes_seen = 0usize;
while let Some(item) = reads.next() {
match item {
Ok(seq) => {
batch_bytes_seen += seq.event.as_bytes().len();
batch.push((seq.position, seq.event.to_owned()));
if batch.len() >= batch_events || batch_bytes_seen >= batch_bytes {
if !send_item(&out, Ok(mem::take(&mut batch)), shutdown) {
return;
}
batch_bytes_seen = 0;
}
}
Err(err) => {
let _ = send_item(&out, Err(err), shutdown);
return;
}
}
}
if !batch.is_empty() {
let _ = send_item(&out, Ok(batch), shutdown);
}
}
fn send_item(out: &flume::Sender<ReadItem>, mut item: ReadItem, shutdown: &AtomicBool) -> bool {
loop {
if shutdown.load(Ordering::Acquire) {
return false;
}
match out.send_timeout(item, SHUTDOWN_POLL) {
Ok(()) => return true,
Err(SendTimeoutError::Timeout(returned)) => item = returned,
Err(SendTimeoutError::Disconnected(_)) => return false,
}
}
}
enum StreamState {
Submitting {
submit: Pin<Box<dyn Future<Output = ()> + Send>>,
out: flume::Receiver<ReadItem>,
},
Streaming(flume::r#async::RecvStream<'static, ReadItem>),
Done,
}
pub struct ReadStream(StreamState);
impl Stream for ReadStream {
type Item = ReadItem;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
match mem::replace(&mut this.0, StreamState::Done) {
StreamState::Submitting { mut submit, out } => match submit.as_mut().poll(cx) {
Poll::Ready(()) => {
this.0 = StreamState::Streaming(out.into_stream());
}
Poll::Pending => {
this.0 = StreamState::Submitting { submit, out };
return Poll::Pending;
}
},
StreamState::Streaming(mut stream) => {
let polled = Pin::new(&mut stream).poll_next(cx);
this.0 = StreamState::Streaming(stream);
return polled;
}
StreamState::Done => return Poll::Ready(None),
}
}
}
}
#[cfg(all(test, feature = "async"))]
mod tests {
use std::future;
use std::time::Instant;
use super::*;
use tempfile::TempDir;
use crate::event::{Event, EventType, Tag, Tags};
use crate::log::set::{SegmentConfig, SegmentSet};
use crate::query::{Query, QueryItem};
use crate::writer::{WriteCoordinator, WriteHandle, WriterConfig};
fn block_on<F: future::Future>(fut: F) -> F::Output {
use std::pin::pin;
use std::sync::Arc;
use std::task::{Context, Poll, Wake, Waker};
use std::thread::{self, Thread};
struct ThreadWaker(Thread);
impl Wake for ThreadWaker {
fn wake(self: Arc<Self>) {
self.0.unpark();
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.unpark();
}
}
let waker = Waker::from(Arc::new(ThreadWaker(thread::current())));
let mut cx = Context::from_waker(&waker);
let mut fut = pin!(fut);
loop {
match fut.as_mut().poll(&mut cx) {
Poll::Ready(out) => return out,
Poll::Pending => thread::park(),
}
}
}
fn tags(items: &[&str]) -> Tags {
Tags::new(items.iter().map(|s| Tag::new(*s).unwrap())).unwrap()
}
fn event(ty: &str, tag_strs: &[&str]) -> Event {
Event::new(&EventType::new(ty).unwrap(), &tags(tag_strs), b"data").unwrap()
}
fn store() -> (TempDir, WriteCoordinator, WriteHandle) {
let dir = TempDir::new().unwrap();
let set = SegmentSet::open(dir.path(), SegmentConfig::new(64 << 20)).unwrap();
let (coord, handle) = WriteCoordinator::start(set, WriterConfig::default()).unwrap();
(dir, coord, handle)
}
fn direct(handle: &WriteHandle, query: &Query) -> Vec<Position> {
let mut reads = handle.read(query, Position::ZERO, None);
let mut out = Vec::new();
while let Some(item) = reads.next() {
out.push(item.expect("read failed").position);
}
out
}
fn drain(mut stream: ReadStream) -> (Vec<Position>, usize) {
block_on(async {
let mut positions = Vec::new();
let mut batches = 0;
future::poll_fn(|cx| {
loop {
match Pin::new(&mut stream).poll_next(cx) {
Poll::Ready(Some(item)) => {
let batch = item.expect("read failed");
batches += 1;
positions.extend(batch.into_iter().map(|(p, _)| p));
}
Poll::Ready(None) => return Poll::Ready((positions.clone(), batches)),
Poll::Pending => return Poll::Pending,
}
}
})
.await
})
}
#[test]
fn read_all_returns_every_event_in_position_order() {
let (_dir, _coord, handle) = store();
for i in 0..50 {
handle
.append(
vec![event("Enrolled", &[&format!("course:c{}", i % 5)])],
None,
)
.unwrap();
}
let pool = ReadPool::new(handle.reader(), 4);
let got = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
let positions: Vec<Position> = got.iter().map(|(p, _)| *p).collect();
assert_eq!(positions, direct(&handle, &Query::all()));
assert!(positions.windows(2).all(|w| w[0] < w[1]));
}
#[test]
fn read_back_streams_newest_first_and_matches_read_all_reversed() {
let (_dir, _coord, handle) = store();
for i in 0..50 {
handle
.append(
vec![event("Enrolled", &[&format!("course:c{}", i % 5)])],
None,
)
.unwrap();
}
let pool = ReadPool::new(handle.reader(), 4);
let forward = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
let back = block_on(pool.read_all_back(Query::all(), Position::MAX, None)).unwrap();
let want: Vec<Position> = forward.iter().rev().map(|(p, _)| *p).collect();
let got: Vec<Position> = back.iter().map(|(p, _)| *p).collect();
assert_eq!(got, want);
assert!(got.windows(2).all(|w| w[0] > w[1]), "descending");
let (streamed, _) = drain(pool.read_back(Query::all(), Position::MAX, None));
assert_eq!(streamed, got);
let capped = block_on(pool.read_all_back(Query::all(), Position::MAX, Some(5))).unwrap();
let capped: Vec<Position> = capped.iter().map(|(p, _)| *p).collect();
assert_eq!(capped, want[..5]);
}
#[test]
fn streamed_batches_concatenate_to_read_all() {
let (_dir, _coord, handle) = store();
for _ in 0..30 {
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
}
let pool = ReadPool::with_config(
handle.reader(),
ReadPoolConfig {
workers: 2,
batch_events: 4,
channel_depth: 2,
..ReadPoolConfig::default()
},
);
let (positions, batches) = drain(pool.read(Query::all(), Position::ZERO, None));
assert_eq!(positions, direct(&handle, &Query::all()));
assert!(batches >= 2, "expected several batches, got {batches}");
}
#[test]
fn query_and_limit_are_honored() {
let (_dir, _coord, handle) = store();
for i in 0..20 {
let course = if i % 2 == 0 {
"course:even"
} else {
"course:odd"
};
handle
.append(vec![event("Enrolled", &[course])], None)
.unwrap();
}
let pool = ReadPool::new(handle.reader(), 2);
let query = Query::item(QueryItem::with_tags(tags(&["course:even"])));
let limited = block_on(pool.read_all(query.clone(), Position::ZERO, Some(3))).unwrap();
assert_eq!(limited.len(), 3);
let all_even = block_on(pool.read_all(query.clone(), Position::ZERO, None)).unwrap();
assert_eq!(all_even.len(), 10);
assert_eq!(
all_even.iter().map(|(p, _)| *p).collect::<Vec<_>>(),
direct(&handle, &query),
);
}
#[test]
fn concurrent_reads_all_complete_while_a_writer_appends() {
use std::sync::Arc;
let (_dir, _coord, handle) = store();
for _ in 0..40 {
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
}
let pool = Arc::new(ReadPool::new(handle.reader(), 4));
let writer = {
let handle = handle.clone();
thread::spawn(move || {
for _ in 0..40 {
handle
.append(vec![event("Enrolled", &["course:c2"])], None)
.unwrap();
}
})
};
let readers: Vec<_> = (0..8)
.map(|_| {
let pool = Arc::clone(&pool);
thread::spawn(move || {
let got = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
assert!(got.len() >= 40);
let positions: Vec<Position> = got.iter().map(|(p, _)| *p).collect();
assert!(positions.windows(2).all(|w| w[0] < w[1]));
})
})
.collect();
writer.join().unwrap();
for reader in readers {
reader.join().unwrap();
}
}
#[test]
fn shutdown_joins_workers_and_a_later_read_is_empty() {
let (_dir, _coord, handle) = store();
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
let pool = ReadPool::new(handle.reader(), 2);
let before = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
assert_eq!(before.len(), 1);
pool.shutdown();
}
#[test]
fn a_tiny_bounded_queue_still_reads_correctly() {
let (_dir, _coord, handle) = store();
for _ in 0..25 {
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
}
let pool = ReadPool::with_config(
handle.reader(),
ReadPoolConfig {
workers: 1,
queue_capacity: 1,
channel_depth: 1,
batch_events: 3,
..ReadPoolConfig::default()
},
);
let a = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
let b = block_on(pool.read_all(Query::all(), Position::ZERO, None)).unwrap();
assert_eq!(a.len(), 25);
assert_eq!(
a.iter().map(|(p, _)| *p).collect::<Vec<_>>(),
direct(&handle, &Query::all())
);
assert_eq!(a.len(), b.len());
}
#[test]
fn a_panicking_read_surfaces_as_aborted_not_a_clean_end() {
let (_dir, _coord, handle) = store();
for _ in 0..5 {
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
}
let pool = ReadPool::new(handle.reader(), 1);
let rx = pool.submit_panicking(Query::all());
let first = block_on(rx.recv_async()).expect("channel closed without reporting an item");
match first {
Err(ReadError::Aborted) => {}
Err(other) => panic!("expected ReadError::Aborted, got a different error: {other}"),
Ok(_) => panic!("expected ReadError::Aborted, got a successful batch"),
}
assert!(
block_on(rx.recv_async()).is_err(),
"expected the channel to close after the abort",
);
}
#[test]
fn dropping_the_pool_with_an_undrained_stream_does_not_hang() {
let (_dir, _coord, handle) = store();
for _ in 0..50 {
handle
.append(vec![event("Enrolled", &["course:c1"])], None)
.unwrap();
}
let pool = ReadPool::with_config(
handle.reader(),
ReadPoolConfig {
workers: 1,
channel_depth: 1,
batch_events: 1,
..ReadPoolConfig::default()
},
);
let mut stream = pool.read(Query::all(), Position::ZERO, None);
let first = block_on(async {
future::poll_fn(|cx| match Pin::new(&mut stream).poll_next(cx) {
Poll::Ready(item) => Poll::Ready(item),
Poll::Pending => Poll::Pending,
})
.await
});
assert!(first.is_some());
thread::sleep(Duration::from_millis(50));
let start = Instant::now();
drop(pool);
assert!(
start.elapsed() < Duration::from_secs(5),
"drop hung on a parked worker",
);
drop(stream);
}
}