Skip to main content

apalis_diesel_postgres/
shared.rs

1use std::{
2    collections::HashMap,
3    marker::PhantomData,
4    pin::Pin,
5    sync::{
6        Arc, Mutex,
7        atomic::{AtomicBool, Ordering},
8    },
9    task::{Context, Poll},
10};
11
12/// `SharedRegistry` now stores **multiple senders per queue** instead of a
13/// single `Arc<Mutex<Receiver>>` shared by clones. Each consumer
14/// (`make_shared_with_config` call or `SharedFetcher::clone`) gets its own
15/// mpsc channel; the listener broadcasts to every sender bound to the matching
16/// `job_type`. This removes the `Arc<Mutex<Receiver>>` contention smell and
17/// lets fetcher polls run without mutex acquisition.
18type RegistrySender = Sender<Result<PgTaskId, Error>>;
19
20use apalis_codec::json::JsonCodec;
21use apalis_core::{backend::shared::MakeShared, worker::context::WorkerContext};
22use diesel::RunQueryDsl;
23use futures::{
24    Stream,
25    channel::mpsc::{self, Receiver, Sender},
26};
27use ulid::Ulid;
28
29use crate::{
30    CompactType, Config, Error, PgPool, PgTask, PgTaskId, PostgresStorage, queries, sink::PgSink,
31};
32
33/// Per-registration sender entry.
34///
35/// `id` uniquely identifies the `SharedRegistration` that owns this sender so
36/// that `SharedRegistration::drop` can prune only its own sender from the
37/// queue's Vec instead of wiping the whole entry (which would silently sever
38/// every other consumer on the same queue).
39type RegistryEntry = (Ulid, RegistrySender);
40type RegistryMap = HashMap<String, Vec<RegistryEntry>>;
41type SharedRegistry = Arc<Mutex<RegistryMap>>;
42
43/// Factory for shared notify-backed PostgreSQL storage instances.
44///
45/// A shared storage factory owns one listener thread and one pooled PostgreSQL
46/// connection for notifications. A queue may be registered multiple times: the
47/// single listener broadcasts every notification to all consumers registered on
48/// the matching queue.
49pub struct SharedPostgresStorage<Codec = JsonCodec<CompactType>> {
50    pool: PgPool,
51    registry: SharedRegistry,
52    /// Single source of truth for «listener thread is alive». `make_shared_…`
53    /// CAS-swaps it to `true` and spawns a listener only on the `false → true`
54    /// transition; the listener clears it on exit. Replaces the prior
55    /// `registry.is_empty()` heuristic, which had a race window: when the last
56    /// registration dropped and a new one was added before the old listener's
57    /// next empty-check, the new caller saw `is_empty == false` and skipped
58    /// spawning — leaving the new registration with no listener (or, in the
59    /// mirror case, briefly running two listeners and double-delivering).
60    listener_alive: Arc<AtomicBool>,
61    _marker: PhantomData<Codec>,
62}
63
64impl<Codec> SharedPostgresStorage<Codec> {
65    /// Create a shared storage factory.
66    #[must_use]
67    pub fn new(pool: PgPool) -> Self {
68        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
69        Self {
70            pool,
71            registry,
72            listener_alive: Arc::new(AtomicBool::new(false)),
73            _marker: PhantomData,
74        }
75    }
76
77    fn spawn_registry_listener(&self) {
78        let pool = self.pool.clone();
79        let registry = self.registry.clone();
80        let listener_alive = self.listener_alive.clone();
81        if let Err(error) = std::thread::Builder::new()
82            .name("apalis-postgres-shared-listener".to_owned())
83            .spawn(move || {
84                let mut conn = match pool.get() {
85                    Ok(conn) => conn,
86                    Err(error) => {
87                        exit_listener(
88                            &registry,
89                            &listener_alive,
90                            Some(format!(
91                                "failed to get pooled connection for shared LISTEN: {error}"
92                            )),
93                        );
94                        return;
95                    }
96                };
97                if let Err(error) =
98                    diesel::sql_query("LISTEN \"apalis::job::insert\"").execute(&mut conn)
99                {
100                    exit_listener(
101                        &registry,
102                        &listener_alive,
103                        Some(format!("failed to start shared LISTEN listener: {error}")),
104                    );
105                    return;
106                }
107                run_listener_loop(&mut conn, &registry, &listener_alive);
108                // Remove the subscription before the pooled connection drops
109                // back into r2d2: the next pool user would otherwise inherit
110                // it and notifications would silently accumulate in libpq's
111                // receive buffer with no consumer draining them (mirrors the
112                // single-queue listener cleanup in `queries/notify.rs`).
113                // Best-effort: on failure the connection is in an unknown
114                // state and r2d2's checkout health check decides its fate.
115                let _ = diesel::sql_query("UNLISTEN \"apalis::job::insert\"").execute(&mut conn);
116            })
117        {
118            exit_listener(
119                &self.registry,
120                &self.listener_alive,
121                Some(format!("failed to spawn listener: {error}")),
122            );
123        }
124    }
125}
126
127/// Body of the shared listener thread, between a successful `LISTEN` and the
128/// `UNLISTEN` cleanup. Factored out of `spawn_registry_listener` so every exit
129/// path — notification error, poisoned registry, empty registry — funnels
130/// through a single return boundary, after which the caller removes the
131/// LISTEN subscription before the pooled connection is returned to r2d2.
132fn run_listener_loop(
133    conn: &mut diesel::r2d2::PooledConnection<
134        diesel::r2d2::ConnectionManager<diesel::PgConnection>,
135    >,
136    registry: &SharedRegistry,
137    listener_alive: &AtomicBool,
138) {
139    loop {
140        for notification in conn.notifications_iter() {
141            let notification = match notification {
142                Ok(notification) => notification,
143                Err(error) => {
144                    exit_listener(
145                        registry,
146                        listener_alive,
147                        Some(format!("failed to receive shared notification: {error}")),
148                    );
149                    return;
150                }
151            };
152            let Ok(event) = serde_json::from_str::<crate::InsertEvent>(&notification.payload)
153            else {
154                continue;
155            };
156            let (event_queue, ids) = event.into_ids();
157            let Ok(mut registry) = registry.lock() else {
158                // Poisoned: we cannot synchronize with registrants
159                // any longer, fall back to a bare store.
160                listener_alive.store(false, Ordering::Release);
161                return;
162            };
163            deliver_to_queue(&mut registry, &event_queue, &ids);
164        }
165        match registry.lock() {
166            Ok(registry) => {
167                // Store `false` while still holding the registry
168                // lock: a concurrent `make_shared_with_config`
169                // must observe either (a) `listener_alive == true`
170                // (we haven't exited yet) AND see itself appended
171                // to the registry on our next loop iteration, or
172                // (b) `listener_alive == false` AND therefore
173                // spawn a fresh listener. The exit decision is
174                // factored into `listener_should_exit` (unit-tested)
175                // and applied as a plain `if` rather than a match
176                // guard so the only mutable point is that function,
177                // which the tests pin — an in-thread guard would be
178                // unreachable from a unit test.
179                if listener_should_exit(&registry) {
180                    listener_alive.store(false, Ordering::Release);
181                    drop(registry);
182                    return;
183                }
184            }
185            Err(_) => {
186                // Poisoned: synchronization is no longer possible.
187                listener_alive.store(false, Ordering::Release);
188                return;
189            }
190        }
191        std::thread::sleep(queries::NOTIFY_LISTENER_POLL_INTERVAL);
192    }
193}
194
195/// Drop the listener under the registry lock so a concurrent
196/// `make_shared_with_config` cannot observe `listener_alive == true` AFTER the
197/// listener has decided to exit. The same lock serializes registrants'
198/// `swap(true)` against our `store(false)`, leaving exactly two possible
199/// orderings: (a) registrant runs first and sees `listener_alive == false`,
200/// spawning a fresh listener; (b) listener runs first, sees an empty registry,
201/// stores `false`, and a subsequent registrant spawns. Without this serialization
202/// a registrant could observe stale `true` and skip spawn.
203fn exit_listener(registry: &SharedRegistry, listener_alive: &AtomicBool, error: Option<String>) {
204    match registry.lock() {
205        Ok(mut guard) => {
206            if let Some(message) = error {
207                broadcast_notify_error_locked(&mut guard, message);
208            }
209            listener_alive.store(false, Ordering::Release);
210            drop(guard);
211        }
212        Err(_) => {
213            // Poisoned: best-effort store; we cannot synchronize with
214            // registrants any more.
215            listener_alive.store(false, Ordering::Release);
216        }
217    }
218}
219
220/// Broadcast every id in `ids` to all senders registered on `queue`.
221///
222/// A sender whose receiver has been dropped (`disconnected`) is pruned; a sender
223/// whose channel is merely full is **kept** — the job is durable and the poll
224/// fetcher will pick it up, so transient back-pressure must not sever the
225/// consumer. When the queue's last sender is pruned the queue entry is removed
226/// so the listener's empty-registry exit check can fire.
227///
228/// Extracted from the listener closure so the broadcast/prune decision is
229/// unit-testable against a hand-built registry, without spawning the listener
230/// thread (mirrors `broadcast_notify_error_locked`).
231fn deliver_to_queue(registry: &mut RegistryMap, queue: &str, ids: &[PgTaskId]) {
232    if let Some(senders) = registry.get_mut(queue) {
233        for &id in ids {
234            senders.retain_mut(|(_, sender)| match sender.try_send(Ok(id)) {
235                Ok(()) => true,
236                Err(error) if error.is_disconnected() => false,
237                // Channel full: keep the sender (the job is durable, the poll
238                // fetcher will pick it up) but stop pushing this event into a
239                // saturated channel.
240                Err(_) => true,
241            });
242        }
243        if senders.is_empty() {
244            registry.remove(queue);
245        }
246    }
247}
248
249/// Whether the shared listener should exit. It stops once no consumers remain
250/// registered, releasing its pooled connection and thread. Extracted so the
251/// empty-registry decision is unit-testable without spawning the listener.
252fn listener_should_exit(registry: &RegistryMap) -> bool {
253    registry.is_empty()
254}
255
256/// Whether THIS registration must spawn the shared listener: it does iff it
257/// observed `listener_alive` transition false→true, i.e. it is the first live
258/// registration. The atomic swap is the single source of truth, so exactly one
259/// registration spawns. Extracted so the false→true edge is unit-testable — the
260/// notify integration tests register two-or-more consumers, which masks an
261/// off-by-one in *which* registration spawns (the listener still ends up
262/// running, just claimed by the wrong call).
263fn claim_listener_spawn(listener_alive: &AtomicBool) -> bool {
264    !listener_alive.swap(true, Ordering::AcqRel)
265}
266
267#[cfg(test)]
268fn broadcast_notify_error(registry: &SharedRegistry, message: String) {
269    let Ok(mut guard) = registry.lock() else {
270        return;
271    };
272    broadcast_notify_error_locked(&mut guard, message);
273}
274
275fn broadcast_notify_error_locked(registry: &mut RegistryMap, message: String) {
276    registry.retain(|_, senders| {
277        senders.retain_mut(|(_, sender)| {
278            match sender.try_send(Err(Error::NotifyListener(message.clone()))) {
279                Ok(()) => true,
280                Err(error) => !error.is_disconnected(),
281            }
282        });
283        !senders.is_empty()
284    });
285}
286
287impl<Codec> std::fmt::Debug for SharedPostgresStorage<Codec> {
288    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289        f.debug_struct("SharedPostgresStorage")
290            .finish_non_exhaustive()
291    }
292}
293
294/// Errors returned while creating shared storage instances.
295#[derive(Debug, thiserror::Error)]
296#[non_exhaustive]
297pub enum SharedPostgresError {
298    /// Shared registry lock is poisoned.
299    #[error("registry lock poisoned")]
300    RegistryLocked,
301}
302
303impl<Args, Codec> MakeShared<Args> for SharedPostgresStorage<Codec> {
304    type Backend = PostgresStorage<Args, Codec, SharedFetcher>;
305    type Config = Config;
306    type MakeError = SharedPostgresError;
307
308    fn make_shared(&mut self) -> Result<Self::Backend, Self::MakeError>
309    where
310        Self::Config: Default,
311    {
312        self.make_shared_with_config(Config::new(std::any::type_name::<Args>()))
313    }
314
315    fn make_shared_with_config(
316        &mut self,
317        config: Self::Config,
318    ) -> Result<Self::Backend, Self::MakeError> {
319        let (sender, receiver) =
320            mpsc::channel(crate::queries::clamp_notify_capacity(config.buffer_size()));
321        let mut registry = self
322            .registry
323            .lock()
324            .map_err(|_| SharedPostgresError::RegistryLocked)?;
325        let queue = config.queue().to_string();
326        // Broadcast redesign: multiple consumers per queue are now allowed —
327        // each call appends its own sender to the queue's Vec, and the
328        // listener broadcasts to all of them. Previously the registry held a
329        // single Sender per queue and clones shared the Receiver via
330        // `Arc<Mutex<Receiver>>`, which serialized polls on a mutex.
331        //
332        // `listener_alive` is the single source of truth for «is a listener
333        // currently running». The swap and the registry mutation must happen
334        // under the *same* registry lock as the listener's exit decision
335        // (see `spawn_registry_listener`), otherwise the listener could store
336        // `false` between our `swap(true)` and our push, leaving a non-empty
337        // registry with no listener. By doing both under the lock we serialize
338        // the two state transitions onto the mutex.
339        let registration_id = Ulid::new();
340        registry
341            .entry(queue)
342            .or_default()
343            .push((registration_id, sender));
344        let should_spawn_listener = claim_listener_spawn(&self.listener_alive);
345        drop(registry);
346
347        if should_spawn_listener {
348            self.spawn_registry_listener();
349        }
350
351        let registration = Arc::new(SharedRegistration {
352            id: registration_id,
353            queue: config.queue().to_string(),
354            registry: self.registry.clone(),
355            pool: self.pool.clone(),
356        });
357
358        Ok(PostgresStorage {
359            _marker: PhantomData,
360            sink: PgSink::new(&self.pool, &config),
361            pool: self.pool.clone(),
362            config,
363            fetcher: SharedFetcher {
364                receiver,
365                _registration: registration,
366            },
367            lease_token: crate::queries::worker::mint_lease_token().into(),
368        })
369    }
370}
371
372struct SharedRegistration {
373    /// Identity of this registration's sender inside the queue's Vec.
374    /// `Drop` uses it to prune only this entry — wiping the whole queue
375    /// would silently sever every other consumer registered on the same
376    /// queue (broadcast design allows N senders per queue).
377    id: Ulid,
378    queue: String,
379    registry: SharedRegistry,
380    pool: PgPool,
381}
382
383impl std::fmt::Debug for SharedRegistration {
384    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
385        f.debug_struct("SharedRegistration")
386            .field("queue", &self.queue)
387            .finish_non_exhaustive()
388    }
389}
390
391impl Drop for SharedRegistration {
392    fn drop(&mut self) {
393        let became_empty = match self.registry.lock() {
394            Ok(mut registry) => {
395                // Prune only this registration's sender from the queue's
396                // Vec. If the Vec becomes empty (we were the last consumer
397                // on this queue), drop the queue entry too.
398                if let Some(senders) = registry.get_mut(&self.queue) {
399                    senders.retain(|(id, _)| *id != self.id);
400                    if senders.is_empty() {
401                        registry.remove(&self.queue);
402                    }
403                }
404                registry.is_empty()
405            }
406            Err(_) => false,
407        };
408        // When the registry becomes empty the shared listener thread will exit
409        // on its next loop iteration, but it is parked inside
410        // `notifications_iter`. Send a best-effort NOTIFY so the iterator
411        // returns and the empty-registry check runs immediately. The empty
412        // payload fails `serde_json::from_str::<InsertEvent>`, so any other
413        // listener simply ignores it.
414        if became_empty {
415            // Detach the blocking NOTIFY so the dropping task — which may be
416            // running on an async executor — never blocks on libpq.
417            let pool = self.pool.clone();
418            let _ = std::thread::Builder::new()
419                .name("apalis-postgres-shared-drop".to_owned())
420                .spawn(move || {
421                    if let Ok(mut conn) = pool.get() {
422                        let _ = diesel::sql_query("SELECT pg_notify('apalis::job::insert', '')")
423                            .execute(&mut conn);
424                    }
425                });
426        }
427    }
428}
429
430/// Fetcher used by shared storage instances.
431///
432/// After the broadcast redesign each `SharedFetcher` owns its own mpsc
433/// `Receiver` — no `Arc<Mutex<Receiver>>` indirection. The listener broadcasts
434/// every notification to every registered fetcher for that queue. As a
435/// consequence `SharedFetcher` is **not** `Clone`: cloning would require
436/// either splitting one receiver into two (impossible without locking) or
437/// silently producing a fetcher that never receives events. Use
438/// [`SharedPostgresStorage::make_shared_with_config`] to spawn additional
439/// consumers explicitly.
440pub struct SharedFetcher {
441    receiver: Receiver<Result<PgTaskId, Error>>,
442    _registration: Arc<SharedRegistration>,
443}
444
445impl std::fmt::Debug for SharedFetcher {
446    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
447        f.debug_struct("SharedFetcher").finish_non_exhaustive()
448    }
449}
450
451impl Stream for SharedFetcher {
452    type Item = Result<PgTaskId, Error>;
453
454    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
455        Pin::new(&mut self.get_mut().receiver).poll_next(cx)
456    }
457}
458
459impl crate::fetcher::PgFetcherSource for SharedFetcher {
460    const STORAGE_NAME: &'static str = "SharedPostgresStorage";
461
462    fn into_compact_stream(
463        self,
464        pool: PgPool,
465        config: Config,
466        worker: WorkerContext,
467        lease_token: std::sync::Arc<str>,
468    ) -> apalis_core::backend::TaskStream<PgTask<CompactType>, Error> {
469        crate::fetcher::notify_backed_compact_stream(
470            Self::STORAGE_NAME,
471            self,
472            pool,
473            config,
474            worker,
475            lease_token,
476        )
477    }
478}
479
480#[cfg(test)]
481mod tests {
482    use apalis_core::backend::{Backend, BackendExt, shared::MakeShared};
483    use diesel::{
484        PgConnection,
485        r2d2::{ConnectionManager, Pool},
486    };
487    use lets_expect::{AssertionError, AssertionResult, *};
488
489    use super::*;
490
491    struct SharedObservation {
492        queue: String,
493        buffer_size: usize,
494        debug: String,
495    }
496
497    fn unchecked_pool() -> PgPool {
498        let manager = ConnectionManager::<PgConnection>::new("postgres://127.0.0.1:1/not-used");
499        Pool::builder()
500            .max_size(1)
501            .connection_timeout(std::time::Duration::from_millis(10))
502            .build_unchecked(manager)
503    }
504
505    fn shared_debug() -> String {
506        let shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
507        format!("{shared:?}")
508    }
509
510    fn make_default_shared() -> Result<SharedObservation, SharedPostgresError> {
511        let mut shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
512        let storage = <SharedPostgresStorage as MakeShared<String>>::make_shared(&mut shared)?;
513        Ok(SharedObservation {
514            queue: storage.config.queue().to_string(),
515            buffer_size: storage.config.buffer_size(),
516            debug: format!("{storage:?}"),
517        })
518    }
519
520    fn make_configured_shared() -> Result<SharedObservation, SharedPostgresError> {
521        let mut shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
522        let config = Config::new("shared-unit").set_buffer_size(3);
523        let storage = <SharedPostgresStorage as MakeShared<String>>::make_shared_with_config(
524            &mut shared,
525            config,
526        )?;
527        Ok(SharedObservation {
528            queue: storage.get_queue().to_string(),
529            buffer_size: storage.config.buffer_size(),
530            debug: format!("{:?}", storage.fetcher),
531        })
532    }
533
534    fn shared_trait_surfaces() -> Result<(String, String), SharedPostgresError> {
535        let mut shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
536        let config = Config::new("shared-traits");
537        let storage = <SharedPostgresStorage as MakeShared<String>>::make_shared_with_config(
538            &mut shared,
539            config,
540        )?;
541        let worker = WorkerContext::new::<()>("shared-trait-worker");
542        let middleware_name = std::any::type_name_of_val(&storage.middleware()).to_owned();
543        let stream_name = std::any::type_name_of_val(&storage.poll_compact(&worker)).to_owned();
544        Ok((middleware_name, stream_name))
545    }
546
547    fn registration_debug_and_drop() -> (String, bool) {
548        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
549        let (sender, _receiver) = mpsc::channel(1);
550        let id = Ulid::new();
551        registry
552            .lock()
553            .expect("fresh shared registry is not poisoned")
554            .insert("shared-registration".to_owned(), vec![(id, sender)]);
555
556        let debug = {
557            let registration = SharedRegistration {
558                id,
559                queue: "shared-registration".to_owned(),
560                registry: registry.clone(),
561                pool: unchecked_pool(),
562            };
563            format!("{registration:?}")
564        };
565
566        let removed = registry
567            .lock()
568            .expect("fresh shared registry is not poisoned")
569            .is_empty();
570        (debug, removed)
571    }
572
573    /// Build a registry that contains `target_queue` plus optional sibling
574    /// queues, then drop a `SharedRegistration` that points at `target_queue`.
575    /// Returns the number of entries left in the registry after the drop —
576    /// zero when the dropped registration was the last one (the empty-branch
577    /// triggers the best-effort NOTIFY wake-up), positive when siblings remain.
578    fn drop_leaves_remaining(target_queue: &str, sibling_queues: &[&str]) -> usize {
579        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
580        let target_id = Ulid::new();
581        {
582            let mut reg = registry
583                .lock()
584                .expect("fresh shared registry is not poisoned");
585            let (sender, _r) = mpsc::channel(1);
586            reg.insert(target_queue.to_owned(), vec![(target_id, sender)]);
587            for sibling in sibling_queues {
588                let (sender, _r) = mpsc::channel(1);
589                reg.insert((*sibling).to_owned(), vec![(Ulid::new(), sender)]);
590            }
591        }
592
593        {
594            let registration = SharedRegistration {
595                id: target_id,
596                queue: target_queue.to_owned(),
597                registry: registry.clone(),
598                pool: unchecked_pool(),
599            };
600            drop(registration);
601        }
602
603        registry
604            .lock()
605            .expect("fresh shared registry is not poisoned")
606            .len()
607    }
608
609    fn drop_when_registry_empties() -> usize {
610        drop_leaves_remaining("shared-only", &[])
611    }
612
613    fn drop_when_registry_has_siblings() -> usize {
614        drop_leaves_remaining("shared-target", &["shared-other-a", "shared-other-b"])
615    }
616
617    /// Identity-strengthened sibling-drop check: a bare count cannot tell a
618    /// correct drop (the target's entry removed, both siblings kept) from a
619    /// broken one that removed the wrong key but landed on the same total.
620    /// Build the same target + two siblings, drop the `shared-target`
621    /// registration, and return `(remaining_len,
622    /// only_the_target_was_removed)` where the bool asserts the *target* key is
623    /// gone while *both* named siblings remain. Mirrors
624    /// `broadcast_notify_error_observation`'s identity assertion.
625    fn drop_siblings_keeps_their_identities() -> (usize, bool) {
626        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
627        let target_id = Ulid::new();
628        {
629            let mut reg = registry
630                .lock()
631                .expect("fresh shared registry is not poisoned");
632            let (sender, _r) = mpsc::channel(1);
633            reg.insert("shared-target".to_owned(), vec![(target_id, sender)]);
634            for sibling in ["shared-other-a", "shared-other-b"] {
635                let (sender, _r) = mpsc::channel(1);
636                reg.insert(sibling.to_owned(), vec![(Ulid::new(), sender)]);
637            }
638        }
639
640        drop(SharedRegistration {
641            id: target_id,
642            queue: "shared-target".to_owned(),
643            registry: registry.clone(),
644            pool: unchecked_pool(),
645        });
646
647        let reg = registry
648            .lock()
649            .expect("fresh shared registry is not poisoned");
650        let only_target_removed = !reg.contains_key("shared-target")
651            && reg.contains_key("shared-other-a")
652            && reg.contains_key("shared-other-b");
653        (reg.len(), only_target_removed)
654    }
655
656    /// Drop of one registration on a queue with two consumers must leave the
657    /// other sender intact. Regression test for the bug where
658    /// `registry.remove(&queue)` wiped the whole entry, severing the second
659    /// consumer's notify stream.
660    fn drop_one_of_two_keeps_sibling_sender() -> usize {
661        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
662        let queue = "shared-coexist".to_owned();
663        let first_id = Ulid::new();
664        let second_id = Ulid::new();
665        let (first_sender, _first_rx) = mpsc::channel(1);
666        let (second_sender, _second_rx) = mpsc::channel(1);
667        registry
668            .lock()
669            .expect("fresh registry is not poisoned")
670            .insert(
671                queue.clone(),
672                vec![(first_id, first_sender), (second_id, second_sender)],
673            );
674
675        drop(SharedRegistration {
676            id: first_id,
677            queue: queue.clone(),
678            registry: registry.clone(),
679            pool: unchecked_pool(),
680        });
681
682        let guard = registry.lock().expect("registry is not poisoned");
683        guard.get(&queue).map(Vec::len).unwrap_or(0)
684    }
685
686    /// Poison the registry mutex (panic while holding the lock, like
687    /// `make_shared_with_poisoned_registry`), then drop a `SharedRegistration`
688    /// pointing at an existing queue. The `Err(_) => false` arm of `Drop`
689    /// (src/shared.rs:356) must run: it prunes nothing and fires no best-effort
690    /// wake-up NOTIFY, so the queue's sender Vec is left untouched at length 1.
691    /// The helper returning at all proves the drop did not panic on the poison.
692    fn drop_with_poisoned_registry() -> usize {
693        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
694        let queue = "shared-poisoned-drop".to_owned();
695        let id = Ulid::new();
696        let (sender, _receiver) = mpsc::channel(1);
697        registry
698            .lock()
699            .expect("fresh registry is not poisoned")
700            .insert(queue.clone(), vec![(id, sender)]);
701
702        let poison_target = registry.clone();
703        let join = std::thread::spawn(move || {
704            let _guard = poison_target
705                .lock()
706                .expect("fresh registry lock is not poisoned");
707            panic!("synthetic poisoning panic");
708        });
709        let _ = join.join();
710
711        drop(SharedRegistration {
712            id,
713            queue: queue.clone(),
714            registry: registry.clone(),
715            pool: unchecked_pool(),
716        });
717
718        // Recover past the poison to confirm the sender was left in place.
719        registry
720            .lock()
721            .unwrap_or_else(std::sync::PoisonError::into_inner)
722            .get(&queue)
723            .map(Vec::len)
724            .unwrap_or(0)
725    }
726
727    /// Re-registering a namespace that already lives in the registry now
728    /// succeeds: the broadcast redesign allows multiple consumers per queue, so
729    /// the second `make_shared_with_config` must also return `Ok`.
730    fn double_make_shared_same_queue() -> Result<(), SharedPostgresError> {
731        let mut shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
732        let config = Config::new("double-make-shared");
733        let _first = <SharedPostgresStorage as MakeShared<String>>::make_shared_with_config(
734            &mut shared,
735            config.clone(),
736        )?;
737        let _second = <SharedPostgresStorage as MakeShared<String>>::make_shared_with_config(
738            &mut shared,
739            config,
740        )?;
741        Ok(())
742    }
743
744    /// `broadcast_notify_error` walks the registry, delivers the error to each
745    /// sender, then keeps only the queues whose sender vec is still non-empty.
746    /// The returned tuple is `(retained_after_broadcast,
747    /// alive_is_the_sole_survivor)`: the count alone cannot catch an inverted
748    /// keep-predicate (it would retain the wrong queue but the same total), so
749    /// the bool pins that the *live* queue — not the pruned-empty "dead" one —
750    /// is what survives. No listener thread is touched.
751    fn broadcast_notify_error_observation() -> (usize, bool) {
752        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
753        let (alive_sender, _alive_receiver) = mpsc::channel(1);
754        let (dead_sender, dead_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
755        drop(dead_receiver);
756        {
757            let mut reg = registry.lock().expect("fresh registry is not poisoned");
758            reg.insert("alive".to_owned(), vec![(Ulid::new(), alive_sender)]);
759            reg.insert("dead".to_owned(), vec![(Ulid::new(), dead_sender)]);
760        }
761
762        broadcast_notify_error(&registry, "synthetic listener failure".to_owned());
763        let reg = registry.lock().expect("registry is not poisoned");
764        let retained = reg.len();
765        let alive_is_the_sole_survivor = reg.contains_key("alive") && !reg.contains_key("dead");
766        (retained, alive_is_the_sole_survivor)
767    }
768
769    fn new_task_id() -> PgTaskId {
770        PgTaskId::new(Ulid::new())
771    }
772
773    /// `deliver_to_queue` must prune a sender whose receiver was dropped
774    /// (disconnected) and, being the queue's last consumer, remove the now-empty
775    /// queue so the listener's empty-registry exit can fire. Returns the registry
776    /// length after delivery — 0 once the dead sender and its queue are gone.
777    fn deliver_prunes_disconnected_sender() -> usize {
778        let mut registry: RegistryMap = HashMap::new();
779        let (dead_sender, dead_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
780        drop(dead_receiver);
781        registry.insert(
782            "shared-deliver-dead".to_owned(),
783            vec![(Ulid::new(), dead_sender)],
784        );
785        deliver_to_queue(&mut registry, "shared-deliver-dead", &[new_task_id()]);
786        registry.len()
787    }
788
789    /// `deliver_to_queue` must KEEP a sender whose channel is full (transient
790    /// back-pressure) — only a disconnected sender is pruned, because the job is
791    /// durable and the poll fetcher will pick it up. The receiver is held open so
792    /// `try_send` reports `Full`, not `Disconnected`; more ids than the channel
793    /// can ever buffer are delivered so at least one send hits the full path.
794    /// Returns the number of senders still registered on the queue.
795    fn deliver_keeps_full_sender() -> usize {
796        let mut registry: RegistryMap = HashMap::new();
797        let (full_sender, _full_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
798        registry.insert(
799            "shared-deliver-full".to_owned(),
800            vec![(Ulid::new(), full_sender)],
801        );
802        let ids = [new_task_id(), new_task_id(), new_task_id(), new_task_id()];
803        deliver_to_queue(&mut registry, "shared-deliver-full", &ids);
804        registry
805            .get("shared-deliver-full")
806            .map(Vec::len)
807            .unwrap_or(0)
808    }
809
810    /// `deliver_to_queue`'s central contract per its doc comment — broadcasting
811    /// one id to *every* sender bound to the queue — fanned out to MORE than one
812    /// live consumer. Both receivers are held open so neither is pruned; the
813    /// single id is broadcast and each receiver is drained, asserting it carries
814    /// exactly that id. Returns `(first_got_the_id, second_got_the_id)` — `(true,
815    /// true)` when both senders received the wake-up. A delivery that targeted
816    /// only one (e.g. the last) sender would surface here as a `false`.
817    fn deliver_broadcasts_id_to_every_live_sender() -> (bool, bool) {
818        let mut registry: RegistryMap = HashMap::new();
819        let (first_sender, mut first_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
820        let (second_sender, mut second_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
821        registry.insert(
822            "shared-deliver-fanout".to_owned(),
823            vec![(Ulid::new(), first_sender), (Ulid::new(), second_sender)],
824        );
825        let id = new_task_id();
826        deliver_to_queue(&mut registry, "shared-deliver-fanout", &[id]);
827
828        let got_id = |receiver: &mut Receiver<Result<PgTaskId, Error>>| matches!(receiver.try_recv(), Ok(Ok(got)) if got == id);
829        (got_id(&mut first_receiver), got_id(&mut second_receiver))
830    }
831
832    /// `deliver_to_queue` aimed at a queue that has NO registered consumers (a
833    /// real runtime state: a NOTIFY arrives for a job_type no one is polling) is
834    /// a no-op — the `get_mut(queue)` None arm. It must not panic, must not
835    /// create an entry for the absent queue, and must leave the one unrelated
836    /// queue's live sender untouched. Returns `(unrelated_queue_sender_count,
837    /// absent_queue_was_created)`.
838    fn deliver_to_absent_queue_leaves_others_intact() -> (usize, bool) {
839        let mut registry: RegistryMap = HashMap::new();
840        let (sender, _receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
841        registry.insert("shared-other".to_owned(), vec![(Ulid::new(), sender)]);
842        deliver_to_queue(&mut registry, "shared-absent", &[new_task_id()]);
843        let unrelated_len = registry.get("shared-other").map(Vec::len).unwrap_or(0);
844        let absent_created = registry.contains_key("shared-absent");
845        (unrelated_len, absent_created)
846    }
847
848    /// `broadcast_notify_error_locked` must KEEP a sender whose channel is full
849    /// but still connected — only a disconnected sender is pruned. The receiver
850    /// is held open and the single-slot buffer is pre-saturated so the broadcast's
851    /// one `try_send` reports `Full`, not `Disconnected`. Returns the number of
852    /// senders still registered on the queue — 1 when the back-pressured sender
853    /// is retained. Mirrors `deliver_keeps_full_sender`, the symmetric branch on
854    /// `deliver_to_queue`.
855    fn broadcast_notify_error_keeps_full_sender() -> usize {
856        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
857        let (mut full_sender, _full_receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
858        // Saturate the single-slot channel so the broadcast's `try_send` hits
859        // the `Full` (kept) path rather than the empty `Ok` path.
860        while full_sender.try_send(Ok(new_task_id())).is_ok() {}
861        registry
862            .lock()
863            .expect("fresh registry is not poisoned")
864            .insert(
865                "shared-error-full".to_owned(),
866                vec![(Ulid::new(), full_sender)],
867            );
868        broadcast_notify_error(&registry, "synthetic listener failure".to_owned());
869        let reg = registry.lock().expect("registry is not poisoned");
870        reg.get("shared-error-full").map(Vec::len).unwrap_or(0)
871    }
872
873    fn empty_registry() -> RegistryMap {
874        HashMap::new()
875    }
876
877    fn registry_with_one_consumer() -> RegistryMap {
878        let mut registry: RegistryMap = HashMap::new();
879        let (sender, _receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
880        registry.insert(
881            "shared-still-active".to_owned(),
882            vec![(Ulid::new(), sender)],
883        );
884        registry
885    }
886
887    /// `exit_listener` must both flip `listener_alive` to false and broadcast the
888    /// failure to every registered sender. Returns `(alive_after, error_delivered)`
889    /// — `(false, true)` when both effects happen; the live receiver is kept so the
890    /// broadcast lands.
891    fn exit_listener_observation() -> (bool, bool) {
892        let registry: SharedRegistry = Arc::new(Mutex::new(HashMap::new()));
893        let (sender, mut receiver) = mpsc::channel::<Result<PgTaskId, Error>>(1);
894        registry
895            .lock()
896            .expect("fresh registry is not poisoned")
897            .insert("shared-exit".to_owned(), vec![(Ulid::new(), sender)]);
898        let listener_alive = AtomicBool::new(true);
899        exit_listener(
900            &registry,
901            &listener_alive,
902            Some("synthetic listener spawn failure".to_owned()),
903        );
904        let alive_after = listener_alive.load(Ordering::Acquire);
905        let error_delivered = matches!(receiver.try_recv(), Ok(Err(_)));
906        (alive_after, error_delivered)
907    }
908
909    /// First claim on a fresh flag must spawn (false→true edge), the second must
910    /// not. Returns `(first_claim, second_claim)` — `(true, false)` when exactly
911    /// the first registration owns the spawn.
912    fn listener_spawn_claims() -> (bool, bool) {
913        let listener_alive = AtomicBool::new(false);
914        let first = claim_listener_spawn(&listener_alive);
915        let second = claim_listener_spawn(&listener_alive);
916        (first, second)
917    }
918
919    // Q6-rest removed `Arc<Mutex<Receiver>>`: each fetcher owns its receiver
920    // directly. Poisoned-mutex and locked-receiver paths from the previous
921    // architecture no longer exist; their dedicated tests have been removed.
922
923    fn debug_mentions_type(expected: &'static str) -> impl Fn(&String) -> AssertionResult {
924        move |debug| {
925            if debug.contains(expected) {
926                Ok(())
927            } else {
928                Err(AssertionError::new(vec![format!(
929                    "expected debug output containing {expected:?}, got {debug}"
930                )]))
931            }
932        }
933    }
934
935    fn uses_default_queue(result: &SharedObservation) -> AssertionResult {
936        if result.queue == std::any::type_name::<String>()
937            && result.buffer_size == 10
938            && result.debug.contains("SharedFetcher")
939        {
940            Ok(())
941        } else {
942            Err(AssertionError::new(vec![format!(
943                "unexpected default shared storage: queue={:?}, buffer={}, debug={}",
944                result.queue, result.buffer_size, result.debug
945            )]))
946        }
947    }
948
949    fn uses_configured_queue(result: &SharedObservation) -> AssertionResult {
950        if result.queue == "shared-unit"
951            && result.buffer_size == 3
952            && result.debug.contains("SharedFetcher")
953        {
954            Ok(())
955        } else {
956            Err(AssertionError::new(vec![format!(
957                "unexpected configured shared storage: queue={:?}, buffer={}, debug={}",
958                result.queue, result.buffer_size, result.debug
959            )]))
960        }
961    }
962
963    fn constructs_backend_traits(result: &(String, String)) -> AssertionResult {
964        if result.0.contains("PgMiddleware") && result.1.contains("Stream") {
965            Ok(())
966        } else {
967            Err(AssertionError::new(vec![format!(
968                "unexpected shared trait surfaces: {result:?}"
969            )]))
970        }
971    }
972
973    fn removes_registration(result: &(String, bool)) -> AssertionResult {
974        if result.0.contains("SharedRegistration") && result.1 {
975            Ok(())
976        } else {
977            Err(AssertionError::new(vec![format!(
978                "expected registration debug and drop cleanup, got {result:?}"
979            )]))
980        }
981    }
982
983    /// Drive `make_shared_with_config` against a deliberately-poisoned
984    /// registry mutex and surface the resulting error variant. The poisoning
985    /// is forced by panicking inside a thread that holds the lock; that is
986    /// the only documented way `make_shared_with_config` can return
987    /// `SharedPostgresError::RegistryLocked` (shared.rs:170-173).
988    fn make_shared_with_poisoned_registry() -> Result<(), SharedPostgresError> {
989        let mut shared: SharedPostgresStorage = SharedPostgresStorage::new(unchecked_pool());
990        let registry = shared.registry.clone();
991        let join = std::thread::spawn(move || {
992            let _guard = registry
993                .lock()
994                .expect("fresh registry lock is not poisoned");
995            panic!("synthetic poisoning panic");
996        });
997        // The poisoning thread panics while holding the lock, leaving the
998        // mutex in PoisonError state for the next caller.
999        let _ = join.join();
1000        let config = Config::new("poisoned-registry");
1001        <SharedPostgresStorage as MakeShared<String>>::make_shared_with_config(&mut shared, config)
1002            .map(|_| ())
1003    }
1004
1005    fn is_registry_locked(error: &SharedPostgresError) -> AssertionResult {
1006        // `RegistryLocked` is currently the only variant, so the match is
1007        // exhaustive without a catch-all; adding a variant later will surface
1008        // here as a compile error, which is the right place to revisit this.
1009        match error {
1010            SharedPostgresError::RegistryLocked => Ok(()),
1011        }
1012    }
1013
1014    lets_expect! {
1015        expect(shared_debug()) {
1016            to describes_the_shared_factory { debug_mentions_type("SharedPostgresStorage") }
1017        }
1018
1019        expect(make_default_shared()) {
1020            when no_config_is_supplied {
1021                to uses_the_task_type_as_the_namespace { be_ok_and uses_default_queue }
1022            }
1023        }
1024
1025        expect(make_configured_shared()) {
1026            when config_is_supplied {
1027                to exposes_the_queue_and_fetcher { be_ok_and uses_configured_queue }
1028            }
1029        }
1030
1031        expect(shared_trait_surfaces()) {
1032            when backend_traits_are_requested {
1033                to builds_middleware_and_compact_stream { be_ok_and constructs_backend_traits }
1034            }
1035        }
1036
1037        expect(registration_debug_and_drop()) {
1038            when registration_is_dropped {
1039                to removes_the_namespace_from_the_registry { removes_registration }
1040            }
1041        }
1042
1043        expect(drop_when_registry_empties()) {
1044            when dropping_the_last_registration_empties_the_registry {
1045                to leaves_no_remaining_registrations { equal(0) }
1046            }
1047        }
1048
1049        expect(drop_when_registry_has_siblings()) {
1050            when dropping_one_of_several_registrations {
1051                to keeps_sibling_registrations_intact { equal(2) }
1052            }
1053        }
1054
1055        expect(drop_siblings_keeps_their_identities()) {
1056            when dropping_one_of_several_registrations_with_siblings_present {
1057                to removes_only_the_dropped_queue_and_keeps_both_named_siblings {
1058                    equal((2_usize, true))
1059                }
1060            }
1061        }
1062
1063        expect(drop_one_of_two_keeps_sibling_sender()) {
1064            when dropping_one_of_two_consumers_on_the_same_queue {
1065                to leaves_the_other_senders_sender_in_place { equal(1) }
1066            }
1067        }
1068
1069        expect(drop_with_poisoned_registry()) {
1070            when the_registry_mutex_is_poisoned {
1071                to leaves_the_registration_in_place_without_waking_the_listener {
1072                    equal(1)
1073                }
1074            }
1075        }
1076
1077        expect(double_make_shared_same_queue()) {
1078            when the_same_queue_is_registered_twice {
1079                // Q6-rest broadcast redesign: multiple consumers per queue
1080                // are now allowed (a second registration used to be rejected).
1081                // The listener broadcasts each event to every registered sender.
1082                to accepts_the_second_registration { be_ok }
1083            }
1084        }
1085
1086        expect(broadcast_notify_error_observation()) {
1087            when listener_broadcasts_an_error_to_a_mixed_registry {
1088                to keeps_the_live_queue_and_drops_the_disconnected_one { equal((1_usize, true)) }
1089            }
1090        }
1091
1092        expect(broadcast_notify_error_keeps_full_sender()) {
1093            when a_senders_channel_is_full_but_still_connected {
1094                to keeps_the_back_pressured_sender_registered { equal(1) }
1095            }
1096        }
1097
1098        expect(deliver_prunes_disconnected_sender()) {
1099            when a_queues_only_sender_has_a_dropped_receiver {
1100                to prunes_the_disconnected_sender_and_removes_the_empty_queue { equal(0) }
1101            }
1102        }
1103
1104        expect(deliver_keeps_full_sender()) {
1105            when a_queues_sender_channel_is_full_but_still_connected {
1106                to keeps_the_back_pressured_sender_registered { equal(1) }
1107            }
1108        }
1109
1110        expect(deliver_broadcasts_id_to_every_live_sender()) {
1111            when a_queue_has_multiple_live_consumers {
1112                to broadcasts_the_wakeup_id_to_every_sender { equal((true, true)) }
1113            }
1114        }
1115
1116        expect(deliver_to_absent_queue_leaves_others_intact()) {
1117            when a_notification_targets_a_queue_with_no_registered_consumers {
1118                to leaves_other_queues_untouched_and_creates_no_entry {
1119                    equal((1_usize, false))
1120                }
1121            }
1122        }
1123
1124        expect(listener_should_exit(&registry)) {
1125            let registry = empty_registry();
1126
1127            when no_consumers_remain_registered {
1128                to signals_the_listener_to_exit { be_true }
1129            }
1130
1131            when a_consumer_is_still_registered {
1132                let registry = registry_with_one_consumer();
1133                to keeps_the_listener_running { be_false }
1134            }
1135        }
1136
1137        expect(exit_listener_observation()) {
1138            when the_listener_exits_after_a_spawn_or_connection_failure {
1139                to clears_the_alive_flag_and_broadcasts_the_error { equal((false, true)) }
1140            }
1141        }
1142
1143        expect(listener_spawn_claims()) {
1144            when two_registrations_race_to_claim_the_listener_spawn {
1145                to spawns_on_the_first_registration_only { equal((true, false)) }
1146            }
1147        }
1148
1149        // Q6-rest: removed `locked_fetcher_poll` / `poisoned_fetcher_poll`
1150        // assertions — the `Arc<Mutex<Receiver>>` they exercised no longer
1151        // exists. Each fetcher owns its receiver directly after the broadcast
1152        // redesign.
1153
1154        expect(make_shared_with_poisoned_registry()) {
1155            when the_registry_mutex_is_poisoned_by_a_panic_in_another_thread {
1156                // Sibling to "the_same_queue_is_registered_twice" — covers
1157                // the other failure mode of make_shared_with_config: the
1158                // mutex lock itself is unrecoverable rather than the queue
1159                // being already taken.
1160                to surfaces_registry_locked_rather_than_panicking_or_succeeding {
1161                    be_err_and is_registry_locked
1162                }
1163            }
1164        }
1165    }
1166}