Skip to main content

nautilus_common/msgbus/
stubs.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
16use std::{
17    any::Any,
18    cell::RefCell,
19    fmt::Debug,
20    rc::Rc,
21    sync::{
22        Arc,
23        atomic::{AtomicBool, Ordering},
24    },
25};
26
27use ahash::AHashMap;
28use nautilus_core::UUID4;
29use ustr::Ustr;
30
31use crate::msgbus::{
32    Handler, IntoHandler, ShareableMessageHandler, TypedHandler, TypedIntoHandler,
33    typed_handler::shareable_handler,
34};
35
36/// Minimal message payload received by [`StubMessageHandler`].
37#[derive(Debug, Clone)]
38pub struct StubMessage;
39
40/// Stub handler which logs messages it receives.
41#[derive(Clone)]
42pub struct StubMessageHandler {
43    id: Ustr,
44    callback: Arc<dyn Fn(StubMessage) + Send>,
45}
46
47impl Debug for StubMessageHandler {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct(stringify!(StubMessageHandler))
50            .field("id", &self.id)
51            .finish()
52    }
53}
54
55impl Handler<dyn Any> for StubMessageHandler {
56    fn id(&self) -> Ustr {
57        self.id
58    }
59
60    fn handle(&self, message: &dyn Any) {
61        if let Some(msg) = message.downcast_ref::<StubMessage>() {
62            (self.callback)(msg.clone());
63        }
64    }
65}
66
67#[must_use]
68#[allow(unused_must_use)]
69pub fn get_stub_shareable_handler(id: Option<Ustr>) -> ShareableMessageHandler {
70    let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
71    shareable_handler(Rc::new(StubMessageHandler {
72        id: unique_id,
73        callback: Arc::new(|m: StubMessage| {
74            format!("{m:?}");
75        }),
76    }))
77}
78
79/// Handler that tracks whether it has been called.
80#[derive(Debug, Clone)]
81pub struct CallCheckHandler {
82    id: Ustr,
83    called: Arc<AtomicBool>,
84}
85
86impl CallCheckHandler {
87    #[must_use]
88    pub fn new(id: Option<Ustr>) -> Self {
89        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
90        Self {
91            id: unique_id,
92            called: Arc::new(AtomicBool::new(false)),
93        }
94    }
95
96    #[must_use]
97    pub fn was_called(&self) -> bool {
98        self.called.load(Ordering::SeqCst)
99    }
100
101    /// Returns a `ShareableMessageHandler` for registration.
102    #[must_use]
103    pub fn handler(&self) -> ShareableMessageHandler {
104        shareable_handler(Rc::new(self.clone()))
105    }
106}
107
108impl Handler<dyn Any> for CallCheckHandler {
109    fn id(&self) -> Ustr {
110        self.id
111    }
112
113    fn handle(&self, _message: &dyn Any) {
114        self.called.store(true, Ordering::SeqCst);
115    }
116}
117
118/// Creates a call-checking handler and returns both the handler for registration
119/// and a clone that can be used to check if it was called.
120#[must_use]
121pub fn get_call_check_handler(id: Option<Ustr>) -> (ShareableMessageHandler, CallCheckHandler) {
122    let checker = CallCheckHandler::new(id);
123    let handler = checker.handler();
124    (handler, checker)
125}
126
127/// Handler that saves messages it receives (for Any-based routing).
128#[derive(Debug, Clone)]
129pub struct AnySavingHandler<T> {
130    id: Ustr,
131    messages: Rc<RefCell<Vec<T>>>,
132}
133
134impl<T: Clone + 'static> AnySavingHandler<T> {
135    #[must_use]
136    pub fn new(id: Option<Ustr>) -> Self {
137        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
138        Self {
139            id: unique_id,
140            messages: Rc::new(RefCell::new(Vec::new())),
141        }
142    }
143
144    #[must_use]
145    pub fn get_messages(&self) -> Vec<T> {
146        self.messages.borrow().clone()
147    }
148
149    pub fn clear(&self) {
150        self.messages.borrow_mut().clear();
151    }
152
153    /// Returns a `ShareableMessageHandler` for registration.
154    #[must_use]
155    pub fn handler(&self) -> ShareableMessageHandler {
156        shareable_handler(Rc::new(self.clone()))
157    }
158}
159
160impl<T: Clone + 'static> Handler<dyn Any> for AnySavingHandler<T> {
161    fn id(&self) -> Ustr {
162        self.id
163    }
164
165    fn handle(&self, message: &dyn Any) {
166        if let Some(m) = message.downcast_ref::<T>() {
167            self.messages.borrow_mut().push(m.clone());
168        } else {
169            log::error!(
170                "AnySavingHandler: expected {} got {:?}",
171                std::any::type_name::<T>(),
172                message.type_id()
173            );
174        }
175    }
176}
177
178/// Creates an Any-based message saving handler and returns both the handler
179/// for registration and a clone that can be used to retrieve messages.
180#[must_use]
181pub fn get_any_saving_handler<T: Clone + 'static>(
182    id: Option<Ustr>,
183) -> (ShareableMessageHandler, AnySavingHandler<T>) {
184    let saver = AnySavingHandler::new(id);
185    let handler = saver.handler();
186    (handler, saver)
187}
188
189// Type alias for backward compatibility
190pub type MessageSavingHandler<T> = AnySavingHandler<T>;
191
192/// Typed handler which saves the messages it receives (no downcast needed).
193#[derive(Debug, Clone)]
194pub struct TypedMessageSavingHandler<T> {
195    id: Ustr,
196    messages: Rc<RefCell<Vec<T>>>,
197}
198
199impl<T: Clone + 'static> TypedMessageSavingHandler<T> {
200    #[must_use]
201    pub fn new(id: Option<Ustr>) -> Self {
202        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
203        Self {
204            id: unique_id,
205            messages: Rc::new(RefCell::new(Vec::new())),
206        }
207    }
208
209    #[must_use]
210    pub fn get_messages(&self) -> Vec<T> {
211        self.messages.borrow().clone()
212    }
213
214    /// Returns a `TypedHandler` that can be used for subscriptions.
215    #[must_use]
216    pub fn handler(&self) -> TypedHandler<T> {
217        TypedHandler::new(self.clone())
218    }
219}
220
221impl<T: Clone + 'static> Handler<T> for TypedMessageSavingHandler<T> {
222    fn id(&self) -> Ustr {
223        self.id
224    }
225
226    fn handle(&self, message: &T) {
227        self.messages.borrow_mut().push(message.clone());
228    }
229}
230
231/// Creates a typed message saving handler and returns both the handler for subscriptions
232/// and a clone that can be used to retrieve messages.
233#[must_use]
234pub fn get_typed_message_saving_handler<T: Clone + 'static>(
235    id: Option<Ustr>,
236) -> (TypedHandler<T>, TypedMessageSavingHandler<T>) {
237    let saving_handler = TypedMessageSavingHandler::new(id);
238    let typed_handler = saving_handler.handler();
239    (typed_handler, saving_handler)
240}
241
242/// Ownership-based typed handler which saves the messages it receives.
243///
244/// Unlike [`TypedMessageSavingHandler`] which borrows messages, this handler
245/// takes ownership which is required for `IntoEndpointMap` endpoints.
246#[derive(Debug, Clone)]
247pub struct TypedIntoMessageSavingHandler<T> {
248    id: Ustr,
249    messages: Rc<RefCell<Vec<T>>>,
250}
251
252impl<T: 'static> TypedIntoMessageSavingHandler<T> {
253    #[must_use]
254    pub fn new(id: Option<Ustr>) -> Self {
255        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
256        Self {
257            id: unique_id,
258            messages: Rc::new(RefCell::new(Vec::new())),
259        }
260    }
261
262    /// Creates a handler backed by an existing shared messages vec.
263    #[must_use]
264    pub fn new_with_messages(id: Option<Ustr>, messages: Rc<RefCell<Vec<T>>>) -> Self {
265        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
266        Self {
267            id: unique_id,
268            messages,
269        }
270    }
271
272    #[must_use]
273    pub fn get_messages(&self) -> Vec<T>
274    where
275        T: Clone,
276    {
277        self.messages.borrow().clone()
278    }
279
280    /// Returns a `TypedIntoHandler` that can be used for endpoint registration.
281    #[must_use]
282    pub fn handler(&self) -> TypedIntoHandler<T> {
283        TypedIntoHandler::new(Self {
284            id: self.id,
285            messages: self.messages.clone(),
286        })
287    }
288
289    pub fn clear(&self) {
290        self.messages.borrow_mut().clear();
291    }
292}
293
294impl<T: 'static> IntoHandler<T> for TypedIntoMessageSavingHandler<T> {
295    fn id(&self) -> Ustr {
296        self.id
297    }
298
299    fn handle(&self, message: T) {
300        self.messages.borrow_mut().push(message);
301    }
302}
303
304/// Creates an ownership-based typed message saving handler and returns both the handler
305/// for endpoint registration and a clone that can be used to retrieve messages.
306#[must_use]
307pub fn get_typed_into_message_saving_handler<T: 'static>(
308    id: Option<Ustr>,
309) -> (TypedIntoHandler<T>, TypedIntoMessageSavingHandler<T>) {
310    let saving_handler = TypedIntoMessageSavingHandler::new(id);
311    let typed_handler = saving_handler.handler();
312    (typed_handler, saving_handler)
313}
314
315// Legacy API for tests that use the old pattern with thread_local storage.
316// These wrap AnySavingHandler in thread_local for simpler test usage.
317
318thread_local! {
319    static SAVING_HANDLERS: RefCell<AHashMap<Ustr, Box<dyn std::any::Any>>> = RefCell::new(AHashMap::new());
320}
321
322/// Creates a message saving handler and stores it for later retrieval.
323#[must_use]
324pub fn get_message_saving_handler<T: Clone + 'static>(id: Option<Ustr>) -> ShareableMessageHandler {
325    let (handler, saver) = get_any_saving_handler::<T>(id);
326    let handler_id = handler.0.id();
327    SAVING_HANDLERS.with(|handlers| {
328        handlers.borrow_mut().insert(handler_id, Box::new(saver));
329    });
330    handler
331}
332
333/// Retrieves saved messages from a handler created by `get_message_saving_handler`.
334#[must_use]
335pub fn get_saved_messages<T: Clone + 'static>(handler: &ShareableMessageHandler) -> Vec<T> {
336    let handler_id = handler.0.id();
337    SAVING_HANDLERS.with(|handlers| {
338        let handlers = handlers.borrow();
339        if let Some(saver) = handlers.get(&handler_id)
340            && let Some(saver) = saver.downcast_ref::<AnySavingHandler<T>>()
341        {
342            return saver.get_messages();
343        }
344        Vec::new()
345    })
346}