nautilus_common/msgbus/
stubs.rs1use 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#[derive(Debug, Clone)]
38pub struct StubMessage;
39
40#[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#[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 #[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#[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#[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 #[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#[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
189pub type MessageSavingHandler<T> = AnySavingHandler<T>;
191
192#[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 #[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#[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#[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 #[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 #[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#[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
315thread_local! {
319 static SAVING_HANDLERS: RefCell<AHashMap<Ustr, Box<dyn std::any::Any>>> = RefCell::new(AHashMap::new());
320}
321
322#[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#[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}