use std::cell::Cell;
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, Weak};
use std::thread::{self, JoinHandle};
use std::time::{SystemTime, UNIX_EPOCH};
use liminal::durability::{StoreCounts, StoreCountsHandle};
use crate::ServerError;
use crate::config::types::SelfStreamsConfig;
use crate::server::connection::services::{LiminalConnectionServices, PublishOutcome};
use crate::server::shutdown::ShutdownHandle;
use super::cost::{CostWake, PreparedCost, SelfWork, ServerMetadata};
use super::events::ServerEvent;
use super::publisher::CostSink;
use super::queue::{QueuedRecord, RecordPayload, RecordQueue};
mod bridge;
mod drainer;
mod stop;
use bridge::ExportObservers;
use drainer::RecordDrainer;
thread_local! {
static FINAL_FLUSH: Cell<bool> = const { Cell::new(false) };
}
#[derive(Debug)]
pub(in crate::server) struct OriginPublication {
pub(in crate::server) outcome: PublishOutcome,
pub(in crate::server) self_row_appended: bool,
pub(in crate::server) self_row_at: Option<u64>,
}
impl OriginPublication {
pub(in crate::server) const fn connection(outcome: PublishOutcome) -> Self {
Self {
outcome,
self_row_appended: false,
self_row_at: None,
}
}
}
pub(in crate::server) struct SelfStreams {
metadata: ServerMetadata,
counts: StoreCountsHandle,
queue: Arc<RecordQueue>,
services: Weak<LiminalConnectionServices>,
drainer: Mutex<DrainerOwner>,
wake: Arc<CostWake>,
retention_work: Arc<Mutex<SelfWork>>,
readers: Arc<ExportObservers>,
}
struct DrainerOwner {
worker: Option<RecordDrainer>,
thread: Option<JoinHandle<()>>,
stopped: bool,
}
impl std::fmt::Debug for SelfStreams {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("SelfStreams")
.field("metadata", &self.metadata)
.finish_non_exhaustive()
}
}
impl SelfStreams {
pub(in crate::server) fn new(
metadata: ServerMetadata,
services: &Arc<LiminalConnectionServices>,
wake: Arc<CostWake>,
shutdown: ShutdownHandle,
bounds: SelfStreamsConfig,
) -> Self {
let counts = services.store_counts();
let queue = Arc::new(RecordQueue::new(shutdown));
let retention_work = Arc::new(Mutex::new(SelfWork::default()));
let readers = Arc::new(ExportObservers::default());
let worker = RecordDrainer {
metadata: metadata.clone(),
services: Arc::downgrade(services),
queue: Arc::clone(&queue),
own_work: Arc::clone(&retention_work),
bounds,
readers: Arc::clone(&readers),
retained_from: Mutex::new(BTreeMap::new()),
};
Self {
metadata,
counts,
queue,
services: Arc::downgrade(services),
wake,
retention_work,
readers,
drainer: Mutex::new(DrainerOwner {
worker: Some(worker),
thread: None,
stopped: false,
}),
}
}
pub(in crate::server) fn observe_bridge(
&self,
services: &LiminalConnectionServices,
feeds: &Arc<crate::link_bridge::subscriptions::ExportSubscriptions>,
) -> Result<(), ServerError> {
let work = super::cost::ReaderWork::enter();
if feeds.install_stored(&services.durable_store())? {
self.readers.install(feeds)?;
for subject in ["liminal.cost", "liminal.events"] {
feeds.stored_published(subject, services.self_stream_head(subject)?);
}
}
drop(work);
Ok(())
}
pub(in crate::server) fn start(&self) -> Result<(), ServerError> {
let mut owner = self.drainer.lock().map_err(|error| {
Self::failure(&format!("record drainer start lock poisoned: {error}"))
})?;
if owner.stopped {
return Err(Self::failure("record drainer is stopped"));
}
if owner.thread.is_some() {
return Ok(());
}
let worker = owner
.worker
.take()
.ok_or_else(|| Self::failure("record drainer unavailable after start failure"))?;
owner.thread = Some(
thread::Builder::new()
.name("liminal-self-records".to_owned())
.spawn(move || worker.run())
.map_err(|error| Self::failure(&format!("record drainer construction: {error}")))?,
);
drop(owner);
Ok(())
}
pub(in crate::server) fn failure(message: &str) -> ServerError {
ServerError::ListenerAccept {
message: format!("SELF_STREAM_EVENTS_APPEND_FAILED: {message}"),
}
}
pub(in crate::server) fn drain(&self) -> Result<(), ServerError> {
self.queue.drain()
}
#[cfg(test)]
pub(in crate::server) fn hold_drainer(&self, takes: usize) -> Result<(), ServerError> {
self.queue.hold(takes)
}
#[cfg(test)]
pub(in crate::server) fn release_drainer(&self) -> Result<(), ServerError> {
self.queue.release()
}
#[cfg(test)]
pub(in crate::server) fn drain_all(&self) -> Result<(), ServerError> {
self.queue.drain_all()
}
#[cfg(test)]
pub(super) fn queue_for_test(&self) -> Arc<RecordQueue> {
Arc::clone(&self.queue)
}
#[cfg(test)]
pub(super) fn enqueue_event_for_test(&self, event: ServerEvent) -> Result<u64, ServerError> {
let payload = event
.encode(&self.metadata, server_time_ms()?)
.map_err(|error| Self::failure(&error.to_string()))?;
self.enqueue("liminal.events", payload, false)
}
#[cfg(test)]
pub(super) fn join_drainer_for_test(&self) -> Result<(), ServerError> {
self.queue.stop()?;
let thread = {
let mut owner = self
.drainer
.lock()
.map_err(|error| Self::failure(&format!("drainer join lock: {error}")))?;
owner.stopped = true;
drop(owner.worker.take());
owner.thread.take()
};
if let Some(thread) = thread {
thread
.join()
.map_err(|_| Self::failure("record drainer join panicked"))?;
}
Ok(())
}
pub(in crate::server) fn with_flush_barrier<T>(
&self,
flush: impl FnOnce() -> Result<T, ServerError>,
) -> Result<T, ServerError> {
let before = self.queue.drain_all();
if self.queue.was_abandoned() {
before?;
return Err(Self::failure("owned durable drain was abandoned"));
}
if FINAL_FLUSH.with(|active| active.replace(true)) {
return Err(Self::failure("nested final flush barrier"));
}
let reset = FlushBarrier;
let result = flush();
drop(reset);
let after = self.queue.drain_all();
if let Err(error) = &before {
tracing::error!(%error, "self-stream pre-flush drain failed");
}
if let Err(error) = &result {
tracing::error!(%error, "self-stream backing-store flush failed");
}
if let Err(error) = &after {
tracing::error!(%error, "self-stream post-flush drain failed");
}
let value = result?;
before?;
after?;
Ok(value)
}
fn enqueue(
&self,
channel: &'static str,
payload: Vec<u8>,
#[cfg(test)] connection_opened: bool,
) -> Result<u64, ServerError> {
let result = (|| {
self.start()?;
self.queue.enqueue(QueuedRecord {
sequence: 0,
channel,
payload: RecordPayload::Single(payload),
#[cfg(test)]
connection_opened_count: usize::from(connection_opened),
})
})();
if let Err(error) = &result {
self.queue.fail(error);
}
result
}
pub(in crate::server) fn lifecycle_batch<'a>(
&self,
conversation: u64,
receipts: impl IntoIterator<Item = (u64, &'a str)>,
) -> Result<(), ServerError> {
let result = (|| {
let time = server_time_ms()?;
let payloads = receipts
.into_iter()
.map(|(generation, transition)| {
ServerEvent::Lifecycle {
conversation,
generation,
transition,
}
.encode(&self.metadata, time)
.map_err(|error| Self::failure(&error.to_string()))
})
.collect::<Result<Vec<_>, _>>()?;
if payloads.is_empty() {
return Ok(());
}
self.start()?;
let sequence = self.queue.enqueue(QueuedRecord {
sequence: 0,
channel: "liminal.events",
payload: RecordPayload::Batch(payloads),
#[cfg(test)]
connection_opened_count: 0,
})?;
self.wake.arm_for_work();
self.queue.await_append(sequence)
})();
if let Err(error) = &result {
self.queue.fail(error);
}
result
}
pub(in crate::server) fn event(&self, event: ServerEvent<'_>) {
if !matches!(event, ServerEvent::Flush { .. }) {
self.wake.arm_for_work();
}
let wait = matches!(event, ServerEvent::Flush { .. }) && FINAL_FLUSH.with(Cell::get);
let result = (|| {
let payload = event
.encode(&self.metadata, server_time_ms()?)
.map_err(|error| Self::failure(&error.to_string()))?;
let sequence = self.enqueue(
"liminal.events",
payload,
#[cfg(test)]
matches!(event, ServerEvent::ConnectionOpened { .. }),
)?;
if wait {
self.queue.await_append(sequence)?;
let services = self
.services
.upgrade()
.ok_or_else(|| Self::failure("head checkpoint services owner is gone"))?;
services.checkpoint_self_stream_heads(self.queue.heads()?)?;
}
Ok(())
})();
if let Err(error) = result {
self.queue.fail(&error);
}
}
#[cfg(test)]
pub(super) async fn event_completed(&self) {
self.queue.event_completed().await;
}
#[cfg(test)]
pub(super) fn own_work_for_test(&self) -> Result<SelfWork, ServerError> {
self.queue.own_work()
}
#[cfg(test)]
pub(super) fn queue_measurements(&self) -> Result<(usize, usize, Option<u64>), ServerError> {
self.queue.measurements()
}
}
impl CostSink for Arc<SelfStreams> {
fn completion_failure_signal(&self) -> Option<super::publisher::CompletionFailureSignal> {
let queue = Arc::clone(&self.queue);
Some(Arc::new(move |error| queue.abandon(error)))
}
fn before_stop(&mut self) -> Result<(), ServerError> {
let before = self.counts.snapshot();
self.queue.drain_all()?;
tracing::debug!(counts = ?self.counts.snapshot().since(&before), "self-stream final queued drain store calls");
Ok(())
}
fn snapshot(&self) -> Result<(StoreCounts, SelfWork), ServerError> {
self.drain()?;
let mut own_work = self.queue.own_work()?;
let retention = self.retention_work.lock().map_err(|error| {
SelfStreams::failure(&format!("retention work lock poisoned: {error}"))
})?;
own_work.prunes = retention.prunes;
own_work.rows_pruned = retention.rows_pruned;
own_work.prune_failures = retention.prune_failures;
drop(retention);
Ok((self.counts.snapshot(), own_work))
}
fn publish(&mut self, record: &PreparedCost) -> Result<(), ServerError> {
let before = self.counts.snapshot();
let payload = record
.encode(&self.metadata, server_time_ms()?)
.map_err(|error| SelfStreams::failure(&error.to_string()))?;
let sequence = self.enqueue(
"liminal.cost",
payload,
#[cfg(test)]
false,
)?;
tracing::debug!(sequence, "self-stream cost record awaits append");
self.queue.await_append(sequence)?;
tracing::debug!(sequence, "self-stream cost record append completed");
tracing::debug!(counts = ?self.counts.snapshot().since(&before), "self-stream final cost store calls");
Ok(())
}
fn console_wake(&self) -> Option<Arc<tokio::sync::Notify>> {
self.services.upgrade()?.signals().console_wake()
}
fn pump_console(&mut self) {
let Some(services) = self.services.upgrade() else {
return;
};
if let Err(error) = services.console().pump(&*services) {
tracing::warn!(%error, "CONSOLE_PUMP_FAILED");
}
}
}
struct FlushBarrier;
impl Drop for FlushBarrier {
fn drop(&mut self) {
FINAL_FLUSH.with(|active| active.set(false));
}
}
impl Drop for SelfStreams {
fn drop(&mut self) {
if let Err(error) = self.queue.stop() {
self.queue.fail(&error);
}
let drainer = match self.drainer.get_mut() {
Ok(drainer) => drainer.thread.take(),
Err(error) => {
self.queue.fail(&Self::failure(&format!(
"record drainer join lock poisoned: {error}"
)));
None
}
};
if let Some(drainer) = drainer {
if drainer.thread().id() == thread::current().id() {
tracing::debug!(
"self-stream drainer drops its final services owner and is stopping"
);
} else if self.queue.was_abandoned() && !drainer.is_finished() {
tracing::error!(
"self-stream drainer completion remains unconfirmed after owned drain abandonment"
);
} else if drainer.join().is_err() {
self.queue
.fail(&Self::failure("record drainer join panicked"));
}
}
}
}
fn server_time_ms() -> Result<u64, ServerError> {
let elapsed = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|error| ServerError::ListenerAccept {
message: format!("self-stream server clock is before the epoch: {error}"),
})?;
u64::try_from(elapsed.as_millis()).map_err(|error| ServerError::ListenerAccept {
message: format!("self-stream server time cannot fit in milliseconds: {error}"),
})
}