use std::future::Future;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use futures::Stream;
#[cfg(feature = "testing")]
use futures::future::Either;
use ruststream::{
AckError, BatchSubscriber, BufferedSubscriber, HeaderMap, IncomingMessage, Positioned,
Seekable, Seeker, Subscriber,
};
use sea_streamer_file::{FileConsumer, FileErr, MessageSource, SeekTarget, is_end_of_stream};
use sea_streamer_types::{Consumer as _, ShardId, SharedMessage, StreamErr, StreamKey, Timestamp};
use tokio::runtime::Handle;
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::{mpsc, oneshot};
use crate::batching::BATCH_MAX_WAIT;
use crate::error::{SeaFileError, box_err};
#[cfg(feature = "testing")]
use crate::in_process::{FileQueue, LogSeeker};
use crate::message::{FilePosition, SeaMessage};
const CHANNEL_CAPACITY: usize = 64;
pub(crate) struct SeekCmd {
position: FilePosition,
done: oneshot::Sender<Result<(), SeaFileError>>,
}
pub(crate) struct Stamped {
epoch: u64,
item: Option<Result<SeaMessage, SeaFileError>>,
}
pub struct FileSubscriber {
stream: Arc<str>,
inner: BufferedSubscriber<Deliveries>,
}
impl std::fmt::Debug for FileSubscriber {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FileSubscriber")
.field("stream", &self.stream)
.finish_non_exhaustive()
}
}
impl FileSubscriber {
#[must_use]
pub fn stream_key(&self) -> &str {
&self.stream
}
pub(crate) fn spawn(
runtime: &Handle,
stream: String,
reader: impl Source,
closed: impl Fn() -> bool + Send + Sync + 'static,
) -> Self {
let (out_tx, out_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
let epoch = Arc::new(AtomicU64::new(0));
runtime.spawn(drive(
reader,
out_tx,
cmd_rx,
Arc::clone(&epoch),
stream.clone(),
closed,
));
let stream: Arc<str> = Arc::from(stream);
Self {
stream: Arc::clone(&stream),
inner: BufferedSubscriber::new(Deliveries::Driver(Driven {
stream,
rx: out_rx,
cmd: cmd_tx,
epoch,
}))
.max_wait(BATCH_MAX_WAIT),
}
}
#[cfg(feature = "testing")]
pub(crate) fn in_process(queue: FileQueue) -> Self {
Self {
stream: Arc::clone(queue.stream()),
inner: BufferedSubscriber::new(Deliveries::InProcess(queue)).max_wait(BATCH_MAX_WAIT),
}
}
}
impl Subscriber for FileSubscriber {
type Message = FileMessage;
type Error = SeaFileError;
fn stream(&mut self) -> impl Stream<Item = Result<FileMessage, SeaFileError>> + Send + '_ {
self.inner.stream()
}
}
impl BatchSubscriber for FileSubscriber {
type Batch = Vec<FileMessage>;
fn batches(
&mut self,
size: NonZeroUsize,
) -> impl Stream<Item = Result<Self::Batch, SeaFileError>> + Send + '_ {
self.inner.batches(size)
}
}
impl Seekable for FileSubscriber {
type Seeker = FileSeeker;
fn seeker(&self) -> FileSeeker {
self.inner.seeker()
}
}
enum Deliveries {
Driver(Driven),
#[cfg(feature = "testing")]
InProcess(FileQueue),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Deliveries>() == size_of::<Driven>());
impl Subscriber for Deliveries {
type Message = FileMessage;
type Error = SeaFileError;
#[cfg(not(feature = "testing"))]
fn stream(&mut self) -> impl Stream<Item = Result<FileMessage, SeaFileError>> + Send + '_ {
let Self::Driver(driven) = self;
driven.stream()
}
#[cfg(feature = "testing")]
fn stream(&mut self) -> impl Stream<Item = Result<FileMessage, SeaFileError>> + Send + '_ {
match self {
Self::Driver(driven) => Either::Left(driven.stream()),
Self::InProcess(queue) => Either::Right(in_process_stream(queue)),
}
}
}
impl Seekable for Deliveries {
type Seeker = FileSeeker;
fn seeker(&self) -> FileSeeker {
match self {
Self::Driver(driven) => driven.seeker(),
#[cfg(feature = "testing")]
Self::InProcess(queue) => {
FileSeeker::in_process(Arc::clone(queue.stream()), queue.seeker())
}
}
}
}
#[cfg(feature = "testing")]
fn in_process_stream(
queue: &mut FileQueue,
) -> impl Stream<Item = Result<FileMessage, SeaFileError>> + Send + '_ {
let seeker = Arc::new(FileSeeker::in_process(
Arc::clone(queue.stream()),
queue.seeker(),
));
futures::stream::poll_fn(move |cx| {
queue.poll_next(cx).map(|next| {
next.map(|message| {
Ok(FileMessage {
message,
seeker: Arc::clone(&seeker),
})
})
})
})
}
struct Driven {
stream: Arc<str>,
rx: mpsc::Receiver<Stamped>,
cmd: mpsc::UnboundedSender<SeekCmd>,
epoch: Arc<AtomicU64>,
}
impl Driven {
fn stream(&mut self) -> impl Stream<Item = Result<FileMessage, SeaFileError>> + Send + '_ {
let seeker = Arc::new(self.seeker());
futures::stream::poll_fn(move |cx| {
loop {
match self.rx.poll_recv(cx) {
std::task::Poll::Ready(Some(stamped)) => {
if stamped.epoch == self.epoch.load(Ordering::Acquire) {
return std::task::Poll::Ready(stamped.item.map(|item| {
item.map(|message| FileMessage {
message,
seeker: Arc::clone(&seeker),
})
}));
}
}
std::task::Poll::Ready(None) => return std::task::Poll::Ready(None),
std::task::Poll::Pending => return std::task::Poll::Pending,
}
}
})
}
fn seeker(&self) -> FileSeeker {
FileSeeker {
stream: Arc::clone(&self.stream),
backend: SeekBackend::Driver {
cmd: self.cmd.clone(),
},
}
}
}
pub struct FileMessage {
message: SeaMessage,
seeker: Arc<FileSeeker>,
}
impl std::fmt::Debug for FileMessage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FileMessage")
.field("message", &self.message)
.finish_non_exhaustive()
}
}
impl FileMessage {
#[must_use]
pub fn stream(&self) -> &str {
self.message.stream()
}
pub(crate) fn seeker(&self) -> &FileSeeker {
&self.seeker
}
}
impl Positioned for FileMessage {
type Position = FilePosition;
fn position(&self) -> FilePosition {
self.message.position()
}
}
impl IncomingMessage for FileMessage {
fn payload(&self) -> &[u8] {
self.message.payload()
}
fn headers(&self) -> &HeaderMap {
self.message.headers()
}
fn ack(self) -> impl Future<Output = Result<(), AckError>> {
self.message.ack()
}
fn nack(self, requeue: bool) -> impl Future<Output = Result<(), AckError>> {
self.message.nack(requeue)
}
}
#[derive(Clone)]
enum SeekBackend {
Driver { cmd: mpsc::UnboundedSender<SeekCmd> },
#[cfg(feature = "testing")]
Log(LogSeeker),
}
#[derive(Clone)]
pub struct FileSeeker {
stream: Arc<str>,
backend: SeekBackend,
}
impl FileSeeker {
#[cfg(feature = "testing")]
pub(crate) fn in_process(stream: Arc<str>, log: LogSeeker) -> Self {
Self {
stream,
backend: SeekBackend::Log(log),
}
}
#[must_use]
pub fn stream_key(&self) -> &str {
&self.stream
}
fn dead(&self, why: &'static str) -> SeaFileError {
SeaFileError::Seek {
stream: self.stream.to_string(),
source: Box::from(why),
}
}
}
impl std::fmt::Debug for FileSeeker {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FileSeeker")
.field("stream", &self.stream)
.finish_non_exhaustive()
}
}
impl Seeker for FileSeeker {
type Position = FilePosition;
type Error = SeaFileError;
async fn seek(&self, to: FilePosition) -> Result<(), SeaFileError> {
match &self.backend {
SeekBackend::Driver { cmd } => {
let (done, wait) = oneshot::channel();
cmd.send(SeekCmd { position: to, done })
.map_err(|_| self.dead("the subscription's driver task has shut down"))?;
wait.await
.map_err(|_| self.dead("the subscription's driver task has shut down"))?
}
#[cfg(feature = "testing")]
SeekBackend::Log(log) => log.request(to),
}
}
}
const SHARD: ShardId = ShardId::new(0);
pub(crate) enum Reader {
Tail(FileConsumer),
Replay {
source: Box<MessageSource>,
key: StreamKey,
},
}
pub(crate) trait Source: Send + 'static {
fn next(&mut self) -> impl Future<Output = Result<SharedMessage, StreamErr<FileErr>>> + Send;
fn reposition(
&mut self,
target: SeekTarget,
) -> impl Future<Output = Result<(), FileErr>> + Send;
fn repositions_after_end(&self) -> bool;
}
impl Source for Reader {
async fn next(&mut self) -> Result<SharedMessage, StreamErr<FileErr>> {
match self {
Self::Tail(consumer) => consumer.next().await,
Self::Replay { source, key } => loop {
let message = source.next().await.map_err(StreamErr::Backend)?.message;
if is_end_of_stream(&message) {
return Err(StreamErr::Backend(FileErr::StreamEnded));
}
if message.header().stream_key() == key {
return Ok(message.to_shared());
}
},
}
}
async fn reposition(&mut self, target: SeekTarget) -> Result<(), FileErr> {
match self {
Self::Tail(consumer) => consumer.seek_to(target).await,
Self::Replay { source, key } => source.seek(key, &SHARD, target).await,
}
}
fn repositions_after_end(&self) -> bool {
matches!(self, Self::Replay { .. })
}
}
fn is_clean_end(err: &StreamErr<FileErr>) -> bool {
matches!(
err,
StreamErr::Backend(FileErr::StreamEnded | FileErr::NotEnoughBytes)
)
}
enum After {
Read,
Idle,
Stop,
}
fn delivery(
next: Result<SharedMessage, StreamErr<FileErr>>,
reader: &impl Source,
stream: &str,
) -> (Option<Result<SeaMessage, SeaFileError>>, After) {
match next {
Ok(message) => (Some(Ok(SeaMessage::new(&message))), After::Read),
Err(err) if is_clean_end(&err) => {
let after = if reader.repositions_after_end() {
After::Idle
} else {
After::Stop
};
(None, after)
}
Err(err) => (
Some(Err(SeaFileError::Receive {
stream: stream.to_owned(),
source: box_err(err),
})),
After::Stop,
),
}
}
async fn drive(
mut reader: impl Source,
out: mpsc::Sender<Stamped>,
mut cmd_rx: mpsc::UnboundedReceiver<SeekCmd>,
epoch: Arc<AtomicU64>,
stream: String,
closed: impl Fn() -> bool + Send + Sync,
) {
let mut generation = 0;
let mut idle = false;
'read: loop {
let next = tokio::select! {
biased;
cmd = cmd_rx.recv() => {
let Some(cmd) = cmd else { break };
if serve(&mut reader, cmd, &mut generation, &epoch, &stream, &closed).await {
idle = false;
}
continue;
}
() = out.closed() => break,
next = reader.next(), if !idle => next,
};
let (item, after) = delivery(next, &reader, &stream);
let permit = match out.try_reserve() {
Ok(permit) => permit,
Err(TrySendError::Closed(())) => break,
Err(TrySendError::Full(())) => loop {
tokio::select! {
biased;
cmd = cmd_rx.recv() => {
let Some(cmd) = cmd else { break 'read };
if serve(&mut reader, cmd, &mut generation, &epoch, &stream, &closed).await {
idle = false;
continue 'read;
}
}
permit = out.reserve() => match permit {
Ok(permit) => break permit,
Err(_) => break 'read,
},
}
},
};
permit.send(Stamped {
epoch: generation,
item,
});
match after {
After::Read => {}
After::Idle => idle = true,
After::Stop => break,
}
}
}
async fn serve(
reader: &mut impl Source,
SeekCmd { position, done }: SeekCmd,
generation: &mut u64,
epoch: &AtomicU64,
stream: &str,
closed: impl Fn() -> bool + Send + Sync,
) -> bool {
if closed() {
let _ = done.send(Err(SeaFileError::Seek {
stream: stream.to_owned(),
source: Box::from("the broker has shut down"),
}));
return false;
}
let mut skip_one = false;
let target = match position {
FilePosition::Beginning | FilePosition::Sequence(0 | 1) => SeekTarget::Beginning,
FilePosition::End => SeekTarget::End,
FilePosition::Sequence(sequence) => {
skip_one = true;
SeekTarget::SeqNo(sequence - 1)
}
FilePosition::Timestamp(millis) => {
let nanos = i128::from(millis) * 1_000_000;
match Timestamp::from_unix_timestamp_nanos(nanos) {
Ok(timestamp) => SeekTarget::Timestamp(timestamp),
Err(err) => {
let _ = done.send(Err(SeaFileError::Invalid(format!(
"'{millis}' is not a valid timestamp: {err}"
))));
return false;
}
}
}
};
let mut result = reader.reposition(target).await.map_err(box_err);
if result.is_ok() && skip_one {
result = reader.next().await.map(drop).map_err(box_err);
}
let moved = result.is_ok();
if moved {
*generation += 1;
epoch.store(*generation, Ordering::Release);
}
let _ = done.send(result.map_err(|source| SeaFileError::Seek {
stream: stream.to_owned(),
source,
}));
moved
}
#[cfg(test)]
mod tests {
use std::future::ready;
use std::time::Duration;
use futures::StreamExt;
use sea_streamer_file::SeekErr;
use sea_streamer_types::{MessageHeader, SeqNo};
use super::*;
use crate::wire;
const STUCK: Duration = Duration::from_mins(1);
struct Log {
messages: Vec<SharedMessage>,
next: usize,
}
impl Log {
fn of(count: u8) -> Self {
let key = StreamKey::new("orders").expect("the key is valid");
let messages = (1..=count)
.map(|body| {
let bytes = wire::encode(&HeaderMap::new(), &[body], false);
let length = bytes.len();
SharedMessage::new(
MessageHeader::new(
key.clone(),
SHARD,
SeqNo::from(body),
Timestamp::now_utc(),
),
bytes,
0,
length,
)
})
.collect();
Self { messages, next: 0 }
}
}
impl Source for Log {
fn next(&mut self) -> impl Future<Output = Result<SharedMessage, StreamErr<FileErr>>> {
let message = self.messages.get(self.next).cloned();
if message.is_some() {
self.next += 1;
}
ready(message.ok_or(StreamErr::Backend(FileErr::NotEnoughBytes)))
}
fn reposition(&mut self, target: SeekTarget) -> impl Future<Output = Result<(), FileErr>> {
let next = match target {
SeekTarget::Beginning => 0,
SeekTarget::End => self.messages.len(),
SeekTarget::SeqNo(sequence) if sequence > self.messages.len() as u64 => {
return ready(Err(FileErr::SeekErr(SeekErr::OutOfBound)));
}
SeekTarget::SeqNo(sequence) => {
usize::try_from(sequence.saturating_sub(1)).expect("the sequence fits")
}
SeekTarget::Timestamp(_) => unreachable!("these tests seek by sequence"),
};
self.next = next;
ready(Ok(()))
}
fn repositions_after_end(&self) -> bool {
true
}
}
async fn next_body<S>(stream: &mut S) -> Option<u8>
where
S: Stream<Item = Result<FileMessage, SeaFileError>> + Unpin,
{
tokio::time::timeout(STUCK, stream.next())
.await
.expect("the subscription moves on")
.map(|next| next.expect("the delivery is ok").payload()[0])
}
#[tokio::test(start_paused = true)]
async fn a_seek_completes_while_the_delivery_channel_is_full() {
let backlog = u8::try_from(CHANNEL_CAPACITY * 2).expect("the backlog fits");
let mut subscriber = FileSubscriber::spawn(
&Handle::current(),
"orders".to_owned(),
Log::of(backlog),
|| false,
);
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
assert_eq!(next_body(&mut stream).await, Some(1));
tokio::task::yield_now().await;
tokio::time::timeout(STUCK, seeker.seek(FilePosition::sequence(100)))
.await
.expect("a seek completes while the delivery channel is full")
.expect("the seek succeeds");
for expected in 100..=backlog {
assert_eq!(next_body(&mut stream).await, Some(expected));
}
assert_eq!(next_body(&mut stream).await, None, "the replay ends");
}
#[tokio::test(start_paused = true)]
async fn a_seek_after_the_replay_reached_the_end_replays_again() {
let mut subscriber =
FileSubscriber::spawn(&Handle::current(), "orders".to_owned(), Log::of(3), || {
false
});
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
for expected in 1..=3 {
assert_eq!(next_body(&mut stream).await, Some(expected));
}
tokio::task::yield_now().await;
tokio::time::timeout(STUCK, seeker.seek(FilePosition::beginning()))
.await
.expect("the seek completes")
.expect("a seek after the end of the file succeeds");
for expected in 1..=3 {
assert_eq!(next_body(&mut stream).await, Some(expected));
}
assert_eq!(next_body(&mut stream).await, None, "the replay ends again");
}
#[tokio::test(start_paused = true)]
async fn a_seek_after_the_stream_ended_reopens_it() {
let mut subscriber =
FileSubscriber::spawn(&Handle::current(), "orders".to_owned(), Log::of(2), || {
false
});
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
assert_eq!(next_body(&mut stream).await, Some(1));
assert_eq!(next_body(&mut stream).await, Some(2));
assert_eq!(next_body(&mut stream).await, None);
tokio::time::timeout(STUCK, seeker.seek(FilePosition::sequence(2)))
.await
.expect("the seek completes")
.expect("a seek after the stream ended succeeds");
assert_eq!(next_body(&mut stream).await, Some(2));
assert_eq!(next_body(&mut stream).await, None);
}
#[tokio::test(start_paused = true)]
async fn a_seek_to_the_sequence_after_the_last_lands_at_the_tip() {
let mut subscriber =
FileSubscriber::spawn(&Handle::current(), "orders".to_owned(), Log::of(3), || {
false
});
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
assert_eq!(next_body(&mut stream).await, Some(1));
tokio::time::timeout(STUCK, seeker.seek(FilePosition::sequence(4)))
.await
.expect("the seek completes")
.expect("the sequence the file writes next is a position");
assert_eq!(next_body(&mut stream).await, None);
}
#[tokio::test(start_paused = true)]
async fn a_seek_to_the_first_sequence_replays_from_the_start() {
let mut subscriber =
FileSubscriber::spawn(&Handle::current(), "orders".to_owned(), Log::of(3), || {
false
});
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
assert_eq!(next_body(&mut stream).await, Some(1));
assert_eq!(next_body(&mut stream).await, Some(2));
tokio::time::timeout(STUCK, seeker.seek(FilePosition::sequence(1)))
.await
.expect("the seek completes")
.expect("the first sequence is a position");
for expected in 1..=3 {
assert_eq!(next_body(&mut stream).await, Some(expected));
}
assert_eq!(next_body(&mut stream).await, None);
}
#[tokio::test(start_paused = true)]
async fn a_refused_seek_keeps_what_was_queued() {
let mut subscriber =
FileSubscriber::spawn(&Handle::current(), "orders".to_owned(), Log::of(5), || {
false
});
let seeker = subscriber.seeker();
let mut stream = std::pin::pin!(subscriber.stream());
assert_eq!(next_body(&mut stream).await, Some(1));
tokio::task::yield_now().await;
tokio::time::timeout(STUCK, seeker.seek(FilePosition::sequence(9)))
.await
.expect("the seek completes")
.expect_err("no message carries that sequence");
for expected in 2..=5 {
assert_eq!(next_body(&mut stream).await, Some(expected));
}
assert_eq!(next_body(&mut stream).await, None);
}
}