Introduce NATS JetStream as event bus

Cargo.toml

@@ -27,6 +27,7 @@ alloy-transport-ws = { version = "1.8.3", default-features = false }
 anyhow = "1.0.100"
 app-data = { path = "crates/app-data" }
 arc-swap = "1.7.1"
+async-nats = "0.48.0"
 async-stream = "0.3.5"
 async-trait = "0.1.80"
 autopilot = { path = "crates/autopilot" }

crates/autopilot/src/run.rs

@@ -36,7 +36,7 @@ use {
     http_client::HttpClientFactory,
     model::DomainSeparator,
     num::ToPrimitive,
-    observe::metrics::LivenessChecking,
+    observe::{config::EventBusConfig, metrics::LivenessChecking},
     price_estimation::{
         config::price_estimation::BalanceOverridesConfigExt,
         factory::{self, PriceEstimatorFactory},
@@ -153,6 +153,17 @@ pub async fn start(args: impl Iterator<Item = String>) {
     );
     observe::tracing::init::initialize(&obs_config);
     observe::panic_hook::install();
+    if let Some(event_bus) = &config.shared.event_bus {
+        observe::event_bus::init(EventBusConfig {
+            url: event_bus.url.clone(),
+            channel: event_bus.channel.clone(),
+            chain_id: config
+                .shared
+                .chain_id
+                .expect("when the event bus is configured 'chain-id' is required"),
+        })
+        .await;
+    }
     #[cfg(unix)]
     observe::heap_dump_handler::spawn_heap_dump_handler();
 

crates/configs/src/shared.rs

@@ -96,6 +96,9 @@ pub struct SharedConfig {
     /// By default, volume fees are NOT applied to same-token trades.
     #[serde(default)]
     pub enable_sell_equals_buy_volume_fee: bool,
+
+    /// Enables publishing events to a global events bus.
+    pub event_bus: Option<EventBusConfig>,
 }
 
 impl Default for SharedConfig {
@@ -112,6 +115,7 @@ impl Default for SharedConfig {
             contracts: Default::default(),
             volume_fee_bucket_overrides: Vec::new(),
             enable_sell_equals_buy_volume_fee: false,
+            event_bus: None,
         }
     }
 }
@@ -311,6 +315,17 @@ where
     serializer.serialize_str(level.as_str())
 }
 
+/// OpenTelemetry tracing configuration.
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+#[derive(serde::Serialize)]
+pub struct EventBusConfig {
+    /// Url of the event bus service.
+    pub url: Url,
+    /// Name of the channel to post events to.
+    pub channel: String,
+}
+
 /// Gas price estimation strategy.
 #[derive(Debug, Clone, Deserialize)]
 #[serde(tag = "type")]
@@ -403,6 +418,10 @@ mod tests {
             "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
             "0x6B175474E89094C44Da98b954EedeAC495271d0F",
         ]
+
+        [event-bus]
+        url = "localhost:4222"
+        channel = "main"
         "#;
 
         let config: SharedConfig = toml::from_str(toml).unwrap();
@@ -430,5 +449,10 @@ mod tests {
         assert!(config.enable_sell_equals_buy_volume_fee);
         assert_eq!(config.volume_fee_bucket_overrides.len(), 1);
         assert_eq!(config.volume_fee_bucket_overrides[0].tokens.len(), 2);
+        assert_eq!(
+            config.event_bus.as_ref().unwrap().url,
+            "localhost:4222".parse().unwrap()
+        );
+        assert_eq!(config.event_bus.as_ref().unwrap().channel, "main");
     }
 }

crates/observe/Cargo.toml

@@ -6,7 +6,9 @@ edition = "2024"
 license = "MIT OR Apache-2.0"
 
 [dependencies]
+async-nats = { workspace = true }
 async-trait = { workspace = true }
+bytes = { workspace = true }
 axum = { workspace = true }
 chrono = { workspace = true, features = ["now"] }
 console-subscriber = { workspace = true, optional = true }
