Skip to main content

nautilus_testkit/
connection.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Deterministic command injection at an asynchronous connection publication boundary.
17
18use 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
28/// Connects while injecting a command immediately after the shared sender becomes available.
29///
30/// `command` constructs the command after acquiring the newly published sender.
31/// The sender must use Tokio's fair read/write lock. This function holds a reader until
32/// `connect` queues its publication writer, then queues the command's reader behind it.
33/// The reader's waker sends the command synchronously during publication, before `connect`
34/// can continue past releasing the writer. No scheduler timing or sleep is required.
35///
36/// # Panics
37///
38/// Panics if connection or publication takes more than five seconds, `connect` finishes before
39/// publication, `active` is true before publication, or the command cannot enter the new channel.
40pub 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}