Checkpointing websocket server.

Author: rohitkulshreshtha

Cargo.lock

@@ -59,6 +59,7 @@ serde_json = "1.0.115"
 slotmap = "1.0.0"
 smallvec = "1.6.1"
 tokio-stream = { version = "0.1.3", default-features = false, features = [ "time", "io-util", "sync" ] }
+tokio-tungstenite = "0.23.1"
 tracing = "0.1.37"
 variadics = { path = "../variadics", version = "^0.0.6" }
 web-time = "1.0.0"

hydroflow/Cargo.toml

@@ -49,17 +49,18 @@ async fn main() {
     let opts = Opts::parse();
 
     match opts.role {
-        Role::Client => {
-            run_client(opts).await;
-        }
+        // Role::Client => {
+        //     run_client(opts).await;
+        // }
         Role::Server => {
             run_server(opts).await;
         }
-        Role::GossipingServer1
-        | Role::GossipingServer2
-        | Role::GossipingServer3
-        | Role::GossipingServer4
-        | Role::GossipingServer5 => run_gossiping_server(opts).await,
+        // Role::GossipingServer1
+        // | Role::GossipingServer2
+        // | Role::GossipingServer3
+        // | Role::GossipingServer4
+        // | Role::GossipingServer5 => run_gossiping_server(opts).await,
+        _ => unimplemented!(),
     }
 }
 

hydroflow/examples/chat/main.rs

@@ -1,6 +1,6 @@
 use hydroflow::hydroflow_syntax;
 use hydroflow::scheduled::graph::Hydroflow;
-use hydroflow::util::bind_udp_bytes;
+use hydroflow::util::{bind_udp_bytes, bind_websocket};
 
 use crate::protocol::{Message, MessageWithAddr};
 use crate::{default_server_address, Opts};
