Continue DataActor in Rust

Author: cjdsellers

crates/common/src/actor/data_actor.rs

@@ -50,7 +50,7 @@ use crate::{
     cache::Cache,
     clock::Clock,
     enums::{ComponentState, ComponentTrigger},
-    logging::{CMD, RECV, SENT},
+    logging::{CMD, RECV, REQ, SEND},
     messages::{
         data::{
             DataCommand, DataResponse, RequestBars, RequestBookSnapshot, RequestCommand,
@@ -731,13 +731,33 @@ impl DataActorCore {
 
     fn send_data_cmd(&self, command: DataCommand) {
         if self.config.log_commands {
-            log::info!("{CMD}{SENT} {command:?}");
+            log::info!("{CMD}{SEND} {command:?}");
         }
 
         let endpoint = MessagingSwitchboard::data_engine_execute();
         msgbus::send(&endpoint, command.as_any())
     }
 
+    fn send_data_req<A: DataActor>(&self, request: RequestCommand) {
+        if self.config.log_commands {
+            log::info!("{REQ}{SEND} {request:?}");
+        }
+
+        let actor_id = self.actor_id.inner();
+        let handler = ShareableMessageHandler(Rc::new(TypedMessageHandler::from(
+            move |response: &DataResponse| {
+                get_actor_unchecked::<A>(&actor_id).handle_data_response(response);
+            },
+        )));
+
+        let msgbus = get_message_bus()
+            .borrow_mut()
+            .register_response_handler(request.request_id(), handler);
+
+        let endpoint = MessagingSwitchboard::data_engine_execute();
+        msgbus::send(&endpoint, request.as_any())
+    }
+
     pub fn start(&mut self) -> anyhow::Result<()> {
         self.transition_state(ComponentTrigger::Start)?; // -> Starting
 

crates/common/src/logging/mod.rs

@@ -39,7 +39,7 @@ use self::{
 use crate::enums::LogLevel;
 
 pub const RECV: &str = "<--";
-pub const SENT: &str = "-->";
+pub const SEND: &str = "-->";
 pub const CMD: &str = "[CMD]";
 pub const EVT: &str = "[EVT]";
 pub const DOC: &str = "[DOC]";

crates/common/src/messages/data.rs

@@ -1008,6 +1008,23 @@ pub enum RequestCommand {
 }
 
 impl RequestCommand {
+    /// Converts the command to a dyn Any trait object for messaging.
+    pub fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    pub fn request_id(&self) -> &UUID4 {
+        match self {
+            Self::Data(cmd) => &cmd.request_id,
+            Self::Instruments(cmd) => &cmd.request_id,
+            Self::Instrument(cmd) => &cmd.request_id,
+            Self::BookSnapshot(cmd) => &cmd.request_id,
+            Self::Quotes(cmd) => &cmd.request_id,
+            Self::Trades(cmd) => &cmd.request_id,
+            Self::Bars(cmd) => &cmd.request_id,
+        }
+    }
+
     pub fn client_id(&self) -> Option<&ClientId> {
         match self {
             Self::Data(cmd) => Some(&cmd.client_id),

crates/common/src/msgbus/handler.rs

@@ -31,6 +31,14 @@ pub trait MessageHandler: Any {
     fn as_any(&self) -> &dyn Any;
 }
 
+impl PartialEq for dyn MessageHandler {
+    fn eq(&self, other: &Self) -> bool {
+        self.id() == other.id()
+    }
+}
+
+impl Eq for dyn MessageHandler {}
+
 #[derive(Debug)]
 pub struct TypedMessageHandler<T: 'static + ?Sized, F: Fn(&T) + 'static> {
     id: Ustr,
@@ -124,6 +132,12 @@ fn generate_deterministic_handler_id<T: 'static + ?Sized, F: 'static + Fn(&T)>(
 #[derive(Clone)]
 pub struct ShareableMessageHandler(pub Rc<dyn MessageHandler>);
 
+impl ShareableMessageHandler {
+    pub fn id(&self) -> Ustr {
+        self.0.id()
+    }
+}
+
 impl Debug for ShareableMessageHandler {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct(stringify!(ShareableMessageHandler))

crates/common/src/msgbus/mod.rs

@@ -39,6 +39,7 @@ use std::{
     sync::OnceLock,
 };
 
+use ahash::AHashMap;
 use handler::ShareableMessageHandler;
 use indexmap::IndexMap;
 use nautilus_core::UUID4;
@@ -86,6 +87,21 @@ pub fn send(endpoint: &Ustr, message: &dyn Any) {
     }
 }
 
+/// Sends the `response` to the handler registered for the `correlation_id` (if found).
+pub fn response(correlation_id: &UUID4, message: &dyn Any) {
+    let handler = get_message_bus()
+        .borrow()
+        .get_response_handler(correlation_id)
+        .cloned();
+    if let Some(handler) = handler {
+        handler.0.handle(message);
+    } else {
+        log::error!(
+            "Failed to handle response: handler not found for correlation_id {correlation_id}"
+        )
+    }
+}
+
 /// Publishes the `message` to the `topic`.
 pub fn publish(topic: &Ustr, message: &dyn Any) {
     log::trace!("Publishing topic '{topic}' {message:?}");
@@ -303,8 +319,10 @@ pub struct MessageBus {
     /// Maps a pattern to all the handlers registered for it
     /// this is updated whenever a new subscription is created.
     patterns: IndexMap<Ustr, Vec<Subscription>>,
-    /// Handles a message or a request destined for a specific endpoint.
+    /// Index of endpoint addresses and their handlers.
     endpoints: IndexMap<Ustr, ShareableMessageHandler>,
+    /// Index of request correlation IDs and their response handlers.
+    correlation_index: AHashMap<UUID4, ShareableMessageHandler>,
 }
 
 // SAFETY: Message bus is not meant to be passed between threads
@@ -328,6 +346,7 @@ impl MessageBus {
             subscriptions: IndexMap::new(),
             patterns: IndexMap::new(),
             endpoints: IndexMap::new(),
+            correlation_index: AHashMap::new(),
             has_backing: false,
         }
     }
@@ -401,12 +420,19 @@ impl MessageBus {
         // TODO: Integrate the backing database
         Ok(())
     }
+
     /// Returns the handler for the given `endpoint`.
     #[must_use]
     pub fn get_endpoint<T: AsRef<str>>(&self, endpoint: T) -> Option<&ShareableMessageHandler> {
         self.endpoints.get(&Ustr::from(endpoint.as_ref()))
     }
 
+    /// Returns the handler for the given `correlation_id`.
+    #[must_use]
+    pub fn get_response_handler(&self, correlation_id: &UUID4) -> Option<&ShareableMessageHandler> {
+        self.correlation_index.get(correlation_id)
+    }
+
     #[must_use]
     pub fn matching_subscriptions(&self, pattern: &Ustr) -> Vec<Subscription> {
         let mut matching_subs: Vec<Subscription> = Vec::new();
@@ -433,6 +459,22 @@ impl MessageBus {
         matching_subs
     }
 
+    pub fn register_response_handler(
+        &mut self,
+        correlation_id: &UUID4,
+        handler: ShareableMessageHandler,
+    ) -> anyhow::Result<()> {
+        if self.correlation_index.contains_key(correlation_id) {
+            return Err(anyhow::anyhow!(
+                "Correlation ID <{correlation_id}> already has a registered handler",
+            ));
+        }
+
+        self.correlation_index.insert(*correlation_id, handler);
+
+        Ok(())
+    }
+
     fn matching_handlers<'a>(
         &'a self,
         pattern: &'a Ustr,

crates/common/src/msgbus/tests.rs

@@ -55,6 +55,44 @@ fn test_is_subscribed_when_no_subscriptions() {
     assert!(!msgbus.borrow().is_subscribed("my-topic", handler));
 }
 
+#[rstest]
+fn test_get_response_handler_when_no_handler() {
+    let msgbus = get_message_bus();
+    let msgbus_ref = msgbus.borrow();
+    let handler = msgbus_ref.get_response_handler(&UUID4::new());
+    assert!(handler.is_none());
+}
+
+#[rstest]
+fn test_get_response_handler_when_already_registered() {
+    let msgbus = get_message_bus();
+    let mut msgbus_ref = msgbus.borrow_mut();
+    let handler = get_stub_shareable_handler(None);
+
+    let request_id = UUID4::new();
+    msgbus_ref
+        .register_response_handler(&request_id, handler.clone())
+        .unwrap();
+
+    let result = msgbus_ref.register_response_handler(&request_id, handler.clone());
+    assert!(result.is_err());
+}
+
+#[rstest]
+fn test_get_response_handler_when_registered() {
+    let msgbus = get_message_bus();
+    let mut msgbus_ref = msgbus.borrow_mut();
+    let handler = get_stub_shareable_handler(None);
+
+    let request_id = UUID4::new();
+    msgbus_ref
+        .register_response_handler(&request_id, handler.clone())
+        .unwrap();
+
+    let handler = msgbus_ref.get_response_handler(&request_id).unwrap();
+    assert_eq!(handler.id(), handler.id());
+}
+
 #[rstest]
 fn test_is_registered_when_no_registrations() {
     let msgbus = get_message_bus();

crates/execution/src/order_manager/manager.rs

@@ -18,7 +18,7 @@ use std::{cell::RefCell, collections::HashMap, rc::Rc};
 use nautilus_common::{
     cache::Cache,
     clock::Clock,
-    logging::{CMD, EVT, SENT},
+    logging::{CMD, EVT, SEND},
     msgbus::{self},
 };
 use nautilus_core::UUID4;
@@ -515,13 +515,13 @@ impl OrderManager {
 
     // Message sending methods
     pub fn send_emulator_command(&self, command: TradingCommand) {
-        log::info!("{CMD}{SENT} {command}");
+        log::info!("{CMD}{SEND} {command}");
 
         msgbus::send(&Ustr::from("OrderEmulator.execute"), &command);
     }
 
     pub fn send_algo_command(&self, command: SubmitOrder, exec_algorithm_id: ExecAlgorithmId) {
-        log::info!("{CMD}{SENT} {command}");
+        log::info!("{CMD}{SEND} {command}");
 
         let endpoint = format!("{exec_algorithm_id}.execute");
         msgbus::send(
@@ -531,22 +531,22 @@ impl OrderManager {
     }
 
     pub fn send_risk_command(&self, command: TradingCommand) {
-        log::info!("{CMD}{SENT} {command}");
+        log::info!("{CMD}{SEND} {command}");
         msgbus::send(&Ustr::from("RiskEngine.execute"), &command);
     }
 
     pub fn send_exec_command(&self, command: TradingCommand) {
-        log::info!("{CMD}{SENT} {command}");
+        log::info!("{CMD}{SEND} {command}");
         msgbus::send(&Ustr::from("ExecEngine.execute"), &command);
     }
 
     pub fn send_risk_event(&self, event: OrderEventAny) {
-        log::info!("{EVT}{SENT} {event}");
+        log::info!("{EVT}{SEND} {event}");
         msgbus::send(&Ustr::from("RiskEngine.process"), &event);
     }
 
     pub fn send_exec_event(&self, event: OrderEventAny) {
-        log::info!("{EVT}{SENT} {event}");
+        log::info!("{EVT}{SEND} {event}");
         msgbus::send(&Ustr::from("ExecEngine.process"), &event);
     }
 }