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
25pub 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 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 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}