use std::collections::HashMap;
use std::sync::Arc;
use futures::StreamExt;
use tokio::pin;
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, info, warn};
use crate::event::{DeliveredEvent, EventVersion, StreamEvent};
use crate::filter::{StreamFilter, SubscribeRequest};
use crate::namespace::{Namespace, StreamNamespace};
use crate::request_id::new_request_id;
use crate::store::EventStore;
use crate::stream::ReadStreamBackend;
use crate::subscription::Subscription;
use crate::version::VersionedDelivery;
use crate::{DEFAULT_MAX_EVENT_BYTES, MAX_STREAM_HISTORY_ITEMS};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ReadStreamLifecycle {
Populating,
Ready,
}
pub struct NStreamsHandler<S, R> {
store: Arc<S>,
read_stream: Arc<R>,
max_event_bytes: u64,
lifecycle: Arc<RwLock<HashMap<String, ReadStreamLifecycle>>>,
population_locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
}
impl<S, R> Clone for NStreamsHandler<S, R> {
fn clone(&self) -> Self {
Self {
store: Arc::clone(&self.store),
read_stream: Arc::clone(&self.read_stream),
max_event_bytes: self.max_event_bytes,
lifecycle: Arc::clone(&self.lifecycle),
population_locks: Arc::clone(&self.population_locks),
}
}
}
impl<S, R> NStreamsHandler<S, R>
where
S: EventStore + 'static,
R: ReadStreamBackend + 'static,
{
pub fn new(store: S, read_stream: R) -> Self {
Self::with_max_event_bytes(store, read_stream, DEFAULT_MAX_EVENT_BYTES)
}
pub fn with_max_event_bytes(store: S, read_stream: R, max_event_bytes: u64) -> Self {
Self {
store: Arc::new(store),
read_stream: Arc::new(read_stream),
max_event_bytes,
lifecycle: Arc::new(RwLock::new(HashMap::new())),
population_locks: Arc::new(Mutex::new(HashMap::new())),
}
}
pub async fn subscribe<N: Namespace>(
&self,
namespace: &N,
history_count: u64,
) -> crate::Result<Subscription> {
self.subscribe_with_filter(
namespace,
SubscribeRequest {
namespace: namespace.as_str().to_string(),
history: history_count,
filter: None,
after_version: None,
},
)
.await
}
pub async fn subscribe_with_request(
&self,
request: SubscribeRequest,
) -> crate::Result<Subscription> {
let namespace = StreamNamespace::from(request.namespace.as_str());
self.subscribe_with_filter(&namespace, request).await
}
async fn subscribe_with_filter<N: Namespace>(
&self,
namespace: &N,
request: SubscribeRequest,
) -> crate::Result<Subscription> {
if let Some(filter) = request.filter.as_ref() {
validate_filter(filter)?;
}
StreamNamespace::validate(namespace.as_str())?;
self.store.ensure_namespace(namespace).await?;
self.ensure_read_stream(namespace).await?;
let history_limit = request.history.min(MAX_STREAM_HISTORY_ITEMS);
let mut history = if let Some(after_version) = request.after_version {
self.store
.load_after_version(namespace, after_version, history_limit)
.await?
} else {
self.store.load_history(namespace, history_limit).await?
};
if let Some(filter) = request.filter.as_ref() {
history.retain(|event| filter.matches_filter_value(event.filter_value.as_deref()));
}
let request_id = new_request_id();
let (tx, rx) = tokio::sync::mpsc::channel(1024);
let handler = self.clone();
let namespace_owned = StreamNamespace::from(namespace.as_str());
let request_id_for_task = request_id.clone();
let filter = request.filter;
tokio::spawn(async move {
if let Err(error) = handler
.run_subscription(request_id_for_task, namespace_owned, history, filter, tx)
.await
{
warn!(%error, "subscription task failed");
}
});
Ok(Subscription::new(request_id, rx))
}
async fn ensure_read_stream<N: Namespace>(&self, namespace: &N) -> crate::Result<()> {
let key = namespace.as_str().to_string();
{
let lifecycle = self.lifecycle.read().await;
if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
return Ok(());
}
}
let lock = {
let mut locks = self.population_locks.lock().await;
locks
.entry(key.clone())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
};
let _guard = lock.lock().await;
{
let lifecycle = self.lifecycle.read().await;
if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
return Ok(());
}
}
if self.read_stream.stream_exists(namespace).await? {
let mut lifecycle = self.lifecycle.write().await;
lifecycle.insert(key, ReadStreamLifecycle::Ready);
return Ok(());
}
{
let mut lifecycle = self.lifecycle.write().await;
lifecycle.insert(key.clone(), ReadStreamLifecycle::Populating);
}
info!(namespace = namespace.as_str(), "creating read stream");
self.read_stream
.create_read_stream(namespace, self.max_event_bytes)
.await?;
let bootstrap_limit = MAX_STREAM_HISTORY_ITEMS;
let history = self.store.load_history(namespace, bootstrap_limit).await?;
if !history.is_empty() {
debug!(
namespace = namespace.as_str(),
count = history.len(),
"bootstrapping read stream from postgres"
);
self.read_stream.populate_stream(&history).await?;
}
let mut lifecycle = self.lifecycle.write().await;
lifecycle.insert(key, ReadStreamLifecycle::Ready);
info!(namespace = namespace.as_str(), "read stream ready");
Ok(())
}
pub async fn is_read_stream_ready<N: Namespace>(&self, namespace: &N) -> bool {
let lifecycle = self.lifecycle.read().await;
lifecycle.get(namespace.as_str()) == Some(&ReadStreamLifecycle::Ready)
}
async fn run_subscription(
&self,
request_id: String,
namespace: StreamNamespace,
history: Vec<StreamEvent>,
filter: Option<StreamFilter>,
tx: tokio::sync::mpsc::Sender<crate::Result<DeliveredEvent>>,
) -> crate::Result<()> {
let mut last_version: Option<EventVersion> = None;
for event in history {
last_version = Some(VersionedDelivery::verify_initial(last_version, &event)?);
let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
if tx.send(Ok(delivered)).await.is_err() {
return Ok(());
}
}
let live = self
.read_stream
.subscribe_live(&namespace, filter.as_ref())
.await?;
pin!(live);
while let Some(item) = live.next().await {
match item {
Ok(event) => {
if let Some(filter) = filter.as_ref() {
if !filter.matches_filter_value(event.filter_value.as_deref()) {
continue;
}
}
if let Some(last) = last_version {
if event.version <= last {
continue;
}
VersionedDelivery { version: last }.verify_next(&event)?;
}
last_version = Some(event.version);
let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
if tx.send(Ok(delivered)).await.is_err() {
break;
}
}
Err(error) => {
if tx.send(Err(error)).await.is_err() {
break;
}
}
}
}
Ok(())
}
}
fn validate_filter(filter: &StreamFilter) -> crate::Result<()> {
if filter.is_active() {
Ok(())
} else {
Err(crate::Error::InvalidFilter(
"filter must include at least one value or match_unfiltered=true".into(),
))
}
}