@@ -14,27 +14,20 @@ pub(crate) async fn run_server(opts: Opts) {
 
     println!("Starting server on {:?}", server_address);
 
-    let (outbound, inbound, actual_server_addr) = bind_udp_bytes(server_address).await;
+    let (outbound, inbound, actual_server_addr) = bind_websocket(server_address).await.unwrap();
 
     println!("Server is live! Listening on {:?}", actual_server_addr);
 
     let mut hf: Hydroflow = hydroflow_syntax! {
         // Define shared inbound and outbound channels
-        outbound_chan = union() -> dest_sink_serde(outbound);
-        inbound_chan = source_stream_serde(inbound)
+        inbound_chan = source_stream(inbound)
             -> map(Result::unwrap)
-            -> map(|(msg, addr)| MessageWithAddr::from_message(msg, addr))
-            -> demux_enum::<MessageWithAddr>();
-        clients = inbound_chan[ConnectRequest] -> map(|(addr,)| addr) -> tee();
-        inbound_chan[ConnectResponse] -> for_each(|(addr,)| println!("Received unexpected `ConnectResponse` as server from addr {}.", addr));
-
-        // Pipeline 1: Acknowledge client connections
-        clients[0] -> map(|addr| (Message::ConnectResponse, addr)) -> [0]outbound_chan;
-
-        // Pipeline 2: Broadcast messages to all clients
-        inbound_chan[ChatMsg] -> map(|(_addr, nickname, message, ts)| Message::ChatMsg { nickname, message, ts }) -> [0]broadcast;
-        clients[1] -> [1]broadcast;
-        broadcast = cross_join::<'tick, 'static>() -> [1]outbound_chan;
+            -> tee();
+
+        inbound_chan -> map(|(msg, addr)| addr) -> [1]broadcast;
+        inbound_chan -> map(|(msg, addr)| msg) -> [0]broadcast;
+
+        broadcast = cross_join::<'tick, 'static>() -> dest_sink(outbound);
     };
 
     #[cfg(feature = "debugging")]

hydroflow/examples/chat/server.rs

@@ -2,7 +2,7 @@ use std::net::SocketAddr;
 
 use clap::{Parser, ValueEnum};
 use client::run_client;
-use hydroflow::util::{bind_udp_bytes, ipv4_resolve};
+use hydroflow::util::{bind_udp_bytes, bind_websocket, ipv4_resolve};
 use server::run_server;
 
 mod client;
@@ -35,16 +35,14 @@ async fn main() {
         .unwrap_or_else(|| ipv4_resolve("localhost:0").unwrap());
 
     // allocate `outbound` sink and `inbound` stream
-    let (outbound, inbound, addr) = bind_udp_bytes(addr).await;
+    let (outbound, inbound, addr) = bind_websocket(addr).await.unwrap();
     println!("Listening on {:?}", addr);
 
     match opts.role {
         Role::Server => {
             run_server(outbound, inbound, opts).await;
         }
-        Role::Client => {
-            run_client(outbound, inbound, opts).await;
-        }
+        _ => panic!("Unsupported!")
     }
 }
 

hydroflow/examples/echoserver/main.rs

@@ -29,6 +29,8 @@ pub use socket::*;
 
 #[cfg(feature = "deploy_integration")]
 pub mod deploy;
+mod websocket;
+pub use websocket::*;
 
 use std::io::Read;
 use std::net::SocketAddr;

hydroflow/src/util/mod.rs

@@ -0,0 +1,77 @@
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::net::SocketAddr;
+use std::pin::pin;
+use std::rc::Rc;
+use futures::{SinkExt, StreamExt};
+use tokio::net::{TcpListener};
+use tokio::task::spawn_local;
+use tokio_tungstenite::tungstenite::{Error, Message};
+use crate::util::unsync::mpsc::{Receiver, Sender};
+use crate::util::unsync_channel;
+
+pub async fn bind_websocket(endpoint: SocketAddr) -> Result<(Sender<(Message, SocketAddr)>, Receiver<Result<(Message, SocketAddr), Error>>, SocketAddr), std::io::Error>{
+    let listener = TcpListener::bind(endpoint).await.unwrap();
+
+    let bound_endpoint = listener.local_addr()?;
+
+    let (tx_egress, mut rx_egress) = unsync_channel(None);
+    let (tx_ingress, rx_ingress) = unsync_channel(None);
+
+    let clients = Rc::new(RefCell::new(HashMap::new()));
+
+    spawn_local({
+        let clients = clients.clone();
+
+        async move {
+            while let Some((payload, addr)) = rx_egress.next().await {
+                let client = clients.borrow_mut().remove(&addr);
+
+                if let Some(mut sender) = client {
+                    let _ = SinkExt::send(&mut sender, payload).await;
+                    clients.borrow_mut().insert(addr, sender);
+                }
+            }
+        }
+    });
+
+    // Spawn the listener
+    spawn_local(async move {
+        loop {
+            let (stream, peer_addr) = if let Ok((stream, _)) = listener.accept().await {
+                if let Ok(peer_addr) = stream.peer_addr() {
+                    (stream, peer_addr)
+                } else {
+                    continue;
+                }
+            } else {
+                continue;
+            };
+
+            // Perform the websocket handshake
+            let ws_stream = tokio_tungstenite::accept_async(stream)
+                .await
+                .expect("Error during the websocket handshake occurred");
+
+            // Split the stream into incoming and outgoing
+            let (outgoing, incoming) = ws_stream.split();
+            let mut tx_ingress = tx_ingress.clone();
+
+            clients.borrow_mut().insert(peer_addr, outgoing);
+
+            spawn_local({
+                let clients = clients.clone();
+                async move {
+                    let mapped = incoming.map(|x| Ok(x.map(|x| (x, peer_addr))));
+                    let _ = tx_ingress.send_all(&mut pin!(mapped)).await;
+
+                    clients.borrow_mut().remove(&peer_addr);
+                }
+            });
+        }
+
+    });
+
+    Ok((tx_egress, rx_ingress, bound_endpoint))
+}
+