feat: Redis Publish Subscribe

crates/routing_engine/src/indexer.rs

@@ -91,7 +91,7 @@ impl<
 
                     Ok::<
                         estimator::DataPoint<f64, f64>,
-                        IndexerErrors<TokenPriceProvider, RouteSource, ModelStore>,
+                        IndexerErrors<TokenPriceProvider, RouteSource, ModelStore, Producer>,
                     >(estimator::DataPoint {
                         x: input_value_in_usd,
                         y: fee_in_usd,
@@ -102,7 +102,7 @@ impl<
             .collect::<Vec<
                 Result<
                     estimator::DataPoint<f64, f64>,
-                    IndexerErrors<TokenPriceProvider, RouteSource, ModelStore>,
+                    IndexerErrors<TokenPriceProvider, RouteSource, ModelStore, Producer>,
                 >,
             >>()
             .await
@@ -125,7 +125,7 @@ impl<
     >(
         &mut self,
         values: Vec<(&&BucketConfig, &Estimator)>,
-    ) -> Result<(), IndexerErrors<TokenPriceProvider, RouteSource, ModelStore>> {
+    ) -> Result<(), IndexerErrors<TokenPriceProvider, RouteSource, ModelStore, Producer>> {
         let values_transformed = values
             .iter()
             .map(|(k, v)| {
@@ -150,7 +150,7 @@ impl<
         &mut self,
     ) -> Result<
         HashMap<&'config BucketConfig, Estimator>,
-        IndexerErrors<TokenPriceProvider, RouteSource, ModelStore>,
+        IndexerErrors<TokenPriceProvider, RouteSource, ModelStore, Producer>,
     > {
         // Build Estimators
         let estimator_map: HashMap<&BucketConfig, Estimator> =
@@ -189,6 +189,7 @@ enum IndexerErrors<
     T: token_price::TokenPriceProvider,
     S: source::RouteSource,
     R: storage::RoutingModelStore,
+    U: storage::MessageQueue,
 > {
     #[display(fmt = "Route build error: {}", _0)]
     RouteBuildError(RouteError),
@@ -206,15 +207,15 @@ enum IndexerErrors<
     PublishEstimatorErrors(Vec<R::Error>),
 
     #[display(fmt = "Indexer update message error: {}", _0)]
-    PublishIndexerUpdateMessageError(String),
+    PublishIndexerUpdateMessageError(U::Error),
 }
 
 #[cfg(test)]
 mod tests {
     use std::env;
 
     use config::Config;
-    use storage::{MessageQueue, RoutingModelStore};
+    use storage::{ControlFlow, MessageQueue, Msg, RoutingModelStore};
 
     use crate::CostType;
     use crate::estimator::{Estimator, LinearRegressionEstimator};
@@ -246,12 +247,18 @@ mod tests {
 
     struct ProducerStub;
     impl MessageQueue for ProducerStub {
-        async fn publish(&mut self, topic: &str, message: &str) -> Result<(), String> {
+        type Error = ();
+
+        async fn publish(&mut self, topic: &str, message: &str) -> Result<(), ()> {
             Ok(())
         }
 
-        async fn subscribe(&mut self, topic: &str) -> Result<String, String> {
-            Ok("Subscribed".to_string())
+        fn subscribe<U>(
+            &mut self,
+            topic: &str,
+            callback: impl FnMut(Msg) -> ControlFlow<U>,
+        ) -> Result<(), Self::Error> {
+            Ok(())
         }
     }
 

crates/routing_engine/src/source/bungee/mod.rs

@@ -122,8 +122,8 @@ impl RouteSource for BungeeClient {
             .iter()
             .map(|route| match estimation_type {
                 CostType::Fee => Some(
-                    route.total_gas_fees_in_usd + route.output_value_in_usd?
-                        - route.input_value_in_usd?,
+                    route.total_gas_fees_in_usd + route.input_value_in_usd?
+                        - route.output_value_in_usd?,
                 ),
                 _ => None,
             })

crates/storage/Cargo.toml

@@ -7,4 +7,4 @@ edition = "2021"
 config = { path = "../config" }
 redis = { version = "0.25.4", features = ["aio", "tokio-comp"] }
 thiserror = "1.0.61"
-tokio = { version = "1.38.0", features = ["macros"] }
+tokio = { version = "1.38.0", features = ["macros", "rt"] }

crates/storage/src/lib.rs

@@ -1,5 +1,7 @@
 use std::fmt::Debug;
 
+pub use ::redis::{ControlFlow, Msg};
+
 mod redis;
 
 pub trait RoutingModelStore {
@@ -15,7 +17,13 @@ pub trait RoutingModelStore {
 }
 
 pub trait MessageQueue {
-    async fn publish(&mut self, topic: &str, message: &str) -> Result<(), String>;
+    type Error: Debug;
+
+    async fn publish(&mut self, topic: &str, message: &str) -> Result<(), Self::Error>;
 
-    async fn subscribe(&mut self, topic: &str) -> Result<String, String>;
+    fn subscribe<U>(
+        &mut self,
+        topic: &str,
+        callback: impl FnMut(Msg) -> ControlFlow<U>,
+    ) -> Result<(), Self::Error>;
 }

crates/storage/src/redis.rs

@@ -1,5 +1,5 @@
 use redis;
-use redis::{aio, AsyncCommands};
+use redis::{aio, AsyncCommands, Commands, ControlFlow, Msg, PubSubCommands};
 use redis::RedisError;
 use thiserror::Error;
 
@@ -41,12 +41,20 @@ impl RoutingModelStore for RedisClient {
 }
 
 impl MessageQueue for RedisClient {
-    async fn publish(&mut self, topic: &str, message: &str) -> Result<(), String> {
-        todo!()
+    type Error = RedisClientError;
+
+    async fn publish(&mut self, topic: &str, message: &str) -> Result<(), Self::Error> {
+        self.connection.publish(topic, message).await.map_err(RedisClientError::RedisLibraryError)
     }
 
-    async fn subscribe(&mut self, topic: &str) -> Result<String, String> {
-        todo!()
+    fn subscribe<U>(
+        &mut self,
+        topic: &str,
+        callback: impl FnMut(Msg) -> ControlFlow<U>,
+    ) -> Result<(), Self::Error> {
+        let mut connection = self.client.get_connection()?;
+        connection.subscribe(topic, callback)?;
+        Ok(())
     }
 }
 
@@ -58,6 +66,8 @@ pub enum RedisClientError {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::mpsc::channel;
+
     use tokio;
 
     use super::*;
@@ -68,7 +78,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_redis_client() {
+    async fn test_key_store() {
         let mut client = setup().await;
 
         let keys = vec!["test_key1".to_string(), "test_key2".to_string()];
@@ -100,4 +110,31 @@ mod tests {
         let values = client.get_multiple(&keys).await.unwrap();
         assert_eq!(values, vec!["test_value1".to_string(), "test_value2".to_string()]);
     }
+
+    #[tokio::test]
+    async fn test_pub_sub() {
+        let (tx, mut rx) = channel::<String>();
+        let mut client = setup().await;
+
+        let handle = tokio::task::spawn_blocking(move || {
+            client
+                .subscribe("TOPIC", |msg: Msg| -> ControlFlow<String> {
+                    let message = msg.get_payload().unwrap();
+                    tx.send(message).expect("Sending message failed");
+
+                    ControlFlow::Break("DONE".to_string())
+                })
+                .unwrap();
+        });
+
+        let mut client = setup().await;
+        client.publish("TOPIC", "HELLO").await.unwrap();
+
+        loop {
+            if let Ok(data) = rx.recv() {
+                assert_eq!(data, "HELLO".to_string());
+                break;
+            }
+        }
+    }
 }