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                after_version: None,
78            },
79        )
80        .await
81    }
82
83    /// Subscribe with an explicit request, including optional RabbitMQ stream filters.
84    pub async fn subscribe_with_request(
85        &self,
86        request: SubscribeRequest,
87    ) -> crate::Result<Subscription> {
88        let namespace = StreamNamespace::from(request.namespace.as_str());
89        self.subscribe_with_filter(&namespace, request).await
90    }
91
92    async fn subscribe_with_filter<N: Namespace>(
93        &self,
94        namespace: &N,
95        request: SubscribeRequest,
96    ) -> crate::Result<Subscription> {
97        if let Some(filter) = request.filter.as_ref() {
98            validate_filter(filter)?;
99        }
100
101        StreamNamespace::validate(namespace.as_str())?;
102        self.store.ensure_namespace(namespace).await?;
103        self.ensure_read_stream(namespace).await?;
104
105        let history_limit = request.history.min(MAX_STREAM_HISTORY_ITEMS);
106        let mut history = if let Some(after_version) = request.after_version {
107            self.store
108                .load_after_version(namespace, after_version, history_limit)
109                .await?
110        } else {
111            self.store.load_history(namespace, history_limit).await?
112        };
113        if let Some(filter) = request.filter.as_ref() {
114            history.retain(|event| filter.matches_filter_value(event.filter_value.as_deref()));
115        }
116
117        let request_id = new_request_id();
118        let (tx, rx) = tokio::sync::mpsc::channel(1024);
119
120        let handler = self.clone();
121        let namespace_owned = StreamNamespace::from(namespace.as_str());
122        let request_id_for_task = request_id.clone();
123        let filter = request.filter;
124        tokio::spawn(async move {
125            if let Err(error) = handler
126                .run_subscription(request_id_for_task, namespace_owned, history, filter, tx)
127                .await
128            {
129                warn!(%error, "subscription task failed");
130            }
131        });
132
133        Ok(Subscription::new(request_id, rx))
134    }
135
136    async fn ensure_read_stream<N: Namespace>(&self, namespace: &N) -> crate::Result<()> {
137        let key = namespace.as_str().to_string();
138
139        {
140            let lifecycle = self.lifecycle.read().await;
141            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
142                return Ok(());
143            }
144        }
145
146        let lock = {
147            let mut locks = self.population_locks.lock().await;
148            locks
149                .entry(key.clone())
150                .or_insert_with(|| Arc::new(Mutex::new(())))
151                .clone()
152        };
153
154        let _guard = lock.lock().await;
155
156        {
157            let lifecycle = self.lifecycle.read().await;
158            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
159                return Ok(());
160            }
161        }
162
163        if self.read_stream.stream_exists(namespace).await? {
164            let mut lifecycle = self.lifecycle.write().await;
165            lifecycle.insert(key, ReadStreamLifecycle::Ready);
166            return Ok(());
167        }
168
169        {
170            let mut lifecycle = self.lifecycle.write().await;
171            lifecycle.insert(key.clone(), ReadStreamLifecycle::Populating);
172        }
173
174        info!(namespace = namespace.as_str(), "creating read stream");
175        self.read_stream
176            .create_read_stream(namespace, self.max_event_bytes)
177            .await?;
178
179        let bootstrap_limit = MAX_STREAM_HISTORY_ITEMS;
180        let history = self.store.load_history(namespace, bootstrap_limit).await?;
181        if !history.is_empty() {
182            debug!(
183                namespace = namespace.as_str(),
184                count = history.len(),
185                "bootstrapping read stream from postgres"
186            );
187            self.read_stream.populate_stream(&history).await?;
188        }
189
190        let mut lifecycle = self.lifecycle.write().await;
191        lifecycle.insert(key, ReadStreamLifecycle::Ready);
192        info!(namespace = namespace.as_str(), "read stream ready");
193        Ok(())
194    }
195
196    pub async fn is_read_stream_ready<N: Namespace>(&self, namespace: &N) -> bool {
197        let lifecycle = self.lifecycle.read().await;
198        lifecycle.get(namespace.as_str()) == Some(&ReadStreamLifecycle::Ready)
199    }
200
201    async fn run_subscription(
202        &self,
203        request_id: String,
204        namespace: StreamNamespace,
205        history: Vec<StreamEvent>,
206        filter: Option<StreamFilter>,
207        tx: tokio::sync::mpsc::Sender<crate::Result<DeliveredEvent>>,
208    ) -> crate::Result<()> {
209        let mut last_version: Option<EventVersion> = None;
210
211        for event in history {
212            last_version = Some(VersionedDelivery::verify_initial(last_version, &event)?);
213            let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
214            if tx.send(Ok(delivered)).await.is_err() {
215                return Ok(());
216            }
217        }
218
219        let live = self
220            .read_stream
221            .subscribe_live(&namespace, filter.as_ref())
222            .await?;
223        pin!(live);
224        while let Some(item) = live.next().await {
225            match item {
226                Ok(event) => {
227                    if let Some(filter) = filter.as_ref() {
228                        if !filter.matches_filter_value(event.filter_value.as_deref()) {
229                            continue;
230                        }
231                    }
232
233                    if let Some(last) = last_version {
234                        if event.version <= last {
235                            continue;
236                        }
237                        VersionedDelivery { version: last }.verify_next(&event)?;
238                    }
239                    last_version = Some(event.version);
240                    let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
241                    if tx.send(Ok(delivered)).await.is_err() {
242                        break;
243                    }
244                }
245                Err(error) => {
246                    if tx.send(Err(error)).await.is_err() {
247                        break;
248                    }
249                }
250            }
251        }
252
253        Ok(())
254    }
255}
256
257fn validate_filter(filter: &StreamFilter) -> crate::Result<()> {
258    if filter.is_active() {
259        Ok(())
260    } else {
261        Err(crate::Error::InvalidFilter(
262            "filter must include at least one value or match_unfiltered=true".into(),
263        ))
264    }
265}