use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use async_trait::async_trait;
use futures::Stream;
use nstreams_core::event::StreamEvent;
use nstreams_core::filter::StreamFilter;
use nstreams_core::namespace::{Namespace, NamespaceResources};
use nstreams_core::stream::ReadStreamBackend;
use nstreams_core::{stream_max_length_bytes, STREAM_MAX_AGE};
use rabbitmq_stream_client::error::StreamCreateError;
use rabbitmq_stream_client::types::{ByteCapacity, Message, OffsetSpecification, ResponseCode};
use rabbitmq_stream_client::{Environment, FilterConfiguration, NoDedup};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::{debug, info};
use crate::config::RabbitConfig;
const FILTER_VALUE_HEADER: &str = "x-nstreams-filter-value";
#[derive(Debug, Error)]
pub enum RabbitStreamError {
#[error("config: {0}")]
Config(#[from] crate::config::ConfigError),
#[error("stream: {0}")]
Stream(String),
}
impl From<RabbitStreamError> for nstreams_core::Error {
fn from(value: RabbitStreamError) -> Self {
Self::ReadStream(value.to_string())
}
}
fn stream_err(error: impl std::fmt::Display) -> nstreams_core::Error {
RabbitStreamError::Stream(error.to_string()).into()
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct WireEvent {
namespace: String,
version: u64,
payload: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
filter_value: Option<String>,
}
fn encode_stream_message(event: &StreamEvent) -> nstreams_core::Result<Message> {
let payload = serde_json::from_slice(&event.payload)
.map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;
let wire = WireEvent {
namespace: event.namespace.clone(),
version: event.version,
payload,
filter_value: event.filter_value.clone(),
};
let body =
serde_json::to_vec(&wire).map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;
let mut builder = Message::builder().body(body);
if let Some(filter_value) = &event.filter_value {
builder = builder
.application_properties()
.insert(FILTER_VALUE_HEADER, filter_value.as_str())
.message_builder();
}
Ok(builder.build())
}
fn decode_stream_message(namespace: &str, message: Message) -> nstreams_core::Result<StreamEvent> {
let body = message
.data()
.ok_or_else(|| stream_err("stream message missing body"))?;
let wire: WireEvent =
serde_json::from_slice(body).map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;
let filter_value = wire
.filter_value
.or_else(|| filter_value_from_message(&message));
Ok(StreamEvent {
namespace: if wire.namespace.is_empty() {
namespace.to_string()
} else {
wire.namespace
},
version: wire.version,
payload: serde_json::to_vec(&wire.payload)
.map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?,
filter_value,
})
}
fn filter_value_from_message(message: &Message) -> Option<String> {
message
.application_properties()?
.get(FILTER_VALUE_HEADER)
.and_then(|value| value.clone().try_into().ok())
}
fn filter_value_extractor(message: &Message) -> String {
filter_value_from_message(message).unwrap_or_default()
}
fn build_filter_configuration(namespace: &str, filter: &StreamFilter) -> FilterConfiguration {
let values = filter.values.clone();
let match_unfiltered = filter.match_unfiltered;
let predicate_filter = filter.clone();
let namespace = namespace.to_string();
FilterConfiguration::new(values, match_unfiltered).post_filter(move |message| {
decode_stream_message(&namespace, message.clone())
.map(|event| predicate_filter.matches_filter_value(event.filter_value.as_deref()))
.unwrap_or(false)
})
}
async fn build_producer(
environment: &Environment,
stream: &str,
) -> Result<rabbitmq_stream_client::Producer<NoDedup>, nstreams_core::Error> {
environment
.producer()
.filter_value_extractor(filter_value_extractor)
.build(stream)
.await
.map_err(stream_err)
}
#[derive(Clone)]
pub struct RabbitReadStream {
environment: Arc<Environment>,
}
impl RabbitReadStream {
pub async fn connect(amqp_url: &str) -> Result<Self, RabbitStreamError> {
let config = RabbitConfig::from_amqp_url(amqp_url)?;
let environment = Environment::builder()
.host(&config.host)
.port(config.stream_port)
.username(&config.user)
.password(&config.password)
.virtual_host(&config.vhost)
.build()
.await
.map_err(|error| RabbitStreamError::Stream(error.to_string()))?;
Ok(Self {
environment: Arc::new(environment),
})
}
pub fn from_environment(environment: Environment) -> Self {
Self {
environment: Arc::new(environment),
}
}
fn stream_name<N: Namespace>(&self, namespace: &N) -> String {
NamespaceResources::new(namespace).read_stream
}
}
#[async_trait]
impl ReadStreamBackend for RabbitReadStream {
async fn stream_exists<N: Namespace>(&self, namespace: &N) -> nstreams_core::Result<bool> {
let stream = self.stream_name(namespace);
match self
.environment
.consumer()
.offset(OffsetSpecification::First)
.build(&stream)
.await
{
Ok(consumer) => {
let _ = consumer.handle().close().await;
Ok(true)
}
Err(_) => Ok(false),
}
}
async fn create_read_stream<N: Namespace>(
&self,
namespace: &N,
max_event_bytes: u64,
) -> nstreams_core::Result<()> {
let stream = self.stream_name(namespace);
let max_length = stream_max_length_bytes(max_event_bytes);
let segment_size = max_event_bytes.saturating_mul(1_024).max(1024 * 1024);
let result = self
.environment
.stream_creator()
.max_length(ByteCapacity::B(max_length))
.max_age(RabbitConfig::stream_max_age())
.max_segment_size(ByteCapacity::B(segment_size))
.create(&stream)
.await;
if let Err(error) = result {
if let StreamCreateError::Create { status, .. } = error {
if status != ResponseCode::StreamAlreadyExists {
return Err(stream_err(format!(
"failed to create stream {stream}: {status:?}"
)));
}
} else {
return Err(stream_err(error));
}
}
info!(
namespace = namespace.as_str(),
stream = %stream,
max_length_bytes = max_length,
max_age = STREAM_MAX_AGE,
"created read stream"
);
Ok(())
}
async fn publish_to_stream(&self, event: &StreamEvent) -> nstreams_core::Result<()> {
let stream = nstreams_core::namespace::read_stream_name(&event.namespace);
let producer = build_producer(&self.environment, &stream).await?;
producer
.send_with_confirm(encode_stream_message(event)?)
.await
.map_err(stream_err)?;
producer.close().await.map_err(stream_err)?;
Ok(())
}
async fn populate_stream(&self, events: &[StreamEvent]) -> nstreams_core::Result<()> {
if events.is_empty() {
return Ok(());
}
let stream = nstreams_core::namespace::read_stream_name(&events[0].namespace);
let producer = build_producer(&self.environment, &stream).await?;
for event in events {
producer
.send_with_confirm(encode_stream_message(event)?)
.await
.map_err(stream_err)?;
}
producer.close().await.map_err(stream_err)?;
debug!(stream = %stream, count = events.len(), "populated read stream");
Ok(())
}
async fn subscribe_live<N: Namespace>(
&self,
namespace: &N,
filter: Option<&StreamFilter>,
) -> nstreams_core::Result<impl Stream<Item = nstreams_core::Result<StreamEvent>> + Send> {
let stream = self.stream_name(namespace);
let namespace = namespace.as_str().to_string();
let mut consumer_builder = self
.environment
.consumer()
.offset(OffsetSpecification::Last);
if let Some(filter) = filter.filter(|filter| filter.is_active()) {
consumer_builder =
consumer_builder.filter_input(Some(build_filter_configuration(namespace.as_str(), filter)));
}
let consumer = consumer_builder.build(&stream).await.map_err(stream_err)?;
Ok(LiveStream { namespace, consumer })
}
}
struct LiveStream {
namespace: String,
consumer: rabbitmq_stream_client::Consumer,
}
impl Stream for LiveStream {
type Item = nstreams_core::Result<StreamEvent>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.as_mut().get_mut();
match Pin::new(&mut this.consumer).poll_next(cx) {
Poll::Ready(Some(Ok(delivery))) => {
let message = delivery.message().clone();
Poll::Ready(Some(decode_stream_message(&this.namespace, message)))
}
Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(stream_err(error)))),
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}