use std::fmt;
use std::num::NonZeroUsize;
use std::time::Duration;
use tracing::error;
use crate::config::{Config, StreamKind};
use super::engine::StreamBroadcast;
use super::replay_cap::ReplayCap;
pub(crate) const STREAM_CHANNEL_CAPACITY: usize = 500;
#[derive(Debug, Clone, Copy)]
pub(crate) struct StreamOptions {
pub(crate) replay_chunk_size: Option<NonZeroUsize>,
/// none for in-process subscribers, which pay for their own replays
pub(crate) replay_cap: Option<ReplayCap>,
pub(crate) pending_event_limit: usize,
pub(crate) send_timeout: Duration,
}
impl StreamOptions {
pub(crate) fn from_config(config: &Config) -> Self {
Self {
replay_chunk_size: NonZeroUsize::new(config.stream_replay_chunk_size),
replay_cap: None,
pending_event_limit: config.stream_pending_event_limit.max(1),
send_timeout: config.stream_send_timeout,
}
}
/// for a subscriber on the network, whose cursor replay is capped when `stream` is
/// throttled.
pub(crate) fn network(config: &Config, stream: StreamKind) -> Self {
Self {
replay_cap: ReplayCap::from_config(config, stream),
..Self::from_config(config)
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct StreamTooSlow {
pub(crate) reason: String,
}
impl StreamTooSlow {
pub(crate) fn send_timeout(timeout: Duration) -> Self {
Self {
reason: format!(
"stream delivery blocked for at least {} seconds",
timeout.as_secs()
),
}
}
}
impl fmt::Display for StreamTooSlow {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "stream consumer too slow: {}", self.reason)
}
}
#[cfg(feature = "indexer_stream")]
impl From<StreamTooSlow> for super::super::StreamError {
fn from(err: StreamTooSlow) -> Self {
Self::ConsumerTooSlow { reason: err.reason }
}
}
#[cfg(feature = "relay")]
impl From<StreamTooSlow> for super::super::RelayStreamError {
fn from(err: StreamTooSlow) -> Self {
Self::ConsumerTooSlow { reason: err.reason }
}
}
#[cfg(feature = "jetstream")]
impl From<StreamTooSlow> for super::super::JetstreamStreamError {
fn from(err: StreamTooSlow) -> Self {
Self::ConsumerTooSlow { reason: err.reason }
}
}
#[cfg(feature = "indexer_stream")]
impl StreamBroadcast for crate::types::BroadcastEvent {
type Position = u64;
fn position(&self) -> u64 {
match self {
Self::Persisted(id) => *id,
Self::LiveRecord(evt) => evt.id,
Self::Ephemeral(evt) => evt.id,
}
}
fn is_marker(&self) -> bool {
matches!(self, Self::Persisted(_))
}
}
#[cfg(feature = "relay")]
impl StreamBroadcast for crate::types::RelayBroadcast {
type Position = u64;
fn position(&self) -> u64 {
match self {
Self::Persisted(seq) | Self::Ephemeral(seq, _) => *seq,
}
}
fn is_marker(&self) -> bool {
matches!(self, Self::Persisted(_))
}
}
#[cfg(feature = "jetstream")]
impl StreamBroadcast for crate::types::JetstreamBroadcast {
type Position = crate::types::JetstreamPosition;
fn position(&self) -> Self::Position {
match self {
Self::Historical(position) => *position,
Self::Live(live) => live.position,
}
}
fn is_marker(&self) -> bool {
matches!(self, Self::Historical(_))
}
}
/// a db read failed partway. the rows past it can't be skipped, so the
/// stream ends and its subscriber resumes from the last event it got.
#[derive(Debug)]
pub(crate) struct ReplayFailed;
/// outputs for the rows one db read scanned.
pub(crate) struct ReplayChunk<T, P> {
pub(crate) events: Vec<T>,
/// the last row scanned, sent or filtered out
pub(crate) last_seen: Option<P>,
/// no rows are left in the range read
pub(crate) exhausted: bool,
/// rows the read went through, sent or not, since a capped replay pays for dropped ones too
pub(crate) scanned: usize,
}
impl<T, P> ReplayChunk<T, P> {
/// read `rows` after `after` into up to `limit` outputs. scans at most
/// four times as many rows, so a range the subscriber mostly filters out
/// still returns promptly. rows whose key doesn't parse are skipped.
pub(crate) fn read(
mut rows: fjall::Iter,
after: Option<P>,
limit: usize,
position: impl Fn(&[u8]) -> Option<P>,
mut output: impl FnMut(P, &[u8]) -> Option<T>,
) -> Result<Self, ReplayFailed>
where
P: Copy,
{
let mut chunk = Self {
events: Vec::with_capacity(limit),
last_seen: after,
exhausted: false,
scanned: 0,
};
let max_scanned = limit.saturating_mul(4).max(limit);
while chunk.events.len() < limit && chunk.scanned < max_scanned {
let Some(row) = rows.next() else {
chunk.exhausted = true;
break;
};
chunk.scanned += 1;
let (key, value) = row.into_inner().map_err(|e| {
error!(err = %e, "failed to read stream rows from db, closing the stream");
ReplayFailed
})?;
let Some(position) = position(&key) else {
continue;
};
chunk.last_seen = Some(position);
chunk.events.extend(output(position, &value));
}
Ok(chunk)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn only_network_subscribers_are_capped() {
let all = [
StreamKind::Indexer,
StreamKind::Jetstream,
StreamKind::Relay,
];
let config = Config {
throttled_streams: all.to_vec(),
..Config::default()
};
assert_eq!(StreamOptions::from_config(&config).replay_cap, None);
for stream in all {
assert!(StreamOptions::network(&config, stream).replay_cap.is_some());
}
}
/// keys 1 through 10 as big-endian positions, values naming them, plus one
/// key that isn't a position between 3 and 4
fn rows() -> (tempfile::TempDir, fjall::Database, fjall::Keyspace) {
let tmp = tempfile::tempdir().unwrap();
let db = fjall::Database::builder(tmp.path()).open().unwrap();
let ks = db
.keyspace("rows", fjall::KeyspaceCreateOptions::default)
.unwrap();
for n in 1..=10_u64 {
ks.insert(n.to_be_bytes(), format!("v{n}")).unwrap();
}
ks.insert([0, 0, 0, 0, 0, 0, 0, 3, 0], "junk").unwrap();
(tmp, db, ks)
}
fn position(key: &[u8]) -> Option<u64> {
<[u8; 8]>::try_from(key).ok().map(u64::from_be_bytes)
}
fn evens(position: u64, value: &[u8]) -> Option<(u64, String)> {
position
.is_multiple_of(2)
.then(|| (position, String::from_utf8(value.to_vec()).unwrap()))
}
#[test]
fn a_chunk_stops_at_its_limit_and_remembers_filtered_rows() {
let (_tmp, _db, ks) = rows();
let chunk = ReplayChunk::read(ks.iter(), None, 2, position, evens).unwrap();
assert_eq!(chunk.events, [(2, "v2".into()), (4, "v4".into())]);
assert_eq!(chunk.last_seen, Some(4));
assert!(!chunk.exhausted);
// filtered rows and the junk key count, since reading them cost the same
assert_eq!(chunk.scanned, 5);
// four rows scanned for one wanted output, the junk key among them
let chunk = ReplayChunk::read(ks.iter(), None, 1, position, |_, _| None::<()>).unwrap();
assert!(chunk.events.is_empty());
assert_eq!(chunk.last_seen, Some(3));
assert!(!chunk.exhausted);
assert_eq!(chunk.scanned, 4);
let chunk = ReplayChunk::read(ks.range(9_u64.to_be_bytes()..), Some(8), 4, position, evens)
.unwrap();
assert_eq!(chunk.events, [(10, "v10".into())]);
assert_eq!(chunk.last_seen, Some(10));
assert!(chunk.exhausted);
let chunk = ReplayChunk::read(
ks.range(11_u64.to_be_bytes()..),
Some(10),
4,
position,
evens,
)
.unwrap();
assert!(chunk.events.is_empty());
assert_eq!(chunk.last_seen, Some(10));
assert!(chunk.exhausted);
}
}