use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Condvar, Mutex, MutexGuard, OnceLock};
use crate::ServerError;
use crate::server::shutdown::ShutdownHandle;
use super::cost::SelfWork;
use super::sink::SelfStreams;
pub(super) const CAPACITY: usize = 1024;
#[derive(Debug)]
pub(super) struct QueuedRecord {
pub(super) sequence: u64,
pub(super) channel: &'static str,
pub(super) payload: RecordPayload,
#[cfg(test)]
pub(super) connection_opened_count: usize,
}
#[derive(Debug)]
pub(super) enum RecordPayload {
Single(Vec<u8>),
Batch(Vec<Vec<u8>>),
}
impl RecordPayload {
pub(super) fn len(&self) -> usize {
match self {
Self::Single(_) => 1,
Self::Batch(rows) => rows.len(),
}
}
}
#[derive(Debug)]
struct QueueState {
pending: VecDeque<QueuedRecord>,
pending_rows: usize,
in_flight: bool,
next_sequence: u64,
acknowledged: Option<u64>,
stopped: bool,
own_work: SelfWork,
heads: [u64; 2],
#[cfg(test)]
takes_left: Option<usize>,
}
#[derive(Debug)]
pub(super) struct RecordQueue {
state: Mutex<QueueState>,
changed: Condvar,
fatal: OnceLock<String>,
abandoned: AtomicBool,
shutdown: ShutdownHandle,
#[cfg(test)]
event_completed: tokio::sync::Notify,
}
impl RecordQueue {
pub(super) fn new(shutdown: ShutdownHandle) -> Self {
let queue = Self {
state: Mutex::new(QueueState {
pending: VecDeque::with_capacity(CAPACITY),
pending_rows: 0,
in_flight: false,
next_sequence: 0,
acknowledged: None,
stopped: false,
own_work: SelfWork::default(),
heads: [0; 2],
#[cfg(test)]
takes_left: None,
}),
changed: Condvar::new(),
fatal: OnceLock::new(),
abandoned: AtomicBool::new(false),
shutdown,
#[cfg(test)]
event_completed: tokio::sync::Notify::new(),
};
if let Err(error) = queue.initialize() {
queue.fail(&error);
}
queue
}
fn initialize(&self) -> Result<(), ServerError> {
let state = self.lock()?;
self.changed.notify_all();
drop(state);
Ok(())
}
fn lock(&self) -> Result<MutexGuard<'_, QueueState>, ServerError> {
self.state
.lock()
.map_err(|error| SelfStreams::failure(&format!("record queue lock poisoned: {error}")))
}
pub(super) fn failure(&self) -> Option<ServerError> {
self.fatal
.get()
.map(|message| SelfStreams::failure(message))
}
pub(super) fn fail(&self, error: &ServerError) {
if self.fatal.set(error.to_string()).is_ok() {
crate::metrics::self_stream_append_failed();
tracing::error!(%error, "SELF_STREAM_EVENTS_APPEND_FAILED");
}
match self.state.lock() {
Ok(state) => {
self.changed.notify_all();
drop(state);
}
Err(error) => {
tracing::error!(%error, "self-stream failure notification lock poisoned");
self.changed.notify_all();
}
}
self.shutdown.initiate();
}
pub(super) fn abandon(&self, error: &ServerError) {
self.abandoned.store(true, Ordering::Release);
self.fail(error);
}
pub(super) fn was_abandoned(&self) -> bool {
self.abandoned.load(Ordering::Acquire)
}
pub(super) fn enqueue(&self, mut record: QueuedRecord) -> Result<u64, ServerError> {
let mut state = self.lock()?;
if let Some(error) = self.failure() {
return Err(error);
}
let rows = record.payload.len();
if rows == 0 {
return Err(SelfStreams::failure("record queue batch is empty"));
}
if state.stopped || rows > CAPACITY - state.pending_rows {
let reason = if state.stopped {
"record queue is closed"
} else {
"SELF_STREAM_QUEUE_OVERFLOW: capacity 1024 pending records"
};
drop(state);
let error = SelfStreams::failure(reason);
self.fail(&error);
return Err(error);
}
let count = u64::try_from(rows)
.map_err(|error| SelfStreams::failure(&format!("record queue row count: {error}")))?;
let next = state
.next_sequence
.checked_add(count)
.ok_or_else(|| SelfStreams::failure("record queue sequence overflow"))?;
let sequence = next - 1;
record.sequence = sequence;
state.pending_rows += rows;
state.pending.push_back(record);
state.next_sequence = next;
self.changed.notify_all();
drop(state);
Ok(sequence)
}
pub(super) fn next(&self) -> Result<Option<QueuedRecord>, ServerError> {
let mut state = self.lock()?;
loop {
if let Some(error) = self.failure() {
return Err(error);
}
if let Some(mut record) = Self::take(&mut state) {
state.pending_rows -= record.payload.len();
if record.channel == "liminal.events"
&& matches!(&record.payload, RecordPayload::Single(_))
&& state.pending.front().is_some_and(single_event)
{
let RecordPayload::Single(first) = record.payload else {
return Err(SelfStreams::failure("event batch starts with a batch"));
};
let mut payloads = Vec::with_capacity(state.pending_rows + 1);
payloads.push(first);
while state.pending.front().is_some_and(single_event) {
let next = state.pending.pop_front().ok_or_else(|| {
SelfStreams::failure("event batch backlog changed while locked")
})?;
let RecordPayload::Single(payload) = next.payload else {
return Err(SelfStreams::failure("event batch contains a batch"));
};
state.pending_rows -= 1;
record.sequence = next.sequence;
#[cfg(test)]
{
record.connection_opened_count = record
.connection_opened_count
.checked_add(next.connection_opened_count)
.ok_or_else(|| {
SelfStreams::failure("connection-open count overflow")
})?;
}
payloads.push(payload);
}
record.payload = RecordPayload::Batch(payloads);
}
state.in_flight = true;
return Ok(Some(record));
}
if state.stopped {
return Ok(None);
}
state = self.changed.wait(state).map_err(|error| {
SelfStreams::failure(&format!("record queue readiness lock poisoned: {error}"))
})?;
}
}
fn take(state: &mut QueueState) -> Option<QueuedRecord> {
#[cfg(test)]
if !state.pending.is_empty() {
match &mut state.takes_left {
Some(0) => return None,
Some(left) => *left -= 1,
None => {}
}
}
state.pending.pop_front()
}
pub(super) fn appended(&self, record: &QueuedRecord, head: u64) -> Result<(), ServerError> {
let mut state = self.lock()?;
let index = match record.channel {
"liminal.events" => 0,
"liminal.cost" => 1,
_ => {
return Err(SelfStreams::failure(
"record queue channel is not a self stream",
));
}
};
if head == 0 || head < state.heads[index] {
return Err(SelfStreams::failure(
"acknowledged self-stream head regressed",
));
}
state.heads[index] = head;
state.own_work.rows = state
.own_work
.rows
.checked_add(
u64::try_from(record.payload.len())
.map_err(|error| SelfStreams::failure(&format!("own-row count: {error}")))?,
)
.ok_or_else(|| SelfStreams::failure("own-row count overflow"))?;
if matches!(record.payload, RecordPayload::Batch(_)) {
state.own_work.flushes = state
.own_work
.flushes
.checked_add(1)
.ok_or_else(|| SelfStreams::failure("own-flush count overflow"))?;
}
state.acknowledged = Some(record.sequence);
self.changed.notify_all();
drop(state);
#[cfg(test)]
for _ in 0..record.connection_opened_count {
self.event_completed.notify_one();
}
Ok(())
}
pub(super) fn await_append(&self, sequence: u64) -> Result<(), ServerError> {
let mut state = self.lock()?;
loop {
if let Some(error) = self.failure() {
return Err(error);
}
if state.acknowledged.is_some_and(|ack| ack >= sequence) {
return Ok(());
}
state = self.changed.wait(state).map_err(|error| {
SelfStreams::failure(&format!("record queue completion lock poisoned: {error}"))
})?;
}
}
pub(super) fn drain(&self) -> Result<(), ServerError> {
let sequence = self.lock()?.next_sequence.checked_sub(1);
sequence.map_or_else(
|| self.failure().map_or(Ok(()), Err),
|sequence| self.await_append(sequence),
)
}
pub(super) const fn append_guard(&self) -> InFlightAppend<'_> {
InFlightAppend(self)
}
fn complete_in_flight(&self) -> Result<(), ServerError> {
let mut state = self.lock()?;
state.in_flight = false;
self.changed.notify_all();
drop(state);
Ok(())
}
fn await_idle(&self) -> Result<(), ServerError> {
let mut state = self.lock()?;
while state.in_flight {
if self.was_abandoned() {
if let Some(error) = self.failure() {
return Err(error);
}
}
state = self.changed.wait(state).map_err(|error| {
SelfStreams::failure(&format!("record drainer completion lock poisoned: {error}"))
})?;
}
drop(state);
Ok(())
}
pub(super) fn drain_all(&self) -> Result<(), ServerError> {
loop {
self.await_idle()?;
if let Some(error) = self.failure() {
return Err(error);
}
let next = self.lock()?.next_sequence;
if let Some(sequence) = next.checked_sub(1) {
self.await_append(sequence)?;
}
self.await_idle()?;
if self.lock()?.next_sequence == next {
return self.failure().map_or(Ok(()), Err);
}
}
}
pub(super) fn heads(&self) -> Result<[u64; 2], ServerError> {
Ok(self.lock()?.heads)
}
pub(super) fn own_work(&self) -> Result<SelfWork, ServerError> {
Ok(self.lock()?.own_work)
}
pub(super) fn stop(&self) -> Result<(), ServerError> {
let mut state = self.lock()?;
state.stopped = true;
self.changed.notify_all();
drop(state);
Ok(())
}
#[cfg(test)]
pub(super) async fn event_completed(&self) {
self.event_completed.notified().await;
}
#[cfg(test)]
pub(super) fn hold(&self, takes: usize) -> Result<(), ServerError> {
let mut state = self.lock()?;
state.takes_left = Some(takes);
self.changed.notify_all();
drop(state);
Ok(())
}
#[cfg(test)]
pub(super) fn release(&self) -> Result<(), ServerError> {
let mut state = self.lock()?;
state.takes_left = None;
self.changed.notify_all();
drop(state);
Ok(())
}
#[cfg(test)]
pub(super) fn measurements(&self) -> Result<(usize, usize, Option<u64>), ServerError> {
let state = self.lock()?;
Ok((
state.pending_rows,
state.pending.capacity(),
state.acknowledged,
))
}
}
fn single_event(record: &QueuedRecord) -> bool {
record.channel == "liminal.events" && matches!(&record.payload, RecordPayload::Single(_))
}
pub(super) struct InFlightAppend<'a>(&'a RecordQueue);
impl Drop for InFlightAppend<'_> {
fn drop(&mut self) {
if let Err(error) = self.0.complete_in_flight() {
self.0.fail(&error);
}
}
}