use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::error;
use crate::db::keys;
use crate::state::AppState;
use crate::types::RelayBroadcast;
use super::{ReplayChunk, ReplayFailed, Start, StreamOptions, StreamSource, run_stream};
use crate::control::RelayStreamError;
pub(crate) fn relay_stream_thread(
state: Arc<AppState>,
tx: mpsc::Sender<Result<bytes::Bytes, RelayStreamError>>,
cursor: Option<u64>,
opts: StreamOptions,
) {
let relay_rx = state.db.relay.broadcast_tx.subscribe();
let head = state.db.relay.seqs.head();
let start = match cursor {
// one from another relay or from before a reset. indigo goes live, which hides the gap,
// but the spec has us say so and close
Some(cursor) if cursor > head.unwrap_or(0) => {
let head = head.unwrap_or(0);
let _ = tx.blocking_send(Err(RelayStreamError::FutureCursor { cursor, head }));
return;
}
Some(cursor) => Start::replay(cursor.checked_sub(1), head),
None => Start::live(head),
};
run_stream(tx, relay_rx, start, opts, RelaySource { state });
}
struct RelaySource {
state: Arc<AppState>,
}
impl StreamSource for RelaySource {
type Broadcast = RelayBroadcast;
type Output = bytes::Bytes;
fn read(
&mut self,
after: Option<u64>,
through: Option<u64>,
limit: usize,
) -> Result<ReplayChunk<bytes::Bytes, u64>, ReplayFailed> {
read_relay_replay_chunk(&self.state.db.relay.events, after, through, limit)
}
fn written(&self) -> u64 {
self.state.db.relay.seqs.next()
}
fn frame_bytes(frame: &bytes::Bytes) -> usize {
frame.len()
}
fn render(&mut self, event: RelayBroadcast) -> Option<bytes::Bytes> {
match event {
RelayBroadcast::Persisted(_) => None,
RelayBroadcast::Ephemeral(_, frame) => Some(frame),
}
}
}
fn read_relay_replay_chunk(
ks: &fjall::Keyspace,
after: Option<u64>,
through: Option<u64>,
limit: usize,
) -> Result<ReplayChunk<bytes::Bytes, u64>, ReplayFailed> {
let start = after.map_or(0, |seq| seq.saturating_add(1));
let end = through.unwrap_or(u64::MAX);
if start > end {
return Ok(ReplayChunk {
events: Vec::new(),
last_seen: after,
exhausted: true,
scanned: 0,
});
}
let rows = ks.range(keys::relay_event_key(start)..=keys::relay_event_key(end));
ReplayChunk::read(
rows,
after,
limit,
|key| {
let seq = key.try_into().map(u64::from_be_bytes);
seq.inspect_err(|_| error!("relay stream: failed to parse relay event seq"))
.ok()
},
|_, frame| Some(bytes::Bytes::copy_from_slice(frame)),
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
#[test]
fn replay_reads_big_endian_relay_seqs() {
let tmp = tempfile::tempdir().unwrap();
let db = crate::db::Db::open(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})
.unwrap();
db.relay
.events
.insert([0, 0, 0, 0, 0, 0, 1, 2], b"frame")
.unwrap();
let chunk = read_relay_replay_chunk(&db.relay.events, Some(257), None, 8).unwrap();
assert_eq!(chunk.last_seen, Some(258));
assert_eq!(chunk.events, [bytes::Bytes::from_static(b"frame")]);
}
/// a consumer holding a cursor from another relay hears so, instead of waiting for nothing
/// until the seqs catch up
#[tokio::test]
async fn a_cursor_past_the_last_seq_is_refused() -> miette::Result<()> {
use futures::StreamExt;
use miette::IntoDiagnostic;
use std::time::Duration;
use tokio::time::timeout;
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = crate::control::Hydrant::new(Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})
.await?;
let head = hydrant.state.db.relay.seqs.head().unwrap_or(0);
let mut future = hydrant.subscribe_repos(Some(head + 1));
let first = timeout(Duration::from_secs(5), future.next())
.await
.into_diagnostic()?;
assert!(
matches!(first, Some(Err(RelayStreamError::FutureCursor { .. }))),
"{first:?}"
);
assert!(
timeout(Duration::from_secs(5), future.next())
.await
.into_diagnostic()?
.is_none()
);
// the last seq itself is still in the window, so that one waits for what comes next
let mut current = hydrant.subscribe_repos(Some(head));
assert!(
timeout(Duration::from_millis(200), current.next())
.await
.is_err()
);
Ok(())
}
}