Merge pull request #11 from dorianim/feat/redis-channels

Feat/redis channels

Author: dorianim

src/main.rs

@@ -11,6 +11,7 @@ use tracing::Span;
 mod color;
 mod routes;
 
+use tokio::sync::broadcast;
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct Segment {
     label: String,
@@ -19,7 +20,7 @@ pub struct Segment {
     color: Option<Color>,
 }
 
-#[derive(Serialize, Deserialize, Clone, Default)]
+#[derive(Serialize, Deserialize, Clone, Default, Debug)]
 pub struct Timer {
     // Return after TimerRequest
     pub segments: Vec<Segment>,
@@ -44,9 +45,9 @@ struct InstanceProperties {
 type SharedState = Arc<AppState>;
 pub struct AppState {
     redis: redis::aio::ConnectionManager,
-    redis_client: redis::Client,
     jwt_key: String,
     instance_properties: InstanceProperties,
+    redis_task_rx: broadcast::Receiver<Timer>,
 }
 
 #[tokio::main]
@@ -58,7 +59,7 @@ async fn main() {
 
     let redis_string = env::var("REDIS_STRING").expect("REDIS_STRING is not set");
     let jwt_key = env::var("JWT_KEY").expect("JWT_KEY is not set");
-    let client = redis::Client::open(redis_string.to_owned()).unwrap();
+    let client = redis::Client::open(redis_string.to_owned()).expect("Could not connect to redis");
     let manager = redis::aio::ConnectionManager::new(client.clone())
         .await
         .unwrap();
@@ -70,23 +71,28 @@ async fn main() {
             .unwrap_or(None),
     };
 
+    let (redis_task_tx, redis_task_rx) = broadcast::channel::<Timer>(10);
+
     let state: SharedState = Arc::new(AppState {
-        redis: manager,
-        redis_client: client,
+        redis: manager.clone(),
         jwt_key,
         instance_properties,
+        redis_task_rx,
     });
 
     let app = Router::new()
-        .nest("/api/ws", routes::ws::routes())
+        .nest(
+            "/api/ws",
+            routes::ws::routes(manager, client, redis_task_tx),
+        )
         .nest("/api/timer", routes::timer::routes(state.clone()))
         .nest("/api/instance", routes::instance::routes())
         .fallback(routes::client::client_assets)
         .layer(cors)
         .layer(
             TraceLayer::new_for_http()
-                .on_request(|_request: &Request<_>, _span: &Span| {
-                    println!("Request {}", _request.uri());
+                .on_request(|request: &Request<_>, _span: &Span| {
+                    println!("Request {} {}", request.method(), request.uri());
                 })
                 .on_response(|_response: &Response, _latency: Duration, _span: &Span| {
                     println!(

src/routes/ws.rs

@@ -17,6 +17,7 @@ use serde::{Deserialize, Serialize};
 use serde_json;
 
 use redis::AsyncCommands;
+use tokio::sync::broadcast;
 use tokio::sync::mpsc::{Receiver, Sender};
 use tokio::task::JoinHandle;
 
@@ -41,16 +42,6 @@ struct WsConnection {}
 
 impl WsConnection {
     async fn new(state: SharedState, socket: WebSocket) {
-        let mut connection = state.redis_client.get_async_connection().await.unwrap();
-        let _: () = redis::cmd("CONFIG")
-            .arg("SET")
-            .arg("notify-keyspace-events")
-            .arg("KEA")
-            .query_async(&mut connection)
-            .await
-            .unwrap();
-        let pubsub = connection.into_pubsub();
-
         let (ws_sender, ws_receiver) = socket.split();
         let (ws_message_tx, ws_message_rx) = tokio::sync::mpsc::channel::<WSMessage>(32);
         let (redis_listen_id_tx, redis_listen_id_rx) = tokio::sync::mpsc::channel::<String>(32);
@@ -63,10 +54,9 @@ impl WsConnection {
             ws_receiver,
         );
         let redis_listener_task = WsConnection::spawn_redis_listener_task(
-            state.redis.clone(),
             ws_message_tx,
             redis_listen_id_rx,
-            pubsub,
+            state.redis_task_rx.resubscribe(),
         );
 
         ws_receiver_task.await.unwrap();
@@ -76,10 +66,9 @@ impl WsConnection {
     }
 
     fn spawn_redis_listener_task(
-        mut redis: redis::aio::ConnectionManager,
         ws_message_tx: Sender<WSMessage>,
         mut redis_listen_id_rx: Receiver<String>,
-        mut pubsub: redis::aio::PubSub,
+        mut redis_task_rx: tokio::sync::broadcast::Receiver<Timer>,
     ) -> JoinHandle<()> {
         tokio::spawn(async move {
             let mut msg = None;
@@ -88,25 +77,14 @@ impl WsConnection {
             }
 
             let timer_id = msg.unwrap();
-            pubsub
-                .psubscribe(format!("__keyspace@*__:{}", &timer_id))
-                .await
-                .unwrap();
-            println!("Redis listening!");
-            redis_listen_id_rx.close();
-
-            let mut pubsub = pubsub.into_on_message();
-
-            while let Some(msg) = pubsub.next().await {
-                println!("Updated! {:?}", msg);
 
-                let timer: Timer = serde_json::from_str(
-                    &redis.get::<String, String>(timer_id.clone()).await.unwrap(),
-                )
-                .unwrap();
+            while let Ok(timer) = redis_task_rx.recv().await {
+                if timer.id == timer_id {
+                    println!("Updated! {:?}", timer);
 
-                let response = WSMessage::Timer(timer.into());
-                ws_message_tx.send(response).await.unwrap();
+                    let response = WSMessage::Timer(timer.into());
+                    ws_message_tx.send(response).await.unwrap();
+                }
             }
         })
     }
@@ -214,6 +192,51 @@ pub async fn ws_handler(
     ws.on_upgrade(move |socket| WsConnection::new(state, socket))
 }
 
-pub fn routes() -> Router<SharedState> {
+fn spawn_global_redis_listener_task(
+    mut redis: redis::aio::ConnectionManager,
+    redis_client: redis::Client,
+    redis_task_tx: broadcast::Sender<Timer>,
+) -> JoinHandle<()> {
+    tokio::spawn(async move {
+        let mut connection = redis_client.get_async_connection().await.unwrap();
+        let _: () = redis::cmd("CONFIG")
+            .arg("SET")
+            .arg("notify-keyspace-events")
+            .arg("KEA")
+            .query_async(&mut connection)
+            .await
+            .unwrap();
+
+        let mut pubsub = connection.into_pubsub();
+
+        pubsub
+            .psubscribe("__keyspace@*__:*")
+            .await
+            .expect("Failed to subscribe to redis channel");
+
+        let mut pubsub = pubsub.into_on_message();
+
+        while let Some(msg) = pubsub.next().await {
+            println!("Updated! {:?}", msg);
+            let timer_id = msg.get_channel_name().split(":").last().unwrap();
+
+            let timer_str = &redis
+                .get::<String, String>(String::from(timer_id))
+                .await
+                .expect("Did not find timer in redis");
+            let timer: Timer = serde_json::from_str(timer_str).unwrap();
+
+            // Broadcast to all listeners
+            redis_task_tx.send(timer).unwrap();
+        }
+    })
+}
+
+pub fn routes(
+    redis: redis::aio::ConnectionManager,
+    redis_client: redis::Client,
+    redis_task_tx: broadcast::Sender<Timer>,
+) -> Router<SharedState> {
+    spawn_global_redis_listener_task(redis, redis_client, redis_task_tx);
     Router::new().route("/", get(ws_handler))
 }