use std::collections::VecDeque;
use std::fs::File;
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
#[cfg(feature = "script")]
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::task::Waker;
use std::time::Duration;
pub(crate) const DEFAULT_QUEUE_DEPTH: usize = 64;
pub(crate) const MAX_WAITING_SUBMITTERS: usize = 64;
pub(crate) const MAX_BATCH_RECORDS: usize = 64;
pub(crate) const MAX_BATCH_BYTES: usize = 256 * 1024;
pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
match mutex.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
#[cfg(feature = "script")]
pub(crate) fn read<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
match lock.read() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
#[cfg(feature = "script")]
pub(crate) fn write<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
match lock.write() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
pub(crate) fn wait<'a, T>(condvar: &Condvar, held: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
match condvar.wait(held) {
Ok(woke) => woke,
Err(poisoned) => poisoned.into_inner(),
}
}
pub(crate) fn wait_timeout<'a, T>(
condvar: &Condvar,
held: MutexGuard<'a, T>,
bound: Duration,
) -> (MutexGuard<'a, T>, std::sync::WaitTimeoutResult) {
match condvar.wait_timeout(held, bound) {
Ok((woke, verdict)) => (woke, verdict),
Err(poisoned) => poisoned.into_inner(),
}
}
#[cfg(test)]
pub(crate) fn wait_timeout_while<'a, T, F>(
condvar: &Condvar,
held: MutexGuard<'a, T>,
bound: Duration,
satisfied: F,
) -> (MutexGuard<'a, T>, std::sync::WaitTimeoutResult)
where
F: FnMut(&mut T) -> bool,
{
match condvar.wait_timeout_while(held, bound, satisfied) {
Ok((woke, verdict)) => (woke, verdict),
Err(poisoned) => poisoned.into_inner(),
}
}
pub(crate) type Job<S, A> =
Box<dyn FnOnce(&mut File, &mut S) -> Result<Stage<A, S>, std::io::Error> + Send>;
pub(crate) enum Stage<A, S> {
Settled(Result<A, std::io::Error>),
Unsynced {
answer: A,
bytes: usize,
settle: Box<dyn FnOnce(&mut S) + Send>,
},
#[cfg(feature = "script")]
Committed(A),
}
impl<A, S> Stage<A, S> {
fn staged_bytes(&self) -> usize {
match *self {
Self::Settled(_) => 0,
#[cfg(feature = "script")]
Self::Committed(_) => 0,
Self::Unsynced { bytes, .. } => bytes,
}
}
fn owes_flush(&self) -> bool {
matches!(self, Self::Unsynced { .. })
}
}
struct Request<S, A> {
job: Mutex<Option<Job<S, A>>>,
slot: Mutex<Option<Result<A, std::io::Error>>>,
arrived: Condvar,
waker: Mutex<Option<Waker>>,
abandoned: Mutex<bool>,
}
impl<S, A> Request<S, A> {
fn new() -> Self {
Self {
job: Mutex::new(None),
slot: Mutex::new(None),
arrived: Condvar::new(),
waker: Mutex::new(None),
abandoned: Mutex::new(false),
}
}
fn arm(&self, job: Job<S, A>) {
*lock(&self.job) = Some(job);
}
fn take_job(&self) -> Result<Job<S, A>, std::io::Error> {
lock(&self.job)
.take()
.ok_or_else(|| std::io::Error::other("the storage owner was handed an unarmed request"))
}
fn abandon(&self) {
*lock(&self.abandoned) = true;
}
fn was_abandoned(&self) -> bool {
*lock(&self.abandoned)
}
fn publish(&self, outcome: Result<A, std::io::Error>) {
*lock(&self.slot) = Some(outcome);
self.wake();
}
fn wake(&self) {
if let Some(waker) = lock(&self.waker).take() {
waker.wake();
}
self.arrived.notify_all();
}
fn take(&self) -> Option<Result<A, std::io::Error>> {
lock(&self.slot).take()
}
fn wait_blocking(&self) -> Option<Result<A, std::io::Error>> {
let mut held = lock(&self.slot);
loop {
if let Some(answer) = held.take() {
return Some(answer);
}
held = wait(&self.arrived, held);
}
}
}
struct Slot<S, A> {
pending: VecDeque<Arc<Request<S, A>>>,
waiting: VecDeque<Arc<Request<S, A>>>,
poisoned: bool,
closed: bool,
released: bool,
fail_next: bool,
fail_next_flush: bool,
flushes: u64,
staged_records: u64,
}
impl<S, A> Slot<S, A> {
fn admission(&self) -> Admission {
if self.poisoned {
Admission::Poisoned
} else if self.pending.len() < DEFAULT_QUEUE_DEPTH {
Admission::Admitted
} else if self.waiting.len() < MAX_WAITING_SUBMITTERS {
Admission::Waiting
} else {
Admission::Refused
}
}
fn has_work(&self) -> bool {
!self.pending.is_empty() || !self.waiting.is_empty()
}
}
enum Admission {
Admitted,
Waiting,
Poisoned,
Refused,
}
impl<S, A> Default for Slot<S, A> {
fn default() -> Self {
Self {
pending: VecDeque::new(),
waiting: VecDeque::new(),
poisoned: false,
closed: false,
released: false,
fail_next: false,
fail_next_flush: false,
flushes: 0,
staged_records: 0,
}
}
}
pub(crate) struct StorageOwner<S, A>
where
S: Send + 'static,
A: Send + 'static,
{
slot: Arc<Mutex<Slot<S, A>>>,
gate: Arc<Mutex<Gate>>,
signal: Arc<Condvar>,
}
#[derive(Clone)]
pub struct StorageGate {
latch: Arc<Mutex<Gate>>,
signal: Arc<Condvar>,
}
#[derive(Default)]
struct Gate {
stalled: bool,
}
impl StorageGate {
pub fn release(&self) {
lock(&self.latch).stalled = false;
self.signal.notify_all();
}
}
#[derive(Debug)]
pub(crate) enum SubmitError {
Poisoned,
QueueFull,
OutcomeUnknown,
Device(std::io::Error),
}
impl SubmitError {
pub(crate) const fn is_outcome_unknown(&self) -> bool {
matches!(self, Self::OutcomeUnknown | Self::Device(_))
}
pub(crate) fn into_io(self) -> std::io::Error {
match self {
Self::Poisoned => std::io::Error::other(
"a previous append's outcome is unknown; reopen the store to replay the truth",
),
Self::QueueFull => std::io::Error::other(
"the store's write queue is full; this append was refused, not queued",
),
Self::OutcomeUnknown => std::io::Error::other(
"the storage owner answered nobody; a prefix of the bytes may be on the disk",
),
Self::Device(cause) => cause,
}
}
}
impl From<std::io::Error> for SubmitError {
fn from(cause: std::io::Error) -> Self {
Self::Device(cause)
}
}
impl<S, A> StorageOwner<S, A>
where
S: Send + 'static,
A: Send + 'static,
{
pub(crate) fn spawn(file: File, state: S, stalled: bool) -> std::io::Result<Self> {
let slot = Arc::new(Mutex::new(Slot::<S, A>::default()));
let gate = Arc::new(Mutex::new(Gate { stalled }));
let signal = Arc::new(Condvar::new());
let owner = Self {
slot: Arc::clone(&slot),
gate: Arc::clone(&gate),
signal: Arc::clone(&signal),
};
let _task = lgwks_std::task::spawn_blocking(move || serve(file, state, slot, gate, signal));
Ok(owner)
}
pub(crate) fn gate(&self) -> StorageGate {
StorageGate {
latch: Arc::clone(&self.gate),
signal: Arc::clone(&self.signal),
}
}
pub(crate) fn poisoned(&self) -> bool {
lock(&self.slot).poisoned
}
#[cfg(feature = "script")]
pub(crate) fn flush_counts(&self) -> (u64, u64) {
let held = lock(&self.slot);
(held.flushes, held.staged_records)
}
#[cfg(test)]
pub(crate) fn fail_next_commit(&self) {
lock(&self.slot).fail_next = true;
}
#[cfg(all(test, feature = "script"))]
pub(crate) fn fail_next_flush(&self) {
lock(&self.slot).fail_next_flush = true;
}
pub(crate) fn submit<F>(&self, job: F) -> Result<A, SubmitError>
where
F: FnOnce(&mut File, &mut S) -> Result<Stage<A, S>, std::io::Error> + Send + 'static,
{
let reply = Arc::new(Request::<S, A>::new());
self.enqueue(Box::new(job), Arc::clone(&reply))?;
match reply.wait_blocking() {
Some(outcome) => outcome.map_err(SubmitError::Device),
None => Err(SubmitError::OutcomeUnknown),
}
}
pub(crate) fn submit_async<F>(&self, job: F) -> crate::BoxFuture<'_, Result<A, SubmitError>>
where
F: FnOnce(&mut File, &mut S) -> Result<Stage<A, S>, std::io::Error> + Send + 'static,
{
let awaiting = self.enqueue_awaiting(job);
Box::pin(awaiting)
}
pub(crate) fn enqueue_awaiting<F>(&self, job: F) -> Awaiting<S, A>
where
F: FnOnce(&mut File, &mut S) -> Result<Stage<A, S>, std::io::Error> + Send + 'static,
{
let reply = Arc::new(Request::<S, A>::new());
let outcome = self.enqueue(Box::new(job), Arc::clone(&reply));
Awaiting {
enqueued: Some(outcome),
reply,
}
}
fn enqueue(&self, job: Job<S, A>, reply: Arc<Request<S, A>>) -> Result<(), SubmitError> {
let mut held = lock(&self.slot);
match held.admission() {
Admission::Poisoned => Err(SubmitError::Poisoned),
Admission::Refused => Err(SubmitError::QueueFull),
Admission::Admitted => {
reply.arm(job);
held.pending.push_back(Arc::clone(&reply));
drop(held);
self.signal.notify_all();
Ok(())
}
Admission::Waiting => {
reply.arm(job);
held.waiting.push_back(Arc::clone(&reply));
drop(held);
self.signal.notify_all();
Ok(())
}
}
}
}
pub(crate) struct Awaiting<S, A> {
enqueued: Option<Result<(), SubmitError>>,
reply: Arc<Request<S, A>>,
}
impl<S, A> std::future::Future for Awaiting<S, A> {
type Output = Result<A, SubmitError>;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.get_mut();
if let Some(Err(cause)) = this.enqueued.take() {
return std::task::Poll::Ready(Err(cause));
}
*lock(&this.reply.waker) = Some(cx.waker().clone());
if let Some(answer) = this.reply.take() {
return std::task::Poll::Ready(answer.map_err(SubmitError::Device));
}
std::task::Poll::Pending
}
}
impl<S, A> Drop for Awaiting<S, A> {
fn drop(&mut self) {
if self.reply.take().is_none() {
self.reply.abandon();
}
}
}
impl<S, A> Drop for StorageOwner<S, A>
where
S: Send + 'static,
A: Send + 'static,
{
fn drop(&mut self) {
let mut slot = lock(&self.slot);
slot.closed = true;
self.signal.notify_all();
while !slot.released {
slot = wait(&self.signal, slot);
}
}
}
fn serve<S, A>(
mut file: File,
mut state: S,
slot: Arc<Mutex<Slot<S, A>>>,
gate: Arc<Mutex<Gate>>,
signal: Arc<Condvar>,
) {
while run_batch(&mut file, &mut state, &slot, &gate, &signal) {}
drop(file);
drop(state);
let mut held = lock(&slot);
held.released = true;
signal.notify_all();
}
struct Member<S, A> {
request: Arc<Request<S, A>>,
staged: Stage<A, S>,
}
fn run_batch<S, A>(
file: &mut File,
state: &mut S,
slot: &Mutex<Slot<S, A>>,
gate: &Mutex<Gate>,
signal: &Condvar,
) -> bool {
promote(&mut lock(slot));
let mut batch: Vec<Member<S, A>> = Vec::new();
let mut staged_bytes = 0usize;
loop {
if batch.is_empty() && !await_work(slot, gate, signal) {
break;
}
if batch.len() >= MAX_BATCH_RECORDS || staged_bytes >= MAX_BATCH_BYTES {
break;
}
let Some(request) = take(slot) else {
break;
};
let outcome = match request.take_job() {
Err(cause) => Err(cause),
Ok(job) => {
let injected = {
let mut held = lock(slot);
std::mem::take(&mut held.fail_next)
};
if injected {
Err(std::io::Error::other("the injected device refusal"))
} else {
job(file, state)
}
}
};
let staged = match outcome {
Ok(staged) => staged,
Err(cause) => Stage::Settled(Err(cause)),
};
staged_bytes = staged_bytes.saturating_add(staged.staged_bytes());
{
let mut held = lock(slot);
held.staged_records = held.staged_records.saturating_add(1);
}
batch.push(Member { request, staged });
}
if batch.is_empty() {
return false;
}
let owes = batch.iter().any(|member| member.staged.owes_flush());
let sync = if owes {
let injected = {
let mut held = lock(slot);
std::mem::take(&mut held.fail_next_flush)
};
if injected {
Err(std::io::Error::other("the injected flush refusal"))
} else {
let outcome = file.sync_all();
let mut held = lock(slot);
held.flushes = held.flushes.saturating_add(1);
outcome
}
} else {
Ok(())
};
let mut failed = false;
let flush_ok = sync.is_ok();
let flush_kind = sync
.as_ref()
.err()
.map_or(std::io::ErrorKind::Other, |cause| cause.kind());
for Member { request, staged } in batch {
let answer = match staged {
Stage::Settled(answer) => answer,
Stage::Unsynced { answer, settle, .. } if flush_ok => {
settle(state);
Ok(answer)
}
Stage::Unsynced { .. } => Err(std::io::Error::new(
flush_kind,
"the batch's flush failed, so this record's outcome is unknown",
)),
#[cfg(feature = "script")]
Stage::Committed(answer) => Ok(answer),
};
if answer.is_err() {
failed = true;
}
if request.was_abandoned() {
failed = true;
}
if failed {
lock(slot).poisoned = true;
}
request.publish(answer);
}
true
}
fn promote<S, A>(held: &mut MutexGuard<'_, Slot<S, A>>) {
while held.pending.len() < DEFAULT_QUEUE_DEPTH {
let Some(request) = held.waiting.pop_front() else {
break;
};
request.wake();
held.pending.push_back(request);
}
}
fn take<S, A>(slot: &Mutex<Slot<S, A>>) -> Option<Arc<Request<S, A>>> {
lock(slot).pending.pop_front()
}
fn await_work<S, A>(slot: &Mutex<Slot<S, A>>, gate: &Mutex<Gate>, signal: &Condvar) -> bool {
let mut held = lock(slot);
loop {
while lock(gate).stalled && !held.closed {
let (guard, _timed_out) = wait_timeout(signal, held, Duration::from_millis(1));
held = guard;
}
promote(&mut held);
if held.has_work() {
return true;
}
if held.closed {
return false;
}
let (guard, _timed_out) = wait_timeout(signal, held, Duration::from_millis(1));
held = guard;
}
}
#[cfg(test)]
mod tests {
use super::{Stage, StorageOwner, lock, wait_timeout_while};
use std::fs::File;
use std::sync::{Arc, Condvar, Mutex};
use std::time::Duration;
#[cfg(feature = "script")]
use std::error::Error;
#[cfg(feature = "script")]
use crate::effect::RunId;
#[cfg(feature = "script")]
use crate::script::run_store::RunRecords;
#[cfg(feature = "script")]
use crate::script::{Scope, Tenant};
#[cfg(feature = "script")]
use crate::task::RunStore;
#[cfg(feature = "script")]
const MEMBERS: usize = 3;
#[cfg(feature = "script")]
const RUN_HEX: [&str; MEMBERS] = [
"01000000000000000000000000000000",
"02000000000000000000000000000000",
"03000000000000000000000000000000",
];
#[cfg(feature = "script")]
const PAYLOAD: [u8; MEMBERS] = [1, 2, 3];
#[cfg(feature = "script")]
#[test]
fn a_failed_batch_flush_acknowledges_nobody_and_folds_nothing() -> Result<(), Box<dyn Error>> {
let dir = crate::journal::file::tests::scratch("owner-flush")?;
std::fs::create_dir_all(&dir)?;
let path = dir.join("store");
let store = RunStore::open_with_stalled_device(&path)?;
let tenant = "flush-tenant";
let step_scope = Scope::root(Tenant::new(tenant)?).enter("step")?;
let key = step_scope.key();
let step_definition = crate::script::run_store::definition_of(&step_scope, "step");
let late_definition = crate::script::run_store::definition_of(&step_scope, "late");
let mut runs = Vec::with_capacity(MEMBERS);
let mut outstanding = Vec::with_capacity(MEMBERS);
for index in 0..MEMBERS {
let run = RunId::from_hex(RUN_HEX[index])?;
runs.push(run);
let record = RunRecords::stage(
&store,
tenant,
run,
key,
"step",
&step_definition,
vec![PAYLOAD[index]],
);
outstanding.push(RunRecords::append_async(&store, record));
}
{
let waker = std::task::Waker::noop();
let mut context = std::task::Context::from_waker(waker);
for future in &mut outstanding {
if future.as_mut().poll(&mut context).is_ready() {
return Err("a parked device answered a request before it was released".into());
}
}
}
store.fail_next_flush();
store.release_device();
for (index, pending) in outstanding.into_iter().enumerate() {
let answered = crate::rt::runtime::block_on(pending);
assert!(
answered.is_err(),
"member {index} of a batch whose flush failed was acknowledged: {answered:?}"
);
}
for (index, run) in runs.iter().enumerate() {
assert!(
!store.knows_run(*run),
"member {index}'s run is in the handle's index after a failed flush"
);
assert_eq!(
store.record_count(*run),
0,
"member {index}'s record was folded into the handle's index after a failed flush"
);
}
for run in runs.iter().take(2) {
let refusal =
RunRecords::append(&store, tenant, *run, key, "late", &late_definition, vec![9])
.err()
.map(|error| error.to_string());
assert!(
refusal
.as_deref()
.is_some_and(|message| message.contains("previous append")),
"a later append must be refused as poisoned, not accepted: {refusal:?}"
);
}
let staged = store.committed_bytes();
drop(store);
let reopened = RunStore::open(&path)?;
assert_eq!(
reopened.committed_bytes(),
std::fs::metadata(&path)?.len(),
"a reopen must consume the whole file, so it replays exactly its durable frames"
);
assert_eq!(
reopened.committed_bytes(),
staged,
"the failed batch's complete frames must be the file's own bytes"
);
for (index, run) in runs.iter().enumerate() {
assert_eq!(
reopened.record_count(*run),
1,
"member {index}'s complete frame is on the file and a reopen reads it, \
though the handle acknowledged none of the batch"
);
}
drop(reopened);
drop(std::fs::remove_dir_all(&dir));
Ok(())
}
const ROUND_TRIPS: u64 = 200_000;
const PATIENCE: Duration = Duration::from_secs(60);
#[test]
fn an_answer_published_during_registration_still_wakes_the_poll()
-> Result<(), Box<dyn std::error::Error>> {
let path = crate::journal::file::tests::scratch("owner-wake")?;
let owner = StorageOwner::<(), u64>::spawn(File::create(&path)?, (), false)?;
let progress = Arc::new((Mutex::new((0_u64, false)), Condvar::new()));
let reported = Arc::clone(&progress);
let driver = lgwks_std::task::spawn_blocking(move || -> Result<(), String> {
let (ref state, ref changed) = *reported;
let mut outcome = Ok(());
for trip in 0..ROUND_TRIPS {
match lgwks_std::task::block_on(
owner.submit_async(move |_, _| Ok(Stage::Settled(Ok(trip)))),
) {
Ok(answer) if answer == trip => {}
other => {
outcome = Err(format!("round trip {trip} answered {other:?}"));
break;
}
}
lock(state).0 = trip + 1;
changed.notify_all();
}
lock(state).1 = true;
changed.notify_all();
outcome
});
let (ref state, ref changed) = *progress;
let (held, waited) =
wait_timeout_while(changed, lock(state), PATIENCE, |progress| !progress.1);
let (completed, done) = *held;
drop(held);
if waited.timed_out() && !done {
return Err(format!(
"round trip {completed} of {ROUND_TRIPS} parked for {PATIENCE:?}: the owner \
published an answer whose poll had not yet registered a waker, so nothing \
woke it"
)
.into());
}
lgwks_std::task::block_on(driver)?;
std::fs::remove_file(&path)?;
Ok(())
}
}