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 },
78 )
79 .await
80 }
81
82 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}