Line data Source code
1 : // SPDX-License-Identifier: EUPL-1.2
2 : //! Group commit for the store calls of in-process deliveries.
3 : //!
4 : //! A fan-out of one change to 1 000 socket subscriptions on one view used
5 : //! to cost 1 000 reads of the same Entities and 1 000 bookkeeping
6 : //! transactions, each waiting for a pool connection. Here a delivery that
7 : //! reaches the store joins the NEXT statement of its tenant: while one
8 : //! statement runs, the deliveries arriving behind it queue, and the next
9 : //! statement serves all of them. A delivery never joins a statement already
10 : //! running, so the Entities it reads are read after it asked
11 : //! rule 9: the view as the table serves it now). A tenant of a store has at
12 : //! most one booking and one read statement in flight, whatever its fan-out.
13 :
14 : use super::*;
15 : use antares_store::Delivery;
16 :
17 : /// What waits for a tenant's next statement: each question with the channel
18 : /// its answer goes back on.
19 : #[cfg(not(target_arch = "wasm32"))]
20 : type Waiting<Q, A> = Vec<(Q, tokio::sync::oneshot::Sender<A>)>;
21 :
22 : /// A queue's key: the store the statement runs on and the tenant. The store
23 : /// part keeps two brokers of one process (tests run several, with the same
24 : /// tenants) apart: a statement runs on the first caller's store, so a
25 : /// question of another store answered there found nothing and its
26 : /// notification was dropped. A runner holds its store, so no other
27 : /// store reuses the address while the entry lives.
28 : #[cfg(not(target_arch = "wasm32"))]
29 : type Key = (usize, String);
30 :
31 : #[cfg(not(target_arch = "wasm32"))]
32 178 : fn key(st: &AppState, tenant: &TenantId) -> Key {
33 178 : (
34 178 : Arc::as_ptr(&st.store).cast::<()>().addr(),
35 178 : tenant.as_str().to_owned(),
36 178 : )
37 178 : }
38 :
39 : /// One queue per store and tenant. An entry exists exactly while a runner
40 : /// serves it, so a caller that finds none starts one, and the map holds only
41 : /// queues with a statement in flight.
42 : #[cfg(not(target_arch = "wasm32"))]
43 : pub(super) struct Coalescer<Q, A> {
44 : queues: std::sync::Mutex<HashMap<Key, Waiting<Q, A>>>,
45 : }
46 :
47 : #[cfg(not(target_arch = "wasm32"))]
48 : impl<Q: Send + 'static, A: Send + 'static> Coalescer<Q, A> {
49 18 : pub(super) fn new() -> Self {
50 18 : Self {
51 18 : queues: std::sync::Mutex::new(HashMap::new()),
52 18 : }
53 18 : }
54 :
55 : /// `q`'s answer from the next statement `run` makes for `tenant`, one
56 : /// answer per question in order. `None` when the runner died before
57 : /// answering (a panic inside `run`).
58 : ///
59 : /// A caller that finds its tenant idle runs its own statement on its own
60 : /// task, so an uncontended delivery costs what it did before; only the
61 : /// questions that queued behind that statement get a runner task.
62 168 : pub(super) async fn call<F, Fut>(
63 168 : &'static self,
64 168 : st: &AppState,
65 168 : tenant: &TenantId,
66 168 : q: Q,
67 168 : run: F,
68 168 : ) -> Option<A>
69 168 : where
70 168 : F: Fn(AppState, TenantId, Vec<Q>) -> Fut + Clone + Send + 'static,
71 168 : Fut: std::future::Future<Output = Vec<A>> + Send,
72 168 : {
73 168 : let key = key(st, tenant);
74 168 : let waiting = {
75 : // A poisoned map is a panic in some past runner; its guard
76 : // already removed that runner's entry.
77 168 : let mut queues = self.queues.lock().unwrap_or_else(|p| p.into_inner());
78 168 : match queues.get_mut(&key) {
79 130 : Some(waiting) => {
80 130 : let (tx, rx) = tokio::sync::oneshot::channel();
81 130 : waiting.push((q, tx));
82 130 : Err(rx)
83 : }
84 : None => {
85 38 : queues.insert(key.clone(), Vec::new());
86 38 : Ok(q)
87 : }
88 : }
89 : };
90 168 : let q = match waiting {
91 38 : Ok(q) => q,
92 130 : Err(rx) => return rx.await.ok(),
93 : };
94 : // Dropped mid-statement (its delivery's deadline) or done, the
95 : // caller hands what queued behind it to a runner: a question that
96 : // waited is never lost because the caller ahead of it went away.
97 38 : let respawn = {
98 38 : let (st, tenant, run) = (st.clone(), tenant.clone(), run.clone());
99 10 : move || {
100 10 : tokio::spawn(async move { self.serve(&st, &tenant, run).await });
101 10 : }
102 : };
103 38 : let _handoff = Retire {
104 38 : queues: &self.queues,
105 38 : key: &key,
106 38 : armed: true,
107 38 : respawn: Some(Box::new(respawn)),
108 38 : };
109 38 : run(st.clone(), tenant.clone(), vec![q]).await.pop()
110 164 : }
111 :
112 : /// Serve `tenant`'s queue until it is empty, one statement per round.
113 10 : async fn serve<F, Fut>(&self, st: &AppState, tenant: &TenantId, run: F)
114 10 : where
115 10 : F: Fn(AppState, TenantId, Vec<Q>) -> Fut,
116 10 : Fut: std::future::Future<Output = Vec<A>>,
117 10 : {
118 10 : let key = key(st, tenant);
119 10 : let mut guard = Retire {
120 10 : queues: &self.queues,
121 10 : key: &key,
122 10 : armed: true,
123 10 : respawn: None,
124 10 : };
125 : loop {
126 10 : let batch = {
127 20 : let mut queues = self.queues.lock().unwrap_or_else(|p| p.into_inner());
128 20 : match queues.get_mut(&key) {
129 20 : Some(waiting) if !waiting.is_empty() => std::mem::take(waiting),
130 : _ => {
131 10 : queues.remove(&key);
132 10 : guard.armed = false;
133 10 : return;
134 : }
135 : }
136 : };
137 10 : let (qs, txs): (Vec<Q>, Vec<_>) = batch.into_iter().unzip();
138 10 : let answers = run(st.clone(), tenant.clone(), qs).await;
139 130 : for (tx, answer) in txs.into_iter().zip(answers) {
140 130 : // the caller gave up (its delivery deadline): nothing to tell
141 130 : let _gone = tx.send(answer);
142 130 : }
143 : }
144 10 : }
145 : }
146 :
147 : /// What a runner leaves behind. With `respawn` (the caller that ran its own
148 : /// statement), questions still waiting go to a new runner, whether the caller
149 : /// finished or was dropped. Without it, or when unwinding from a panic, the
150 : /// entry goes, so the tenant's next caller starts afresh instead of queueing
151 : /// behind a dead runner; the questions still queued then get `None`.
152 : #[cfg(not(target_arch = "wasm32"))]
153 : struct Retire<'a, Q, A> {
154 : queues: &'a std::sync::Mutex<HashMap<Key, Waiting<Q, A>>>,
155 : key: &'a Key,
156 : armed: bool,
157 : respawn: Option<Box<dyn FnOnce() + Send + 'a>>,
158 : }
159 :
160 : #[cfg(not(target_arch = "wasm32"))]
161 : impl<Q, A> Drop for Retire<'_, Q, A> {
162 48 : fn drop(&mut self) {
163 48 : if !self.armed {
164 10 : return;
165 38 : }
166 38 : let mut queues = self.queues.lock().unwrap_or_else(|p| p.into_inner());
167 38 : let waiting = queues.get(self.key).is_some_and(|w| !w.is_empty());
168 38 : match self.respawn.take() {
169 38 : Some(respawn) if waiting && !std::thread::panicking() => {
170 10 : drop(queues);
171 10 : respawn();
172 10 : }
173 28 : _ => {
174 28 : queues.remove(self.key);
175 28 : }
176 : }
177 48 : }
178 : }
179 :
180 : /// The 5.2.14.2 stamp of one in-process delivery, written with every other
181 : /// in-process delivery of its tenant waiting for the same statement. `None`:
182 : /// the subscription is gone or the write failed, and nothing may be sent.
183 116 : pub(super) async fn book_in_process(
184 116 : st: &AppState,
185 116 : tenant: &TenantId,
186 116 : sub_id: &str,
187 116 : now: &str,
188 116 : ) -> Option<Delivery> {
189 : #[cfg(not(target_arch = "wasm32"))]
190 : {
191 : static BOOKINGS: std::sync::LazyLock<Coalescer<(String, String), Option<Delivery>>> =
192 : std::sync::LazyLock::new(Coalescer::new);
193 116 : BOOKINGS
194 116 : .call(
195 116 : st,
196 116 : tenant,
197 116 : (sub_id.to_owned(), now.to_owned()),
198 40 : |st, tenant, attempts| async move { book(&st, &tenant, &attempts).await },
199 : )
200 116 : .await
201 116 : .flatten()
202 : }
203 : #[cfg(target_arch = "wasm32")]
204 : {
205 : book(st, tenant, &[(sub_id.to_owned(), now.to_owned())])
206 : .await
207 : .pop()
208 : .flatten()
209 : }
210 116 : }
211 :
212 : /// One `record_deliveries` per round of distinct subscriptions. Two waiting
213 : /// attempts on one subscription (two notifications in flight for it) go in
214 : /// consecutive rounds, so each books exactly one notification and gets back
215 : /// the `lastSuccess` it overwrote.
216 24 : pub(super) async fn book(
217 24 : st: &AppState,
218 24 : tenant: &TenantId,
219 24 : attempts: &[(String, String)],
220 24 : ) -> Vec<Option<Delivery>> {
221 24 : let mut out: Vec<Option<Delivery>> = vec![None; attempts.len()];
222 24 : let mut left: Vec<usize> = (0..attempts.len()).collect();
223 52 : while !left.is_empty() {
224 28 : let mut seen = std::collections::HashSet::new();
225 136 : let (round, later): (Vec<usize>, Vec<usize>) = left.into_iter().partition(|i| {
226 136 : attempts
227 136 : .get(*i)
228 136 : .is_some_and(|(id, _)| seen.insert(id.as_str()))
229 136 : });
230 28 : let stamps: Vec<(&str, &str)> = round
231 28 : .iter()
232 132 : .filter_map(|i| attempts.get(*i))
233 132 : .map(|(id, now)| (id.as_str(), now.as_str()))
234 28 : .collect();
235 28 : match st
236 28 : .store
237 28 : .record_deliveries(tenant, Kind::Subscription, &stamps)
238 28 : .await
239 : {
240 28 : Ok(found) => {
241 132 : for (i, d) in round.into_iter().zip(found) {
242 132 : if let Some(slot) = out.get_mut(i) {
243 132 : *slot = d;
244 132 : }
245 : }
246 : }
247 0 : Err(e) => tracing::warn!("bookkeeping writeback failed: {e}"),
248 : }
249 28 : left = later;
250 : }
251 24 : out
252 24 : }
253 :
254 : /// The failure write-back of one in-process delivery
255 : /// ([`super::delivery::record_failure`]), queued per tenant behind the one in
256 : /// flight, so overflowing sockets of one fan-out never hold more than one
257 : /// pool connection for it. Rounds write in arrival order, so two failures of
258 : /// one subscription roll `lastSuccess` back in the order they were booked.
259 0 : pub(super) async fn fail_in_process(
260 0 : st: &AppState,
261 0 : tenant: &TenantId,
262 0 : sub_id: &str,
263 0 : prev_success: Option<Value>,
264 0 : ts: &str,
265 0 : ) -> Option<Value> {
266 : #[cfg(not(target_arch = "wasm32"))]
267 : {
268 : type Failure = (String, Option<Value>, String);
269 : static FAILURES: std::sync::LazyLock<Coalescer<Failure, Option<Value>>> =
270 : std::sync::LazyLock::new(Coalescer::new);
271 0 : FAILURES
272 0 : .call(
273 0 : st,
274 0 : tenant,
275 0 : (sub_id.to_owned(), prev_success, ts.to_owned()),
276 0 : |st, tenant, failures: Vec<Failure>| async move {
277 0 : let mut docs = Vec::with_capacity(failures.len());
278 0 : for (id, prev, ts) in failures {
279 0 : docs.push(
280 0 : super::delivery::record_failure(
281 0 : &st,
282 0 : &tenant,
283 0 : Kind::Subscription,
284 0 : &id,
285 0 : prev,
286 0 : &ts,
287 0 : )
288 0 : .await,
289 : );
290 : }
291 0 : docs
292 0 : },
293 : )
294 0 : .await
295 0 : .flatten()
296 : }
297 : #[cfg(target_arch = "wasm32")]
298 : {
299 : super::delivery::record_failure(st, tenant, Kind::Subscription, sub_id, prev_success, ts)
300 : .await
301 : }
302 0 : }
303 :
304 : /// The stored Entities `ids` names, read in one statement with every other
305 : /// delivery of the tenant waiting for a read. `None`: the read failed (logged
306 : /// once for the whole statement).
307 4 : pub(super) async fn read_entities(
308 4 : st: &AppState,
309 4 : tenant: &TenantId,
310 4 : ids: Vec<String>,
311 4 : ) -> Option<Arc<HashMap<String, Value>>> {
312 : #[cfg(not(target_arch = "wasm32"))]
313 : {
314 : type Read = Option<Arc<HashMap<String, Value>>>;
315 : static READS: std::sync::LazyLock<Coalescer<Vec<String>, Read>> =
316 : std::sync::LazyLock::new(Coalescer::new);
317 4 : READS
318 4 : .call(
319 4 : st,
320 4 : tenant,
321 4 : ids,
322 4 : |st, tenant, asks: Vec<Vec<String>>| async move {
323 4 : let found = read(&st, &tenant, asks.iter().flatten()).await;
324 4 : vec![found; asks.len()]
325 8 : },
326 : )
327 4 : .await
328 4 : .flatten()
329 : }
330 : #[cfg(target_arch = "wasm32")]
331 : {
332 : read(st, tenant, ids.iter()).await
333 : }
334 4 : }
335 :
336 : /// The union of `ids` in reads of at most `chunk` ids each:
337 : /// a large fan-out of folded deliveries never becomes one statement.
338 8 : fn read_batches<'a>(ids: impl Iterator<Item = &'a String>, chunk: usize) -> Vec<Vec<&'a str>> {
339 8 : let union: std::collections::BTreeSet<&str> = ids.map(String::as_str).collect();
340 8 : let union: Vec<&str> = union.into_iter().collect();
341 8 : union.chunks(chunk.max(1)).map(<[&str]>::to_vec).collect()
342 8 : }
343 :
344 : /// `get_entities` over the union of `ids`, one batch at a time; a failed
345 : /// batch fails the whole read, since a partial map would read as deletions.
346 4 : async fn read<'a>(
347 4 : st: &AppState,
348 4 : tenant: &TenantId,
349 4 : ids: impl Iterator<Item = &'a String>,
350 4 : ) -> Option<Arc<HashMap<String, Value>>> {
351 4 : let mut found = HashMap::new();
352 4 : for batch in read_batches(ids, *crate::bounds::MAX_BATCH_ITEMS) {
353 4 : match st.store.get_entities(tenant, &batch).await {
354 4 : Ok(part) => found.extend(part),
355 0 : Err(e) => {
356 0 : tracing::warn!("notification path: reading the notified Entities failed: {e}");
357 0 : return None;
358 : }
359 : }
360 : }
361 4 : Some(Arc::new(found))
362 4 : }
363 :
364 : #[cfg(test)]
365 : mod tests {
366 : use super::read_batches;
367 :
368 : /// 1000 asks of 256 ids each, overlapping, read as their union in
369 : /// batches no larger than the chunk.
370 : #[test]
371 4 : fn a_large_union_is_read_in_bounded_batches() {
372 4 : let asks: Vec<Vec<String>> = (0..1000)
373 1024000 : .map(|a| (0..256).map(|i| format!("urn:x:{}", a * 100 + i)).collect())
374 4 : .collect();
375 4 : let batches = read_batches(asks.iter().flatten(), 1000);
376 404 : assert!(batches.iter().all(|b| b.len() <= 1000));
377 4 : let total: usize = batches.iter().map(Vec::len).sum();
378 4 : assert_eq!(total, 999 * 100 + 256, "each id of the union once");
379 4 : }
380 : }
|
