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::namespace::{Namespace, StreamNamespace};
11use crate::request_id::new_request_id;
12use crate::store::EventStore;
13use crate::stream::ReadStreamBackend;
14use crate::subscription::Subscription;
15use crate::version::VersionedDelivery;
16use crate::{DEFAULT_MAX_EVENT_BYTES, MAX_STREAM_HISTORY_ITEMS};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19enum ReadStreamLifecycle {
20    Populating,
21    Ready,
22}
23
24/// Coordinates read-side subscriptions with history replay and live tailing.
25pub struct NStreamsHandler<S, R> {
26    store: Arc<S>,
27    read_stream: Arc<R>,
28    max_event_bytes: u64,
29    lifecycle: Arc<RwLock<HashMap<String, ReadStreamLifecycle>>>,
30    population_locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
31}
32
33impl<S, R> Clone for NStreamsHandler<S, R> {
34    fn clone(&self) -> Self {
35        Self {
36            store: Arc::clone(&self.store),
37            read_stream: Arc::clone(&self.read_stream),
38            max_event_bytes: self.max_event_bytes,
39            lifecycle: Arc::clone(&self.lifecycle),
40            population_locks: Arc::clone(&self.population_locks),
41        }
42    }
43}
44
45impl<S, R> NStreamsHandler<S, R>
46where
47    S: EventStore + 'static,
48    R: ReadStreamBackend + 'static,
49{
50    pub fn new(store: S, read_stream: R) -> Self {
51        Self::with_max_event_bytes(store, read_stream, DEFAULT_MAX_EVENT_BYTES)
52    }
53
54    pub fn with_max_event_bytes(store: S, read_stream: R, max_event_bytes: u64) -> Self {
55        Self {
56            store: Arc::new(store),
57            read_stream: Arc::new(read_stream),
58            max_event_bytes,
59            lifecycle: Arc::new(RwLock::new(HashMap::new())),
60            population_locks: Arc::new(Mutex::new(HashMap::new())),
61        }
62    }
63
64    /// Subscribe to a namespace, replaying up to `history_count` events before live tail.
65    pub async fn subscribe<N: Namespace>(
66        &self,
67        namespace: &N,
68        history_count: u64,
69    ) -> crate::Result<Subscription> {
70        StreamNamespace::validate(namespace.as_str())?;
71        self.store.ensure_namespace(namespace).await?;
72        self.ensure_read_stream(namespace).await?;
73
74        let history_limit = history_count.min(MAX_STREAM_HISTORY_ITEMS);
75        let history = self.store.load_history(namespace, history_limit).await?;
76
77        let request_id = new_request_id();
78        let (tx, rx) = tokio::sync::mpsc::channel(1024);
79
80        let handler = self.clone();
81        let namespace_owned = StreamNamespace::from(namespace.as_str());
82        let request_id_for_task = request_id.clone();
83        tokio::spawn(async move {
84            if let Err(error) = handler
85                .run_subscription(request_id_for_task, namespace_owned, history, tx)
86                .await
87            {
88                warn!(%error, "subscription task failed");
89            }
90        });
91
92        Ok(Subscription::new(request_id, rx))
93    }
94
95    async fn ensure_read_stream<N: Namespace>(&self, namespace: &N) -> crate::Result<()> {
96        let key = namespace.as_str().to_string();
97
98        {
99            let lifecycle = self.lifecycle.read().await;
100            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
101                return Ok(());
102            }
103        }
104
105        let lock = {
106            let mut locks = self.population_locks.lock().await;
107            locks
108                .entry(key.clone())
109                .or_insert_with(|| Arc::new(Mutex::new(())))
110                .clone()
111        };
112
113        let _guard = lock.lock().await;
114
115        {
116            let lifecycle = self.lifecycle.read().await;
117            if lifecycle.get(&key) == Some(&ReadStreamLifecycle::Ready) {
118                return Ok(());
119            }
120        }
121
122        if self.read_stream.stream_exists(namespace).await? {
123            let mut lifecycle = self.lifecycle.write().await;
124            lifecycle.insert(key, ReadStreamLifecycle::Ready);
125            return Ok(());
126        }
127
128        {
129            let mut lifecycle = self.lifecycle.write().await;
130            lifecycle.insert(key.clone(), ReadStreamLifecycle::Populating);
131        }
132
133        info!(namespace = namespace.as_str(), "creating read stream");
134        self.read_stream
135            .create_read_stream(namespace, self.max_event_bytes)
136            .await?;
137
138        let bootstrap_limit = MAX_STREAM_HISTORY_ITEMS;
139        let history = self.store.load_history(namespace, bootstrap_limit).await?;
140        if !history.is_empty() {
141            debug!(
142                namespace = namespace.as_str(),
143                count = history.len(),
144                "bootstrapping read stream from postgres"
145            );
146            self.read_stream.populate_stream(&history).await?;
147        }
148
149        let mut lifecycle = self.lifecycle.write().await;
150        lifecycle.insert(key, ReadStreamLifecycle::Ready);
151        info!(namespace = namespace.as_str(), "read stream ready");
152        Ok(())
153    }
154
155    pub async fn is_read_stream_ready<N: Namespace>(&self, namespace: &N) -> bool {
156        let lifecycle = self.lifecycle.read().await;
157        lifecycle.get(namespace.as_str()) == Some(&ReadStreamLifecycle::Ready)
158    }
159
160    async fn run_subscription(
161        &self,
162        request_id: String,
163        namespace: StreamNamespace,
164        history: Vec<StreamEvent>,
165        tx: tokio::sync::mpsc::Sender<crate::Result<DeliveredEvent>>,
166    ) -> crate::Result<()> {
167        let mut last_version: Option<EventVersion> = None;
168
169        for event in history {
170            last_version = Some(VersionedDelivery::verify_initial(last_version, &event)?);
171            let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
172            if tx.send(Ok(delivered)).await.is_err() {
173                return Ok(());
174            }
175        }
176
177        let live = self.read_stream.subscribe_live(&namespace).await?;
178        pin!(live);
179        while let Some(item) = live.next().await {
180            match item {
181                Ok(event) => {
182                    if let Some(last) = last_version {
183                        if event.version <= last {
184                            continue;
185                        }
186                        VersionedDelivery { version: last }.verify_next(&event)?;
187                    }
188                    last_version = Some(event.version);
189                    let delivered = DeliveredEvent::from_stream_event(&request_id, &event);
190                    if tx.send(Ok(delivered)).await.is_err() {
191                        break;
192                    }
193                }
194                Err(error) => {
195                    if tx.send(Err(error)).await.is_err() {
196                        break;
197                    }
198                }
199            }
200        }
201
202        Ok(())
203    }
204}
nstreams_core/handler.rs

nstreams_core/
handler.rs