@@ -28,6 +30,7 @@ tracing = { workspace = true }
 tracing-opentelemetry = { workspace = true }
 tracing-serde = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["env-filter", "fmt", "time"] }
+url = { workspace = true }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["rt-multi-thread"] }

crates/observe/src/config.rs

@@ -1,4 +1,4 @@
-use {core::time::Duration, tracing::Level};
+use {core::time::Duration, tracing::Level, url::Url};
 
 #[derive(Debug, Clone)]
 pub struct Config {
@@ -88,3 +88,13 @@ impl TracingConfig {
         }
     }
 }
+
+/// Configures a backend wide event bus events can be posted to.
+pub struct EventBusConfig {
+    /// Url of the event bus service
+    pub url: Url,
+    /// Name of the channel to post events to
+    pub channel: String,
+    /// Which chain this service operates on
+    pub chain_id: u64,
+}

crates/observe/src/event_bus/mod.rs

@@ -0,0 +1,146 @@
+//! Implements a simple globally available way to publish events to an event
+//! bus. Under the hood it's using NATS. To support publishing events from
+//! synchronous contexts we use an unbounded channel as an in-memory buffer.
+//! Whenever a message gets posted to this channel a background task wakes
+//! up and forwards it to the NATS service running in a different process.
+//! Messages always get serialized as JSON so you can publish anything that
+//! can be serialized to JSON as well.
+use {
+    crate::config::EventBusConfig,
+    async_nats::jetstream::Context as JetstreamClient,
+    bytes::Bytes,
+    chrono::Utc,
+    serde::Serialize,
+    serde_json::json,
+    tokio::sync::{
+        OnceCell,
+        mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
+    },
+};
+
+struct EventBusConnector {
+    /// Unbounded channel to allow emitting events from synchrounous
+    /// contexts.
+    message_queue: UnboundedSender<Message>,
+    /// Subject prefix to disambiguate messages in globally shared event bus
+    /// service.
+    subject_prefix: String,
+}
+
+struct Message {
+    subject: String,
+    data: Bytes,
+}
+
+/// Singleton event bus connection to allow publishing events
+/// conventiently from everywhere.
+static BUS: OnceCell<EventBusConnector> = OnceCell::const_new();
+
+/// Initializes the event bus and panics if it fails.
+pub async fn init(config: EventBusConfig) {
+    BUS.get_or_init(|| async move {
+        let client = async_nats::connect(config.url.as_str())
+            .await
+            .expect("failed to connect to NATS service");
+        let jetstream = async_nats::jetstream::new(client);
+        let mut stream = jetstream
+            .get_stream(&config.channel)
+            .await
+            .expect("could not connect to jetstream");
+        let info = stream.info().await.expect("failed to fetch stream info");
+        tracing::debug!(?info, "connected to jetstream");
+
+        let (sender, receiver) = unbounded_channel();
+        tokio::task::spawn(forward_messages_to_event_bus_client(receiver, jetstream));
+        EventBusConnector {
+            message_queue: sender,
+            // we prefix every subject with `event` to allow consumers to easily
+            // subscribe to all events without also seeing NATS internal events
+            subject_prefix: format!("event.{}.", config.chain_id),
+        }
+    })
+    .await;
+}
+
+/// Monitors a message queue and forwards all messages to the event bus
+/// service.
+async fn forward_messages_to_event_bus_client(
+    mut receiver: UnboundedReceiver<Message>,
+    client: JetstreamClient,
+) {
+    while let Some(message) = receiver.recv().await {
+        match client.publish(message.subject, message.data).await {
+            Err(err) => {
+                tracing::debug!(?err, "failed to publish event");
+            }
+            Ok(_fut) => {
+                // let's assume the message arrived for now
+            }
+        }
+    }
+}
+
+/// Enqueues the event to be sent to the event bus in a background task.
+pub fn publish(subject: &str, data: impl Serialize) {
+    let Some(bus) = BUS.get() else {
+        return;
+    };
+
+    let mut message = json!({
+        "version": "v1",
+        "timestamp": Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
+        "body": data,
+    });
+    if let Some(id) = crate::tracing::distributed::request_id::from_current_span() {
+        message["requestId"] = id.into();
+    }
+    let body = match serde_json::to_vec(&message) {
+        Ok(body) => body,
+        Err(err) => {
+            tracing::error!(?err, "failed to serialize event");
+            return;
+        }
+    };
+
+    let message = Message {
+        subject: format!("{}{}", bus.subject_prefix, subject),
+        data: body.into(),
+    };
+
+    if let Err(err) = bus.message_queue.send(message) {
+        tracing::error!(?err, "failed to enqueue message");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use {super::*, serde_json::json};
+
+    #[ignore]
+    #[tokio::test]
+    async fn send_messages() {
+        crate::tracing::init::initialize(&crate::Config {
+            env_filter: "warn,observe=debug".to_string(),
+            stderr_threshold: None,
+            use_json_format: false,
+            tracing: None,
+        });
+        init(EventBusConfig {
+            url: "localhost:4222".parse().unwrap(),
+            channel: "main".to_string(),
+            chain_id: 1,
+        })
+        .await;
+
+        for _ in 0..1000 {
+            publish(
+                "name",
+                json!({
+                    "estimator": "baseline",
+                    "outAmount": 1234,
+                }),
+            );
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    }
+}

crates/observe/src/lib.rs

@@ -2,6 +2,7 @@
 //! improve the observability of a system. That includes initialization logic
 //! for metrics and logging as well as logging helper functions.
 pub mod config;
+pub mod event_bus;
 pub mod future;
 #[cfg(unix)]
 pub mod heap_dump_handler;

crates/orderbook/src/run.rs

@@ -29,7 +29,10 @@ use {
     http_client::HttpClientFactory,
     model::DomainSeparator,
     num::ToPrimitive,
-    observe::metrics::{DEFAULT_METRICS_PORT, serve_metrics},
+    observe::{
+        config::EventBusConfig,
+        metrics::{DEFAULT_METRICS_PORT, serve_metrics},
+    },
     order_validation,
     price_estimation::{
         PriceEstimating,
@@ -74,6 +77,17 @@ pub async fn start(args: impl Iterator<Item = String>) {
     tracing::info!("running order book with validated arguments:\n{}", args);
     observe::panic_hook::install();
     observe::metrics::setup_registry(Some("gp_v2_api".into()), None);
+    if let Some(event_bus) = &config.shared.event_bus {
+        observe::event_bus::init(EventBusConfig {
+            url: event_bus.url.clone(),
+            channel: event_bus.channel.clone(),
+            chain_id: config
+                .shared
+                .chain_id
+                .expect("when the event bus is configured 'chain-id' is required"),
+        })
+        .await;
+    }
     #[cfg(unix)]
     observe::heap_dump_handler::spawn_heap_dump_handler();
     tracing::info!("file configuration:\n{:#?}", config);

crates/price-estimation/src/competition/mod.rs

@@ -108,9 +108,10 @@ impl<T: Send + Sync + 'static> CompetitionEstimator<T> {
 
             while stage_index < self.stages.len() && requests.len() < requests_for_batch {
                 let stage = &self.stages.get(stage_index).expect("index checked by loop");
-                let futures = stage.iter().enumerate().map(|(index, (_name, estimator))| {
+                let futures = stage.iter().enumerate().map(|(index, (name, estimator))| {
                     get_single_result(Context {
                         estimator,
+                        name,
                         query: query.clone(),
                         remaining_stages: Arc::clone(&remaining_stages),
                     })
@@ -175,6 +176,8 @@ struct Context<'a, ESTIMATOR, QUERY> {
     /// the number of stages that are left after the queries
     /// produced by this Context's stages.
     remaining_stages: Arc<OnceLock<usize>>,
+    /// Name of the estimator
+    name: &'a str,
 }
 
 impl<'a, E, Q> Context<'a, E, Q> {