connection type specific readers/writers

server/kafka/kafka-script-proxy/src/main.rs

@@ -8,7 +8,7 @@ use tokio::fs::OpenOptions;
 use tokio::io::AsyncWriteExt;
 use tokio::process::Command;
 
-const TOPICS: [&str; 3] = ["move", "status", "input"];
+const TOPICS: [&str; 4] = ["move", "status", "input", "session"];
 
 #[tokio::main]
 pub async fn main() {

server/src/http.rs

@@ -2,6 +2,7 @@ use std::convert::Infallible;
 use std::fs::read;
 use std::io::ErrorKind::NotFound;
 use std::net::SocketAddr;
+use std::str::FromStr;
 use std::sync::Arc;
 use hyper::{Body, body, Method, Request, Response, Server, StatusCode};
 use hyper::server::conn::AddrStream;
@@ -16,17 +17,17 @@ use pong::event::event::{Event, EventReader, EventWriter};
 use futures::{sink::SinkExt, stream::StreamExt};
 use crate::kafka::{KafkaEventReaderImpl, KafkaSessionEventWriterImpl};
 use crate::player::Player;
-use crate::session::{CachingSessionManager, SessionManager};
+use crate::session::{SessionManager};
 use crate::utils::http_utils::{get_query_params, read_json_body};
 
 pub struct HttpServer {
     addr: [u8; 4],
     port: u16,
-    session_manager: Arc<Mutex<CachingSessionManager>>
+    session_manager: Arc<Mutex<SessionManager>>
 }
 impl HttpServer {
     pub fn new(addr: [u8; 4], port: u16, kafka_host: &str) -> HttpServer {
-        let session_manager = Arc::new(Mutex::new(CachingSessionManager::new(kafka_host)));
+        let session_manager = Arc::new(Mutex::new(SessionManager::new(kafka_host)));
         HttpServer {addr, port, session_manager}
     }
 
@@ -39,12 +40,40 @@ impl HttpServer {
                     let mut session_manager = Arc::clone(&session_manager);
                     async move {
                         if hyper_tungstenite::is_upgrade_request(&req) {
+                            println!("Received request to upgrade to websocket connection: {:?}", req);
+                            let params = get_query_params(&req);
+                            println!("Ws request params: {:?}", params);
+                            if !params.contains_key("session_id") {
+                                eprintln!("Missing session id request param for websocket connection, don't upgrade connection to ws.");
+                                return build_error_res("Missing request param: session_id", StatusCode::BAD_REQUEST);
+                            }
+                            if !params.contains_key("connection_type") {
+                                eprintln!("Missing connection type request param for websocket connection, don't upgrade connection to ws.");
+                                let res = build_error_res("Missing request param: connection_type", StatusCode::BAD_REQUEST);
+                                return res;
+                            }
+                            let session_id = params.get("session_id").unwrap();
+                            let connection_type_raw = params.get("connection_type").unwrap();
+                            let connection_type = WebSocketConnectionType::from_str(connection_type_raw);
+                            if let Err(_) = connection_type {
+                                let error = format!("Invalid connection type: {}", connection_type_raw);
+                                eprintln!("{}", error);
+                                return build_error_res(error.as_str(), StatusCode::BAD_REQUEST);
+                            }
+                            if let None = session_manager.lock().await.get_session(session_id) {
+                                let error = format!("Session does not exist: {}", session_id);
+                                eprintln!("{}", error);
+                                return build_error_res(error.as_str(), StatusCode::NOT_FOUND);
+                            }
+                            let websocket_session = WebSocketSession {session_id: session_id.to_string(), connection_type: connection_type.unwrap()};
+                            println!("Websocket upgrade request is valid, will now upgrade to websocket: {:?}", req);
+
                             let (response, websocket) = hyper_tungstenite::upgrade(req, None).unwrap();
 
                             // Spawn a task to handle the websocket connection.
                             tokio::spawn(async move {
-                                if let Err(e) = serve_websocket(websocket, session_manager).await {
-                                    eprintln!("Error in websocket connection: {}", e);
+                                if let Err(e) = serve_websocket(websocket_session, websocket, session_manager).await {
+                                    eprintln!("Error in websocket connection: {:?}", e);
                                 }
                             });
 
@@ -68,10 +97,21 @@ impl HttpServer {
 }
 
 /// Handle a websocket connection.
-async fn serve_websocket(websocket: HyperWebsocket, session_manager: Arc<Mutex<CachingSessionManager>>) -> Result<(), Error> {
+async fn serve_websocket(websocket_session: WebSocketSession, websocket: HyperWebsocket, session_manager: Arc<Mutex<SessionManager>>) -> Result<(), Error> {
     let mut websocket = websocket.await?;
+    let (mut write, mut read) = websocket.split();
+
+    let session_manager_reader = Arc::clone(&session_manager);
+    let websocket_session_read_copy = websocket_session.clone();
     tokio::spawn(async move {
-        while let Some(message) = websocket.next().await {
+        let mut locked = session_manager_reader.lock().await;
+        let writer = locked.get_session_writer(&websocket_session_read_copy.session_id);
+        if let Err(e) = writer {
+            eprintln!("Failed to retrieve session writer for session {}: {}", websocket_session_read_copy.session_id, e);
+            return
+        }
+        let mut writer = writer.unwrap();
+        while let Some(message) = read.next().await {
             match message.unwrap() {
                 Message::Text(msg) => {
                     let events = serde_json::from_str::<SessionEventListDTO>(&msg);
@@ -80,15 +120,12 @@ async fn serve_websocket(websocket: HyperWebsocket, session_manager: Arc<Mutex<C
                         continue;
                     }
                     let event_wrapper = events.unwrap();
-                    let mut locked = session_manager.lock().await;
-                    let writer = locked.get_session_writer(&event_wrapper.session_id);
-                    if let Err(e) = writer {
-                        eprintln!("Failed to retrieve session writer for session {}: {}", event_wrapper.session_id, e);
+                    if event_wrapper.session_id != websocket_session_read_copy.session_id {
+                        eprintln!("Websocket has session {:?} but was asked to write to session {} - skip.", websocket_session_read_copy, event_wrapper.session_id);
                         continue;
                     }
-                    let writer = writer.unwrap();
-                    let mut writer = writer.lock().await;
                     let mut any_error = false;
+                    let event_count = event_wrapper.events.len();
                     for event in event_wrapper.events {
                         let write_res = writer.write_to_session(&event.topic, &event.msg);
                         if let Err(e) = write_res {
@@ -98,6 +135,8 @@ async fn serve_websocket(websocket: HyperWebsocket, session_manager: Arc<Mutex<C
                     }
                     if any_error {
                         eprintln!("Failed to write at least one message for session {}", event_wrapper.session_id);
+                    } else {
+                        println!("Successfully wrote {} messages to kafka for session {:?}", event_count, websocket_session_read_copy)
                     }
                 },
                 Message::Close(msg) => {
@@ -112,16 +151,28 @@ async fn serve_websocket(websocket: HyperWebsocket, session_manager: Arc<Mutex<C
             }
         }
     });
+    let session_manager_writer = Arc::clone(&session_manager);
+    let websocket_session_write_copy = websocket_session.clone();
     tokio::spawn(async move {
-        // TODO: Get session writer and query consumer
-        // TODO: Differentiate player1 (receives only input of player2 and session updates) vs player2 (receives all messages)
+        let consumer = {
+            let consumer_async = async {
+                let mut session_manager_writer_locked = session_manager_writer.lock().await;
+                return session_manager_writer_locked.get_session_reader(&websocket_session_write_copy.session_id, websocket_session_write_copy.connection_type.get_topics());
+            };
+            consumer_async.await
+        };
+        let mut consumer = consumer.unwrap();
+        loop {
+            let messages = consumer.read_from_session();
+            println!("Read messages for websocket_session {:?} from consumer: {:?}", websocket_session_write_copy, messages);
+        }
     });
     Ok(())
 }
 
 // TODO: How to handle event writes/reads? This must be a websocket, but how to implement in hyper (if possible)?
 // https://github.com/de-vri-es/hyper-tungstenite-rs
-async fn handle_request(session_manager: &Arc<Mutex<CachingSessionManager>>, req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
+async fn handle_request(session_manager: &Arc<Mutex<SessionManager>>, req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
     println!("req to {} with method {}", req.uri().path(), req.method());
     match (req.method(), req.uri().path()) {
         (&Method::GET, "/session") => handle_get_session(session_manager, req).await,
@@ -133,7 +184,7 @@ async fn handle_request(session_manager: &Arc<Mutex<CachingSessionManager>>, req
     }
 }
 
-async fn handle_get_session(session_manager: &Arc<Mutex<CachingSessionManager>>, req: Request<Body>) -> Result<Response<Body>, Infallible> {
+async fn handle_get_session(session_manager: &Arc<Mutex<SessionManager>>, req: Request<Body>) -> Result<Response<Body>, Infallible> {
     let mut locked = session_manager.lock().await;
     let query_params = get_query_params(&req);
     let session_id = query_params.get("session_id");
@@ -148,7 +199,7 @@ async fn handle_get_session(session_manager: &Arc<Mutex<CachingSessionManager>>,
     return build_success_res(&serde_json::to_string(&session.unwrap()).unwrap());
 }
 
-async fn handle_session_create(session_manager: &Arc<Mutex<CachingSessionManager>>, req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
+async fn handle_session_create(session_manager: &Arc<Mutex<SessionManager>>, req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
     let mut locked = session_manager.lock().await;
     let session_create_res = locked.create_session(Player {id: addr.to_string()}).await;
     if let Err(e) = session_create_res {
@@ -158,7 +209,8 @@ async fn handle_session_create(session_manager: &Arc<Mutex<CachingSessionManager
     return build_success_res(&serialized.to_string());
 }
 
-async fn handle_session_join(session_manager: &Arc<Mutex<CachingSessionManager>>, mut req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
+async fn handle_session_join(session_manager: &Arc<Mutex<SessionManager>>, mut req: Request<Body>, addr: SocketAddr) -> Result<Response<Body>, Infallible> {
+    println!("Received request to join session: {:?}", req);
     let mut locked = session_manager.lock().await;
     let body = read_json_body::<SessionJoinDto>(&mut req).await;
     let session_join_res = locked.join_session(body.session_id, Player {id: addr.to_string()}).await;
@@ -166,11 +218,13 @@ async fn handle_session_join(session_manager: &Arc<Mutex<CachingSessionManager>>
         eprintln!("Failed to join session: {:?}", e);
         return Ok(Response::builder().status(StatusCode::INTERNAL_SERVER_ERROR).body(Body::from(e)).unwrap());
     }
-    let serialized = json!(session_join_res.unwrap());
+    let session = session_join_res.unwrap();
+    println!("Successfully joined session: {:?}", session);
+    let serialized = json!(session);
     return build_success_res(&serialized.to_string());
 }
 
-async fn handle_event_write(session_manager: &Arc<Mutex<CachingSessionManager>>, mut req: Request<Body>) -> Result<Response<Body>, Infallible> {
+async fn handle_event_write(session_manager: &Arc<Mutex<SessionManager>>, mut req: Request<Body>) -> Result<Response<Body>, Infallible> {
     let mut locked = session_manager.lock().await;
     let event = read_json_body::<SessionEventWriteDTO>(&mut req).await;
     let writer = locked.get_session_writer(&event.session_id);
@@ -182,9 +236,8 @@ async fn handle_event_write(session_manager: &Arc<Mutex<CachingSessionManager>>,
         return Ok(res);
     }
     let mut writer = writer.unwrap();
-    let mut writer_locked = writer.lock().await;
     println!("Writing session event to kafka: {:?}", event);
-    let write_res = writer_locked.write_to_session(&event.topic, &event.msg);
+    let write_res = writer.write_to_session(&event.topic, &event.msg);
     if let Err(e) = write_res {
         let err = format!("Failed to write event: {}", e);
         println!("{}", err);
@@ -196,10 +249,10 @@ async fn handle_event_write(session_manager: &Arc<Mutex<CachingSessionManager>>,
     build_success_res(&serde_json::to_string(&event).unwrap())
 }
 
-async fn handle_event_read(session_manager: &Arc<Mutex<CachingSessionManager>>, mut req: Request<Body>) -> Result<Response<Body>, Infallible> {
+async fn handle_event_read(session_manager: &Arc<Mutex<SessionManager>>, mut req: Request<Body>) -> Result<Response<Body>, Infallible> {
     let mut locked = session_manager.lock().await;
     let read_payload = read_json_body::<SessionReadDTO>(&mut req).await;
-    let reader = locked.get_session_reader(&read_payload.session_id);
+    let reader = locked.get_session_reader(&read_payload.session_id, &["move", "status", "input", "session"]);
     if let Err(e) = reader {
         let err = format!("Failed to read events: {}", e);
         println!("{}", err);
@@ -208,9 +261,8 @@ async fn handle_event_read(session_manager: &Arc<Mutex<CachingSessionManager>>,
         return Ok(res);
     }
     let mut reader = reader.unwrap();
-    let mut reader_locked = reader.lock().await;
     println!("Reading session events from kafka for session: {}", read_payload.session_id);
-    let events = reader_locked.read_from_session();
+    let events = reader.read_from_session();
     if let Err(e) = events {
         let err = format!("Failed to read events: {}", e);
         println!("{}", err);
@@ -268,3 +320,36 @@ struct SessionReadDTO {
 struct SessionJoinDto {
     session_id: String
 }
+
+#[derive(Debug, Clone)]
+enum WebSocketConnectionType {
+    HOST, PEER, OBSERVER
+}
+
+impl FromStr for WebSocketConnectionType {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "host" => Ok(WebSocketConnectionType::HOST),
+            "peer" => Ok(WebSocketConnectionType::PEER),
+            "observer" => Ok(WebSocketConnectionType::OBSERVER),
+            _ => Err(())
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+struct WebSocketSession {
+    pub connection_type: WebSocketConnectionType,
+    pub session_id: String
+}
+
+impl WebSocketConnectionType {
+    pub fn get_topics(&self) -> &[&str] {
+        match self {
+            WebSocketConnectionType::HOST => &["input", "session"],
+            WebSocketConnectionType::PEER | WebSocketConnectionType::OBSERVER => &["move", "input", "status", "session"],
+        }
+    }
+}

server/src/kafka.rs

@@ -119,8 +119,8 @@ impl KafkaEventReaderImpl {
         }
     }
 
-    pub fn for_partitions(host: &str, partitions: &[i32], topics: &[&str]) -> KafkaEventReaderImpl {
-        println!("Connecting partition specific consumer to kafka host: {}", host);
+    pub fn for_partitions(host: &str, partitions: &[i32], topics: &[&str]) -> Result<KafkaEventReaderImpl, String> {
+        println!("Connecting partition specific consumer to kafka host {} with topics {:?} / partitions {:?}", host, topics, partitions);
         let mut builder = Consumer::from_hosts(vec!(host.to_owned()));
         for topic in topics.iter() {
             builder = builder.with_topic_partitions(topic.parse().unwrap(), partitions);
@@ -131,11 +131,15 @@ impl KafkaEventReaderImpl {
             .with_offset_storage(GroupOffsetStorage::Kafka);
 
         let consumer = builder
-            .create()
-            .unwrap();
-        KafkaEventReaderImpl {
-            consumer
+            .create();
+        if let Err(e) = consumer {
+            eprintln!("Failed to connect consumer: {:?}", e);
+            return Err("Failed to connect consumer".to_string());
         }
+        let consumer = consumer.unwrap();
+        Ok(KafkaEventReaderImpl {
+            consumer
+        })
     }
 }
 impl EventReaderImpl for KafkaEventReaderImpl {
@@ -172,11 +176,16 @@ pub struct KafkaSessionEventReaderImpl {
 }
 
 impl KafkaSessionEventReaderImpl {
-    pub fn new(host: &str, session: &Session, topics: &[&str]) -> KafkaSessionEventReaderImpl {
+    pub fn new(host: &str, session: &Session, topics: &[&str]) -> Result<KafkaSessionEventReaderImpl, String> {
         let partitions = [session.id as i32];
-        KafkaSessionEventReaderImpl {
-            inner: KafkaEventReaderImpl::for_partitions(host, &partitions, topics)
+        let reader = KafkaEventReaderImpl::for_partitions(host, &partitions, topics);
+        if let Err(e) = reader {
+            return Err("Failed to create kafka session event reader".to_string());
         }
+        let reader = reader.unwrap();
+        Ok(KafkaSessionEventReaderImpl {
+            inner: reader
+        })
     }
 }
 
@@ -254,7 +263,7 @@ impl Partitioner for SessionPartitioner {
             Some(key) => {
                 let key = std::str::from_utf8(key).unwrap();
                 msg.partition = key.parse::<i32>().unwrap();
-                println!("Overriding message partition with key: {}", msg.partition);
+                // println!("Overriding message partition with key: {}", msg.partition);
             },
             None => panic!("Producing message without key not allowed!")
         }

server/src/session.rs

@@ -87,13 +87,18 @@ impl SessionManager {
         return Ok(());
     }
 
-    pub fn get_session_reader(&self, session_id: &str) -> Result<SessionReader, String> {
+    pub fn get_session_reader(&self, session_id: &str, topics: &[&str]) -> Result<SessionReader, String> {
         let session = self.find_session(&session_id);
         if let None = session {
             return Err(format!("Unable to find session with hash {}", session_id))
         }
         let session = session.unwrap();
-        let event_reader = EventReader::new(Box::new(KafkaSessionEventReaderImpl::new(&self.kafka_host, &session, &["move", "status", "input"])));
+        let kafka_reader = KafkaSessionEventReaderImpl::new(&self.kafka_host, &session, topics);
+        if let Err(_) = kafka_reader {
+            return Err("Unable to create kafka reader.".to_string())
+        }
+        let kafka_reader = kafka_reader.unwrap();
+        let event_reader = EventReader::new(Box::new(kafka_reader));
         Ok(SessionReader {reader: event_reader, session})
     }
 
@@ -145,7 +150,7 @@ impl CachingSessionManager {
             println!("Reusing existing reader for session: {:?}", session_id);
             return Ok(Arc::clone(reader));
         }
-        let reader = self.inner.get_session_reader(session_id);
+        let reader = self.inner.get_session_reader(session_id, &["move", "input", "status", "session"]);
         if let Err(e) = reader {
             return Err(e);
         }

www/index.js

@@ -88,7 +88,7 @@ window.WASM_PONG.createOnlineSession = () => {
         session_display_tag.style.display = 'block';
         session_display_tag.innerHTML = JSON.stringify(session)
 
-        websocket = new WebSocket("ws://localhost:4000")
+        websocket = new WebSocket(`ws://localhost:4000/ws?session_id=${session.hash}&connection_type=host`)
         websocket.onmessage = (event) => {
             console.log(event)
         }
@@ -109,7 +109,7 @@ window.WASM_PONG.joinOnlineSession = () => {
         session_display_tag.style.display = 'block';
         session_display_tag.innerHTML = JSON.stringify(session)
 
-        websocket = new WebSocket("ws://localhost:4000")
+        websocket = new WebSocket(`ws://localhost:4000/ws?session_id=${session.hash}&connection_type=peer`)
         websocket.onmessage = (event) => {
             console.log(event)
         }