use super::{ExportObservers, SelfStreams};
use crate::ServerError;
use crate::config::types::SelfStreamsConfig;
use crate::server::connection::services::LiminalConnectionServices;
use crate::server::self_streams::cost::{SelfWork, SelfWorkMark, ServerMetadata};
use crate::server::self_streams::events::PublishOrigin;
use crate::server::self_streams::queue::{RecordPayload, RecordQueue};
use liminal::protocol::{CausalContext, MessageEnvelope, SchemaId};
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, MutexGuard, Weak};
use std::thread;
pub(super) struct RecordDrainer {
pub(super) metadata: ServerMetadata,
pub(super) services: Weak<LiminalConnectionServices>,
pub(super) queue: Arc<RecordQueue>,
pub(super) own_work: Arc<Mutex<SelfWork>>,
pub(super) bounds: SelfStreamsConfig,
pub(super) readers: Arc<ExportObservers>,
pub(super) retained_from: Mutex<BTreeMap<&'static str, u64>>,
}
impl RecordDrainer {
pub(super) fn run(self) {
let _self_work = SelfWorkMark::enter();
let panic_guard = DrainerPanicGuard(Arc::clone(&self.queue));
let result = (|| {
while let Some(record) = self.queue.next()? {
let completion = self.queue.append_guard();
let head = match &record.payload {
RecordPayload::Single(payload) => self.publish(record.channel, payload)?,
RecordPayload::Batch(payloads) => self.publish_batch(payloads)?,
};
let sequence = head.checked_sub(1).ok_or_else(|| {
SelfStreams::failure("self-stream append returned an empty head")
})?;
self.retain(record.channel, sequence);
self.queue.appended(&record, head)?;
self.readers.published(record.channel, head);
drop(completion);
}
Ok::<(), ServerError>(())
})();
if let Err(error) = result {
self.queue.fail(&error);
}
drop(panic_guard);
}
fn publish_batch(&self, payloads: &[Vec<u8>]) -> Result<u64, ServerError> {
let services = self
.services
.upgrade()
.ok_or_else(|| SelfStreams::failure("receipt batch services owner is gone"))?;
services.publish_self_stream_batch(payloads, self.metadata.publisher())
}
fn publish(&self, channel: &str, payload: &[u8]) -> Result<u64, ServerError> {
let services = self
.services
.upgrade()
.ok_or_else(|| SelfStreams::failure("connection services owner is gone"))?;
let envelope = MessageEnvelope::new(
SchemaId::from_schema_bytes(b"{}"),
CausalContext::default(),
payload.to_vec(),
);
let mut published = services.publish_with_origin(
channel,
&envelope,
None,
PublishOrigin::Server {
device: self.metadata.device(),
},
);
let mut confirmed_head = None;
let appended = match &mut published {
Ok(publication) => {
confirmed_head = publication
.self_row_at
.map(|sequence| {
sequence
.checked_add(1)
.ok_or_else(|| SelfStreams::failure("self-stream sequence overflow"))
})
.transpose()?;
publication.self_row_appended
}
Err(ServerError::PublishUnresolved {
channel: addressed,
sequence: Some(sequence),
message,
..
}) if addressed.as_str() == channel => {
match self.confirm_unresolved_row(&services, channel, *sequence, &envelope) {
Ok(stored) => {
if stored {
confirmed_head = Some(sequence.checked_add(1).ok_or_else(|| {
SelfStreams::failure("confirmed self-stream sequence overflow")
})?);
}
stored
}
Err(error) => {
*message = format!("{message}; self-row confirmation failed: {error}");
false
}
}
}
Err(_) => false,
};
if appended {
return confirmed_head.map_or_else(|| services.self_stream_head(channel), Ok);
}
published?;
Err(SelfStreams::failure(
"self-stream publication was refused before append",
))
}
fn retain(&self, channel: &'static str, sequence: u64) {
let bound = match channel {
"liminal.cost" => self.bounds.cost_rows,
_ => self.bounds.events_rows,
};
let keep = bound - bound / 8;
let cut = match self.retained_from.lock() {
Ok(retained_from) => {
let first = retained_from.get(channel).copied().unwrap_or(0);
let retained = sequence.saturating_add(1).saturating_sub(first);
if retained <= bound {
return;
}
sequence.saturating_add(1).saturating_sub(keep)
}
Err(error) => {
self.prune_failed(channel, &format!("retention lock poisoned: {error}"));
return;
}
};
let pruned = self
.services
.upgrade()
.ok_or_else(|| "connection services owner is gone".to_owned())
.and_then(|services| {
let key = format!("{channel}:0");
liminal::durability::bridge::block_on(
services.durable_store().truncate_before(&key, cut),
)
.map_err(|error| error.to_string())?
.map_err(|error| error.to_string())
});
match pruned {
Ok(rows) => {
let recorded = self
.retained_from
.lock()
.map(|mut retained_from| {
retained_from.insert(channel, cut);
})
.map_err(|error| format!("retention lock poisoned: {error}"))
.and_then(|()| {
self.own_work()
.map(|mut own_work| {
own_work.prunes = own_work.prunes.saturating_add(1);
own_work.rows_pruned = own_work.rows_pruned.saturating_add(rows);
})
.map_err(|error| error.to_string())
});
if let Err(error) = recorded {
self.prune_failed(channel, &error);
}
}
Err(error) => self.prune_failed(channel, &error),
}
}
fn prune_failed(&self, channel: &str, error: &str) {
tracing::warn!(stream = channel, %error, "SELF_STREAM_PRUNE_FAILED");
match self.own_work() {
Ok(mut own_work) => {
own_work.prune_failures = own_work.prune_failures.saturating_add(1);
}
Err(error) => {
tracing::error!(stream = channel, %error, "SELF_STREAM_PRUNE_FAILED not counted");
}
}
}
fn own_work(&self) -> Result<MutexGuard<'_, SelfWork>, ServerError> {
self.own_work
.lock()
.map_err(|error| SelfStreams::failure(&format!("own-work lock poisoned: {error}")))
}
fn confirm_unresolved_row(
&self,
services: &LiminalConnectionServices,
channel: &str,
sequence: u64,
expected: &MessageEnvelope,
) -> Result<bool, ServerError> {
let store = services.durable_store();
let key = format!("{channel}:0");
let row = liminal::durability::bridge::block_on(store.read_at(&key, sequence))
.map_err(|error| {
SelfStreams::failure(&format!("self-row confirmation bridge: {error}"))
})?
.map_err(|error| {
SelfStreams::failure(&format!("self-row confirmation read: {error}"))
})?;
let Some(row) = row else {
return Ok(false);
};
if row.sequence != sequence {
return Err(SelfStreams::failure(
"self-row confirmation returned a different sequence",
));
}
let stored =
liminal::durability::MessageEnvelope::deserialize(&row.payload).map_err(|error| {
SelfStreams::failure(&format!("self-row confirmation envelope: {error}"))
})?;
Ok(stored.payload == expected.payload
&& stored.publisher_id == self.metadata.publisher()
&& stored.idempotency_key.is_none())
}
}
struct DrainerPanicGuard(Arc<RecordQueue>);
impl Drop for DrainerPanicGuard {
fn drop(&mut self) {
if thread::panicking() {
self.0
.fail(&SelfStreams::failure("record drainer panicked"));
}
}
}