nautilus_testkit/
connection.rs1use std::{
19 future::{Future, poll_fn},
20 pin::Pin,
21 sync::Arc,
22 task::{Context, Poll, Wake, Waker},
23 time::Duration,
24};
25
26use parking_lot::Mutex;
27
28pub async fn race_command_with_publication<C, F, R>(
41 sender: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<C>>>,
42 connect: F,
43 command: impl FnOnce() -> C + Send + 'static,
44 active: impl Fn() -> bool,
45) -> R
46where
47 C: Send + 'static,
48 F: Future<Output = R>,
49{
50 let publication = sender.read().await;
51 let mut connect = Box::pin(connect);
52 tokio::time::timeout(Duration::from_secs(5), async {
53 loop {
54 let polled = poll_fn(|cx| Poll::Ready(connect.as_mut().poll(cx))).await;
55 assert!(polled.is_pending(), "connect finished before publication");
56
57 if sender.try_read().is_err() {
58 break;
59 }
60
61 tokio::task::yield_now().await;
62 }
63 })
64 .await
65 .expect("connect must wait to publish its sender");
66
67 let active_before_publication = active();
68 let command_sender = Arc::clone(&sender);
69 let command = PublicationCommand {
70 future: Mutex::new(Some(Box::pin(async move {
71 let sender = command_sender.read().await;
72 assert!(sender.send(command()).is_ok());
73 }))),
74 };
75
76 let command = Arc::new(command);
77 let waker = Waker::from(Arc::clone(&command));
78 assert!(command.poll(&waker).is_pending());
79 drop(publication);
80
81 let result = tokio::time::timeout(Duration::from_secs(5), connect)
82 .await
83 .expect("connect must complete after publication");
84 assert!(command.future.lock().is_none(), "command was not sent");
85 assert!(
86 !active_before_publication,
87 "active before sender publication"
88 );
89 result
90}
91
92struct PublicationCommand<F> {
93 future: Mutex<Option<Pin<Box<F>>>>,
94}
95
96impl<F: Future<Output = ()>> PublicationCommand<F> {
97 fn poll(&self, waker: &Waker) -> Poll<()> {
98 let mut future = self.future.lock();
99 let polled = future
100 .as_mut()
101 .expect("publication command polled after completion")
102 .as_mut()
103 .poll(&mut Context::from_waker(waker));
104
105 if polled.is_ready() {
106 *future = None;
107 }
108
109 polled
110 }
111}
112
113impl<F: Future<Output = ()> + Send + 'static> Wake for PublicationCommand<F> {
114 fn wake(self: Arc<Self>) {
115 self.wake_by_ref();
116 }
117
118 fn wake_by_ref(self: &Arc<Self>) {
119 assert!(self.poll(Waker::noop()).is_ready());
120 }
121}