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::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;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::{debug, info};
use crate::config::RabbitConfig;
#[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,
}
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,
};
let body =
serde_json::to_vec(&wire).map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;
Ok(Message::builder().body(body).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()))?;
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()))?,
})
}
#[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 = self
.environment
.producer()
.build(&stream)
.await
.map_err(stream_err)?;
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 = self
.environment
.producer()
.build(&stream)
.await
.map_err(stream_err)?;
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,
) -> 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 consumer = self
.environment
.consumer()
.offset(OffsetSpecification::Last)
.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,
}
}
}