Skip to main content

nstreams_core/
handler.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use futures::StreamExt;
5use tokio::pin;
6use tokio::sync::{Mutex, RwLock};
7use tracing::{debug, info, warn};
8
9use crate::event::{DeliveredEvent, EventVersion, StreamEvent};
10use crate::filter::{StreamFilter, SubscribeRequest};
11use crate::namespace::{Namespace, StreamNamespace};
12use crate::request_id::new_request_id;
13use crate::store::EventStore;
14use crate::stream::ReadStreamBackend;
15use crate::subscription::Subscription;
16use crate::version::VersionedDelivery;
17use crate::{DEFAULT_MAX_EVENT_BYTES, MAX_STREAM_HISTORY_ITEMS};
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20enum ReadStreamLifecycle {
21    Populating,
22    Ready,
23}
24
25/// Coordinates read-side subscriptions with history replay and live tailing.
26pub struct NStreamsHandler<S, R> {
27    store: Arc<S>,
28    read_stream: Arc<R>,
29    max_event_bytes: u64,
30    lifecycle: Arc<RwLock<HashMap<String, ReadStreamLifecycle>>>,
31    population_locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
32}
33
34impl<S, R> Clone for NStreamsHandler<S, R> {
35    fn clone(&self) -> Self {
36        Self {
37            store: Arc::clone(&self.store),
38            read_stream: Arc::clone(&self.read_stream),
39            max_event_bytes: self.max_event_bytes,
40            lifecycle: Arc::clone(&self.lifecycle),
41            population_locks: Arc::clone(&self.population_locks),
42        }
43    }
44}
45
46impl<S, R> NStreamsHandler<S, R>
47where
48    S: EventStore + 'static,
49    R: ReadStreamBackend + 'static,
50{
51    pub fn new(store: S, read_stream: R) -> Self {
52        Self::with_max_event_bytes(store, read_stream, DEFAULT_MAX_EVENT_BYTES)
53    }
54
55    pub fn with_max_event_bytes(store: S, read_stream: R, max_event_bytes: u64) -> Self {
56        Self {
57            store: Arc::new(store),
58            read_stream: Arc::new(read_stream),
59            max_event_bytes,
60            lifecycle: Arc::new(RwLock::new(HashMap::new())),
61            population_locks: Arc::new(Mutex::new(HashMap::new())),
62        }
63    }
64
65    /// Subscribe to a namespace, replaying history then tailing the live stream.
66    pub async fn subscribe<N: Namespace>(
67        &self,
68        namespace: &N,
69        history_count: u64,
70    ) -> crate::Result<Subscription> {
71        self.subscribe_with_filter(
72            namespace,
73            SubscribeRequest {
74                namespace: namespace.as_str().to_string(),
75                history: history_count,
76                filter: None,
77            },
78        )
79        .await
80    }
81
82    /// Subscribe with an explicit request, including optional RabbitMQ stream filters.
83    pub async fn subscribe_with_request(
84        &self,
85        request: SubscribeRequest,
86    ) -> crate::Result<Subscription> {
87        let namespace = StreamNamespace::from(request.namespace.as_str());
88        self.subscribe_with_filter(&namespace, request).await
89    }
90
91    async fn subscribe_with_filter<N: Namespace>(
92        &self,
93        namespace: &N,
94        request: SubscribeRequest,
95    ) -> crate::Result<Subscription> {
96        if let Some(filter) = request.filter.as_ref() {
97            validate_filter(filter)?;
98        }
99
100        StreamNamespace::validate(namespace.as_str())?;
101        self.store.ensure_namespace(namespace).await?;
102        self.ensure_read_stream(namespace).await?;
103
104        let history_limit = request.history.min(MAX_STREAM_HISTORY_ITEMS);
105        let mut history = self.store.load_history(namespace, history_limit).await?;
106        if let Some(filter) = request.filter.as_ref() {
107            history.retain(|event| filter.matches_filter_value(event.filter_value.as_deref()));
108        }
109
110        let request_id = new_request_id();
111        let (tx, rx) = tokio::sync::mpsc::channel(1024);
112
113        let handler = self.clone();
114        let namespace_owned = StreamNamespace::from(namespace.as_str());
115        let request_id_for_task = request_id.clone();
116        let filter = request.filter;
117        tokio::spawn(async move {
118            if let Err(error) = handler
119                .run_subscription(request_id_for_task, namespace_owned, history, filter, tx)
120                .await
121            {
122                warn!(%error, "subscription task failed");
123            }
124        });
125
126        Ok(Subscription::new(request_id, rx))
127    }
128
129    async fn ensure_read_stream<N: Namespace>(&self, namespace: &N) -> crate::Result<()> {
130        let key = namespace.as_str().to_string();
131
132        {
133            let lifecycle = self.lifecycle.read().await;
134            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
135                return Ok(());
136            }
137        }
138
139        let lock = {
140            let mut locks = self.population_locks.lock().await;
141            locks
142                .entry(key.clone())
143                .or_insert_with(|| Arc::new(Mutex::new(())))
144                .clone()
145        };
146
147        let _guard = lock.lock().await;
148
149        {
150            let lifecycle = self.lifecycle.read().await;
151            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
152                return Ok(());
153            }
154        }
155
156        if self.read_stream.stream_exists(namespace).await? {
157            let mut lifecycle = self.lifecycle.write().await;
158            lifecycle.insert(key, ReadStreamLifecycle::Ready);
159            return Ok(());
160        }
161
162        {
163            let mut lifecycle = self.lifecycle.write().await;
164            lifecycle.insert(key.clone(), ReadStreamLifecycle::Populating);
165        }
166
167        info!(namespace = namespace.as_str(), "creating read stream");
168        self.read_stream
169            .create_read_stream(namespace, self.max_event_bytes)
170            .await?;
171
172        let bootstrap_limit = MAX_STREAM_HISTORY_ITEMS;
173        let history = self.store.load_history(namespace, bootstrap_limit).await?;
174        if !history.is_empty() {
175            debug!(
176                namespace = namespace.as_str(),
177                count = history.len(),
178                "bootstrapping read stream from postgres"
179            );
180            self.read_stream.populate_stream(&history).await?;
181        }
182
183        let mut lifecycle = self.lifecycle.write().await;
184        lifecycle.insert(key, ReadStreamLifecycle::Ready);
185        info!(namespace = namespace.as_str(), "read stream ready");
186        Ok(())
187    }
188
189    pub async fn is_read_stream_ready<N: Namespace>(&self, namespace: &N) -> bool {
190        let lifecycle = self.lifecycle.read().await;
191        lifecycle.get(namespace.as_str()) == Some(&ReadStreamLifecycle::Ready)
192    }
193
194    async fn run_subscription(
195        &self,
196        request_id: String,
197        namespace: StreamNamespace,
198        history: Vec<StreamEvent>,
199        filter: Option<StreamFilter>,
200        tx: tokio::sync::mpsc::Sender<crate::Result<DeliveredEvent>>,
201    ) -> crate::Result<()> {
202        let mut last_version: Option<EventVersion> = None;
203
204        for event in history {
205            last_version = Some(VersionedDelivery::verify_initial(last_version, &event)?);
206            let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
207            if tx.send(Ok(delivered)).await.is_err() {
208                return Ok(());
209            }
210        }
211
212        let live = self
213            .read_stream
214            .subscribe_live(&namespace, filter.as_ref())
215            .await?;
216        pin!(live);
217        while let Some(item) = live.next().await {
218            match item {
219                Ok(event) => {
220                    if let Some(filter) = filter.as_ref() {
221                        if !filter.matches_filter_value(event.filter_value.as_deref()) {
222                            continue;
223                        }
224                    }
225
226                    if let Some(last) = last_version {
227                        if event.version <= last {
228                            continue;
229                        }
230                        VersionedDelivery { version: last }.verify_next(&event)?;
231                    }
232                    last_version = Some(event.version);
233                    let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
234                    if tx.send(Ok(delivered)).await.is_err() {
235                        break;
236                    }
237                }
238                Err(error) => {
239                    if tx.send(Err(error)).await.is_err() {
240                        break;
241                    }
242                }
243            }
244        }
245
246        Ok(())
247    }
248}
249
250fn validate_filter(filter: &StreamFilter) -> crate::Result<()> {
251    if filter.is_active() {
252        Ok(())
253    } else {
254        Err(crate::Error::InvalidFilter(
255            "filter must include at least one value or match_unfiltered=true".into(),
256        ))
257    }
258}