use std::{
fmt,
future::{Future, ready},
marker::PhantomData,
num::{NonZeroU64, NonZeroUsize},
sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
},
task::{Poll, ready},
time::Duration,
};
use futures::{Stream, task::AtomicWaker};
use thiserror::Error;
use tokio::{sync::mpsc, time::timeout};
use tracing::warn;
use super::{
Bus, Discarding, LogMode, MemoryDelivery, MemoryError, MemoryMessage, MemoryOutbound,
MemoryPublisher, MemoryState, MemorySubscriber, Retaining, SubscriptionShared, log,
};
#[cfg(feature = "testing")]
use crate::testing::coordinator::Coordinator;
use crate::{
BatchSubscriber, BytesMut, IncomingMessage, OutgoingMessage, OwnedTransactions, Partitioned,
Positioned, Publisher, RequestReply, Seekable, Seeker, Str, Subscriber, Take, Transaction,
TransactionalPublisher,
};
pub const PARTITION_KEY_HEADER: &str = "partition-key";
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum RequestError {
#[error("no reply to \"{subject}\" within {timeout:?}")]
Timeout {
subject: String,
timeout: Duration,
},
#[error("the memory broker is shut down")]
ShutDown,
}
#[derive(Clone)]
pub struct MemoryRequester {
state: Arc<MemoryState>,
}
impl MemoryRequester {
pub(super) fn new(state: Arc<MemoryState>) -> Self {
Self { state }
}
}
impl fmt::Debug for MemoryRequester {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MemoryRequester").finish_non_exhaustive()
}
}
impl Publisher for MemoryRequester {
type Payload = Take;
type Error = RequestError;
type Options = ();
fn publish(
&self,
msg: OutgoingMessage<'_, BytesMut>,
_options: Option<&Self::Options>,
) -> impl Future<Output = Result<(), Self::Error>> {
let (name, payload, headers) = msg.into_parts();
ready(
self.state
.fanout(name, payload.freeze(), headers)
.map_err(|_| RequestError::ShutDown),
)
}
}
impl RequestReply for MemoryRequester {
type Reply = MemoryMessage<Discarding>;
async fn request(
&self,
msg: OutgoingMessage<'_, BytesMut>,
wait: Duration,
) -> Result<Self::Reply, Self::Error> {
let id = self.state.inbox_seq.fetch_add(1, Ordering::Relaxed);
let inbox = format!("_inbox.{id}");
let (tx, mut rx) = mpsc::unbounded_channel();
self.state
.register(&inbox, tx.clone())
.map_err(|_| RequestError::ShutDown)?;
let (subject, payload, mut headers) = msg.into_parts();
headers.insert(Str::from_static("reply-to"), inbox.clone());
if self
.state
.fanout(subject, payload.freeze(), headers)
.is_err()
{
self.state.unregister(&inbox);
return Err(RequestError::ShutDown);
}
let outcome = timeout(wait, rx.recv()).await;
self.state.unregister(&inbox);
match outcome {
Ok(Some(reply)) => Ok(MemoryMessage {
delivery: Some(reply),
subscription: None,
#[cfg(feature = "testing")]
coordinator: None,
mode: PhantomData,
}),
Ok(None) => unreachable!("request inbox closed while its sender is held"),
Err(_) => Err(RequestError::Timeout {
subject: subject.to_owned(),
timeout: wait,
}),
}
}
}
impl<Log: LogMode> BatchSubscriber for MemorySubscriber<Log> {
type Batch = Vec<MemoryMessage<Log>>;
fn batches(
&mut self,
size: NonZeroUsize,
) -> impl Stream<Item = Result<Self::Batch, <Self as Subscriber>::Error>> + Send + '_ {
let limit = size.get();
#[cfg(feature = "testing")]
let coordinator = self.coordinator();
futures::stream::poll_fn(move |cx| {
if log::retains::<Log>() {
self.subscription.seek.waker.register(cx.waker());
}
self.apply_pending_seek();
let first = loop {
match ready!(self.rx.poll_recv(cx)) {
Some(delivery)
if log::retains::<Log>()
&& delivery.seq < self.subscription.seek.watermark() =>
{
#[cfg(feature = "testing")]
if let Some(coordinator) = &coordinator {
coordinator.consumed();
}
}
Some(delivery) => break delivery,
None => return Poll::Ready(None),
}
};
let mut batch = Vec::with_capacity(limit);
batch.push(MemoryMessage {
delivery: Some(first),
subscription: Some(Arc::clone(&self.subscription)),
#[cfg(feature = "testing")]
coordinator: coordinator.clone(),
mode: PhantomData,
});
while batch.len() < limit {
match self.rx.poll_recv(cx) {
Poll::Ready(Some(delivery))
if log::retains::<Log>()
&& delivery.seq < self.subscription.seek.watermark() =>
{
#[cfg(feature = "testing")]
if let Some(coordinator) = &coordinator {
coordinator.consumed();
}
}
Poll::Ready(Some(delivery)) => batch.push(MemoryMessage {
delivery: Some(delivery),
subscription: Some(Arc::clone(&self.subscription)),
#[cfg(feature = "testing")]
coordinator: coordinator.clone(),
mode: PhantomData,
}),
Poll::Ready(None) | Poll::Pending => break,
}
}
Poll::Ready(Some(Ok(batch)))
})
}
}
impl TransactionalPublisher for MemoryPublisher {
fn begin_transaction(&self) -> impl Future<Output = Result<(), MemoryError>> {
let mut txn = self.txn.lock().expect("memory broker mutex poisoned");
if txn.is_some() {
return ready(Err(MemoryError::TransactionBusy));
}
*txn = Some(Vec::new());
drop(txn);
ready(Ok(()))
}
fn commit(&self) -> impl Future<Output = Result<(), MemoryError>> {
let buffered = self
.txn
.lock()
.expect("memory broker mutex poisoned")
.take();
let Some(buffered) = buffered else {
return ready(Err(MemoryError::NoTransaction));
};
for outbound in buffered {
if let Err(err) = self
.state
.fanout(&outbound.name, outbound.payload, outbound.headers)
{
return ready(Err(err));
}
}
ready(Ok(()))
}
fn abort(&self) -> impl Future<Output = Result<(), MemoryError>> {
ready(
self.txn
.lock()
.expect("memory broker mutex poisoned")
.take()
.map(|_| ())
.ok_or(MemoryError::NoTransaction),
)
}
}
#[must_use = "a transaction does nothing until settled with commit() or abort()"]
pub struct MemoryTransaction {
state: Arc<MemoryState>,
buffered: Vec<MemoryOutbound>,
settled: bool,
}
impl fmt::Debug for MemoryTransaction {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MemoryTransaction")
.field("buffered", &self.buffered.len())
.field("settled", &self.settled)
.finish_non_exhaustive()
}
}
impl Drop for MemoryTransaction {
fn drop(&mut self) {
if !self.settled {
warn!(
target: "ruststream::memory",
buffered = self.buffered.len(),
"owned transaction dropped without commit or abort; its buffered messages are \
discarded"
);
}
}
}
impl Transaction for MemoryTransaction {
type Payload = Take;
type Error = MemoryError;
type Options = ();
fn publish(
&mut self,
msg: OutgoingMessage<'_, BytesMut>,
_options: Option<&Self::Options>,
) -> impl Future<Output = Result<(), MemoryError>> {
self.buffered.push(MemoryOutbound::from(msg));
ready(Ok(()))
}
fn commit(mut self) -> impl Future<Output = Result<(), MemoryError>> {
self.settled = true;
for outbound in std::mem::take(&mut self.buffered) {
if let Err(err) = self
.state
.fanout(&outbound.name, outbound.payload, outbound.headers)
{
return ready(Err(err));
}
}
ready(Ok(()))
}
fn abort(mut self) -> impl Future<Output = Result<(), MemoryError>> {
self.settled = true;
ready(Ok(()))
}
}
impl OwnedTransactions for MemoryPublisher {
type Transaction = MemoryTransaction;
fn transaction(&self) -> impl Future<Output = Result<MemoryTransaction, MemoryError>> {
ready(Ok(MemoryTransaction {
state: Arc::clone(&self.state),
buffered: Vec::new(),
settled: false,
}))
}
}
impl<Log: LogMode> Partitioned for MemoryMessage<Log> {
fn partition_key(&self) -> Option<&[u8]> {
self.headers().get(PARTITION_KEY_HEADER)
}
}
#[derive(Default)]
pub(super) struct SeekControl {
pending: Mutex<Option<usize>>,
watermark: AtomicUsize,
pub(super) waker: AtomicWaker,
}
impl SeekControl {
pub(super) fn watermark(&self) -> usize {
self.watermark.load(Ordering::Acquire)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[must_use]
pub enum MemoryPosition {
Start,
Sequence(usize),
End,
}
impl MemoryPosition {
pub const fn start() -> Self {
Self::Start
}
pub const fn sequence(n: usize) -> Self {
Self::Sequence(n)
}
pub const fn end() -> Self {
Self::End
}
}
#[derive(Clone)]
pub struct MemorySeeker {
subscription: Arc<SubscriptionShared>,
}
impl MemorySeeker {
fn of(msg: &MemoryMessage<Retaining>) -> Self {
Self {
subscription: Arc::clone(
msg.subscription
.as_ref()
.expect("a seek context builds only off subscription deliveries"),
),
}
}
}
impl fmt::Debug for MemorySeeker {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MemorySeeker")
.field("name", &self.subscription.name)
.finish_non_exhaustive()
}
}
impl Seekable for MemorySubscriber<Retaining> {
type Seeker = MemorySeeker;
fn seeker(&self) -> MemorySeeker {
MemorySeeker {
subscription: Arc::clone(&self.subscription),
}
}
}
impl Seeker for MemorySeeker {
type Position = MemoryPosition;
type Error = MemoryError;
fn seek(&self, to: MemoryPosition) -> impl Future<Output = Result<(), MemoryError>> {
let bus = self
.subscription
.state
.subscribers
.lock()
.expect("memory broker mutex poisoned");
if !matches!(&*bus, Bus::Live(_)) {
return ready(Err(MemoryError::ShutDown));
}
let log = self
.subscription
.state
.log
.lock()
.expect("memory broker mutex poisoned");
let (oldest, tip) = log
.name(&self.subscription.name)
.map_or((0, 0), |name| (name.first_seq(), name.next_seq()));
drop(log);
let target = match to {
MemoryPosition::Start => oldest,
MemoryPosition::End => tip,
MemoryPosition::Sequence(n) if n >= oldest => n.min(tip),
MemoryPosition::Sequence(n) => {
return ready(Err(MemoryError::PositionEvicted {
requested: n,
oldest,
}));
}
};
self.subscription
.seek
.watermark
.store(target, Ordering::Release);
let replaced = self
.subscription
.seek
.pending
.lock()
.expect("memory broker mutex poisoned")
.replace(target);
drop(bus);
#[cfg(feature = "testing")]
if replaced.is_none()
&& let Some(coordinator) = self.subscription.state.coordinator()
{
coordinator.enqueued();
}
#[cfg(not(feature = "testing"))]
let _ = replaced;
self.subscription.seek.waker.wake();
ready(Ok(()))
}
}
impl<Log: LogMode> MemorySubscriber<Log> {
pub(super) fn apply_pending_seek(&mut self) {
if !log::retains::<Log>() {
return;
}
let pending = self
.subscription
.seek
.pending
.lock()
.expect("memory broker mutex poisoned")
.take();
let Some(target) = pending else { return };
#[cfg(feature = "testing")]
let coordinator = self.coordinator();
#[cfg(feature = "testing")]
let seek_token = SeekToken(coordinator.clone());
let bus = self
.subscription
.state
.subscribers
.lock()
.expect("memory broker mutex poisoned");
if !matches!(&*bus, Bus::Live(_)) {
return;
}
let log = self
.subscription
.state
.log
.lock()
.expect("memory broker mutex poisoned");
let retained = log.name(self.subscription.name.as_str());
#[cfg(feature = "testing")]
if let Some(coordinator) = &coordinator {
for _ in retained.into_iter().flat_map(|log| log.replay_from(target)) {
coordinator.enqueued();
}
}
while self.rx.try_recv().is_ok() {
#[cfg(feature = "testing")]
if let Some(coordinator) = &coordinator {
coordinator.consumed();
}
}
if let Some(retained) = retained {
for (seq, entry) in retained.replay_from(target) {
let delivery = MemoryDelivery {
shared: entry.shared(),
seq,
deliveries: NonZeroU64::MIN,
};
let _ = self.subscription.requeue.send(delivery);
}
}
drop(log);
drop(bus);
#[cfg(feature = "testing")]
drop(seek_token);
}
}
#[cfg(feature = "testing")]
struct SeekToken(Option<Coordinator>);
#[cfg(feature = "testing")]
impl Drop for SeekToken {
fn drop(&mut self) {
if let Some(coordinator) = &self.0 {
coordinator.consumed();
}
}
}
impl Positioned for MemoryMessage<Retaining> {
type Position = MemoryPosition;
fn position(&self) -> MemoryPosition {
MemoryPosition::Sequence(self.delivery.as_ref().map_or(0, |d| d.seq))
}
}
#[derive(Debug)]
pub struct MemoryContext {
position: MemoryPosition,
seeker: MemorySeeker,
}
impl crate::BuildContext<MemoryMessage<Retaining>> for MemoryContext {
fn build(msg: &MemoryMessage<Retaining>) -> Self {
Self {
position: Positioned::position(msg),
seeker: MemorySeeker::of(msg),
}
}
}
#[derive(Debug)]
pub struct MemoryBatchContext {
seeker: MemorySeeker,
}
impl crate::BuildBatchContext<MemoryMessage<Retaining>> for MemoryBatchContext {
fn build(first: &MemoryMessage<Retaining>) -> Self {
Self {
seeker: MemorySeeker::of(first),
}
}
}
impl crate::Field<MemoryBatchContext> for SeekHandle {
type Value<'a> = &'a MemorySeeker;
fn get(self, src: &MemoryBatchContext) -> &MemorySeeker {
&src.seeker
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct Position;
impl crate::ContextField for Position {
type Context = MemoryContext;
type Value = MemoryPosition;
fn read(self, src: &MemoryContext) -> MemoryPosition {
src.position
}
}
impl crate::Field<MemoryContext> for Position {
type Value<'a> = MemoryPosition;
fn get(self, src: &MemoryContext) -> MemoryPosition {
src.position
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct SeekHandle;
impl crate::ContextField for SeekHandle {
type Context = MemoryContext;
type Value = MemorySeeker;
fn read(self, src: &MemoryContext) -> MemorySeeker {
src.seeker.clone()
}
}
impl crate::Field<MemoryContext> for SeekHandle {
type Value<'a> = &'a MemorySeeker;
fn get(self, src: &MemoryContext) -> &MemorySeeker {
&src.seeker
}
}
#[cfg(test)]
mod tests;