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
12type 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
33type RegistryEntry = (Ulid, RegistrySender);
40type RegistryMap = HashMap<String, Vec<RegistryEntry>>;
41type SharedRegistry = Arc<Mutex<RegistryMap>>;
42
43pub struct SharedPostgresStorage<Codec = JsonCodec<CompactType>> {
50 pool: PgPool,
51 registry: SharedRegistry,
52 listener_alive: Arc<AtomicBool>,
61 _marker: PhantomData<Codec>,
62}
63
64impl<Codec> SharedPostgresStorage<Codec> {
65 #[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 ®istry,
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 ®istry,
102 &listener_alive,
103 Some(format!("failed to start shared LISTEN listener: {error}")),
104 );
105 return;
106 }
107 run_listener_loop(&mut conn, ®istry, &listener_alive);
108 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
127fn 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>(¬ification.payload)
153 else {
154 continue;
155 };
156 let (event_queue, ids) = event.into_ids();
157 let Ok(mut registry) = registry.lock() else {
158 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 if listener_should_exit(®istry) {
180 listener_alive.store(false, Ordering::Release);
181 drop(registry);
182 return;
183 }
184 }
185 Err(_) => {
186 listener_alive.store(false, Ordering::Release);
188 return;
189 }
190 }
191 std::thread::sleep(queries::NOTIFY_LISTENER_POLL_INTERVAL);
192 }
193}
194
195fn 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 listener_alive.store(false, Ordering::Release);
216 }
217 }
218}
219
220fn 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 Err(_) => true,
241 });
242 }
243 if senders.is_empty() {
244 registry.remove(queue);
245 }
246 }
247}
248
249fn listener_should_exit(registry: &RegistryMap) -> bool {
253 registry.is_empty()
254}
255
256fn 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#[derive(Debug, thiserror::Error)]
296#[non_exhaustive]
297pub enum SharedPostgresError {
298 #[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 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 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 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 if became_empty {
415 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
430pub 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 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 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 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 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 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 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 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(®istry, "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 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 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 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 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 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 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(®istry, "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 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 ®istry,
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 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 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 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 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 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 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(®istry)) {
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 expect(make_shared_with_poisoned_registry()) {
1155 when the_registry_mutex_is_poisoned_by_a_panic_in_another_thread {
1156 to surfaces_registry_locked_rather_than_panicking_or_succeeding {
1161 be_err_and is_registry_locked
1162 }
1163 }
1164 }
1165 }
1166}