From 2042fb0923b864742eb14132a59e1ebd70a09e0c Mon Sep 17 00:00:00 2001 From: Claas Date: Sat, 20 Dec 2025 21:52:49 +0100 Subject: [PATCH] Start work on actor model --- Cargo.lock | 1 + delivery-service/Cargo.toml | 1 + delivery-service/src/actor/mod.rs | 3 + delivery-service/src/actor/switchboard.rs | 166 ++++++++++++++++++++++ delivery-service/src/actor/user.rs | 118 +++++++++++++++ delivery-service/src/actor/web_socket.rs | 104 ++++++++++++++ delivery-service/src/main.rs | 24 +++- 7 files changed, 415 insertions(+), 2 deletions(-) create mode 100644 delivery-service/src/actor/mod.rs create mode 100644 delivery-service/src/actor/switchboard.rs create mode 100644 delivery-service/src/actor/user.rs create mode 100644 delivery-service/src/actor/web_socket.rs diff --git a/Cargo.lock b/Cargo.lock index 2f8b368..d337df8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1156,6 +1156,7 @@ name = "delivery-service" version = "0.1.0" dependencies = [ "axum 0.7.9", + "nanoid", "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions", diff --git a/delivery-service/Cargo.toml b/delivery-service/Cargo.toml index fbee915..4ac37bd 100644 --- a/delivery-service/Cargo.toml +++ b/delivery-service/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] axum = { workspace = true, features = ["macros", "ws"] } +nanoid = "0.4.0" opentelemetry = "0.31.0" opentelemetry-otlp = { version = "0.31.0", features = ["tonic", "grpc-tonic"] } opentelemetry-semantic-conventions = "0.31.0" diff --git a/delivery-service/src/actor/mod.rs b/delivery-service/src/actor/mod.rs new file mode 100644 index 0000000..a901a9c --- /dev/null +++ b/delivery-service/src/actor/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod switchboard; +mod user; +mod web_socket; diff --git a/delivery-service/src/actor/switchboard.rs b/delivery-service/src/actor/switchboard.rs new file mode 100644 index 0000000..286ccbd --- /dev/null +++ b/delivery-service/src/actor/switchboard.rs @@ -0,0 +1,166 @@ +use std::{collections::HashMap, sync::Arc}; + +use axum::extract::ws; +use tokio::sync::{mpsc, oneshot}; + +use crate::actor::{user, web_socket}; + +struct UserNotFound; + +enum Message { + Send { + user_id: Arc, + message: Arc<[u8]>, + response: oneshot::Sender>, + }, + Connect { + user_id: Arc, + web_socket: ws::WebSocket, + }, +} + +struct Switchboard { + receiver: mpsc::Receiver, + users: HashMap, user::Handle>, +} + +impl Switchboard { + fn new(receiver: mpsc::Receiver) -> Self { + Self { + receiver, + users: HashMap::new(), + } + } + + async fn handle_message(&mut self, message: Message) { + match message { + Message::Send { + user_id, + message, + response, + } => { + //TODO might need to store the message for the user to receive later + let Some(user) = self.users.get_mut(&user_id) else { + // We don't care if they stopped waiting for the response + _ = response.send(Err(UserNotFound)); + return; + }; + + let Err(user::HandleError::Closed) = user.send_message(message).await else { + _ = response.send(Ok(())); + return; + }; + + // TODO user is not active. Store message for the user to receive later or send push notification + _ = response.send(Ok(())); + } + Message::Connect { + user_id, + web_socket, + } => { + let socket = web_socket::Handle::new(web_socket); + + let user = self + .users + .entry(user_id.clone()) + .or_insert_with(|| user::Handle::new(user_id.clone())); + + let Err(user::HandleError::Closed) = user.add_socket(socket.clone()).await else { + return; + }; + + tracing::debug!( + "[Switchboard] User actor {} is dead. Rebirthing user.", + user_id + ); + let user = user.rebirth(); + self.users.insert(user_id, user.clone()); + + if let Err(user::HandleError::Closed) = user.add_socket(socket).await { + tracing::error!("[Switchboard] Just rebirthed user actor is already dead") + } + } + } + } + + async fn run(mut self) { + loop { + match self.receiver.recv().await { + Some(message) => self.handle_message(message).await, + None => { + tracing::debug!("[Switchboard] Actor send channel closed"); + return; + } + } + } + } +} + +#[derive(Clone)] +pub(crate) struct Handle { + sender: mpsc::Sender, +} + +pub(crate) enum SendMessageError { + Closed, + UserNotFound, +} + +pub(crate) enum ConnectError { + Closed, +} + +impl From for SendMessageError { + fn from(_: UserNotFound) -> Self { + Self::UserNotFound + } +} + +impl Default for Handle { + fn default() -> Self { + Self::new() + } +} + +impl Handle { + fn new() -> Self { + let (sender, receiver) = mpsc::channel(8); + + let actor = Switchboard::new(receiver); + tokio::spawn(actor.run()); + Self { sender } + } + + pub(in crate::actor) async fn send_message( + &self, + user_id: Arc, + message: Arc<[u8]>, + ) -> Result<(), SendMessageError> { + let (sender, receiver) = oneshot::channel(); + self.sender + .send(Message::Send { + user_id, + message, + response: sender, + }) + .await + .map_err(|_| SendMessageError::Closed)?; + + receiver.await.map_err(|_| SendMessageError::Closed)??; + Ok(()) + } + + pub async fn add_connection( + &self, + user_id: Arc, + web_socket: ws::WebSocket, + ) -> Result<(), ConnectError> { + self.sender + .send(Message::Connect { + user_id, + web_socket, + }) + .await + .map_err(|_| ConnectError::Closed) + } +} diff --git a/delivery-service/src/actor/user.rs b/delivery-service/src/actor/user.rs new file mode 100644 index 0000000..392fff8 --- /dev/null +++ b/delivery-service/src/actor/user.rs @@ -0,0 +1,118 @@ +use std::{collections::HashMap, sync::Arc}; + +use tokio::{sync::mpsc, task::JoinSet}; + +use crate::actor::web_socket; + +enum Message { + Send(Arc<[u8]>), + AddConnection(web_socket::Handle), +} + +struct User { + id: Arc, + receiver: mpsc::Receiver, + sockets: HashMap, web_socket::Handle>, +} + +impl User { + fn new(id: Arc, receiver: mpsc::Receiver) -> Self { + Self { + id, + receiver, + sockets: HashMap::new(), + } + } + + async fn handle_message(&mut self, message: Message) { + match message { + Message::Send(data) => { + let mut sends = JoinSet::new(); + for socket in self.sockets.values().cloned() { + let data = data.clone(); + sends + .spawn(async move { (socket.id.clone(), socket.send_message(data).await) }); + } + + let results = sends.join_all().await; + for (id, result) in results { + let Err(web_socket::HandleError::SocketClosed) = result else { + continue; + }; + + tracing::debug!("[{}] Removing closed websocket receiver {}", self.id, id); + + self.sockets.remove(&id); + } + + // If there are no more sockets, remove the user + if self.sockets.is_empty() { + tracing::debug!("[{}] No more sockets. Stopping user actor", self.id); + return; + } + } + Message::AddConnection(handle) => { + self.sockets.insert(handle.id.clone(), handle); + } + } + } + + async fn run(mut self) { + loop { + match self.receiver.recv().await { + Some(message) => self.handle_message(message).await, + None => { + tracing::debug!("[{}] Actor send channel closed", self.id); + return; + } + } + } + } +} + +#[derive(Clone)] +pub(in crate::actor) struct Handle { + pub(in crate::actor) id: Arc, + sender: mpsc::Sender, +} + +pub(in crate::actor) enum HandleError { + Closed, +} + +impl Handle { + pub(in crate::actor) fn new(id: Arc) -> Self { + let (sender, receiver) = mpsc::channel(8); + + let actor = User::new(id, receiver); + let id = actor.id.clone(); + tokio::spawn(actor.run()); + + Self { id, sender } + } + + /// Revive is the same entity but this creates a new instance so it is a rebirth + pub(in crate::actor) fn rebirth(&self) -> Self { + Self::new(self.id.clone()) + } + + pub(in crate::actor) async fn add_socket( + &self, + socket: web_socket::Handle, + ) -> Result<(), HandleError> { + self.sender + .send(Message::AddConnection(socket)) + .await + .map_err(|_| HandleError::Closed) + } + + pub(in crate::actor) async fn send_message( + &self, + data: impl Into>, + ) -> Result<(), HandleError> { + self.sender + .send(Message::Send(data.into())) + .await + .map_err(|_| HandleError::Closed) + } +} diff --git a/delivery-service/src/actor/web_socket.rs b/delivery-service/src/actor/web_socket.rs new file mode 100644 index 0000000..e087f43 --- /dev/null +++ b/delivery-service/src/actor/web_socket.rs @@ -0,0 +1,104 @@ +use std::{ops::ControlFlow, sync::Arc}; + +use axum::extract::ws::{self}; +use nanoid::nanoid; +use tokio::sync::mpsc; + +enum Message { + Send(Arc<[u8]>), +} +struct WebSocket { + id: Arc, + receiver: mpsc::Receiver, + socket: ws::WebSocket, +} + +impl WebSocket { + fn new(receiver: mpsc::Receiver, socket: ws::WebSocket) -> Self { + Self { + id: nanoid!().into(), + receiver, + socket, + } + } + + async fn handle_message(&mut self, Message::Send(data): Message) -> Result<(), axum::Error> { + // Am I just copying unnecessarily converting from Rc to Vec when I could send a Vec in the first place? + let message = ws::Message::Binary(Vec::from(data.as_ref())); + self.socket.send(message).await?; + + Ok(()) + } + + fn handle_socket_message(&mut self, message: ws::Message) -> ControlFlow<()> { + match message { + ws::Message::Close(_) => { + tracing::debug!("[{}] Websocket closed", self.id); + ControlFlow::Break(()) + } + other => { + tracing::debug!("[{}] Unexpected Websocket message: {:?}", self.id, other); + ControlFlow::Continue(()) + } + } + } + + async fn run(mut self) -> Result<(), axum::Error> { + loop { + tokio::select! { + message = self.receiver.recv() => match message { + Some(message) => self.handle_message(message).await?, + None => { + tracing::debug!("[{}] Actor send channel closed", self.id); + return Ok(()); + }, + }, + message = self.socket.recv() => match message { + Some(Ok(message)) => if self.handle_socket_message(message).is_break() { return Ok(()); }, + Some(Err(error)) => { + tracing::debug!("[{}] websocket error: {}", self.id, error); + return Err(error.into()); + }, + None => { + tracing::debug!("[{}] client closed websocket", self.id); + return Ok(()); + } + }, + } + } + } +} + +#[derive(Clone)] +pub(in crate::actor) struct Handle { + pub(in crate::actor) id: Arc, + sender: mpsc::Sender, +} + +pub(in crate::actor) enum HandleError { + SocketClosed, +} + +impl Handle { + pub(in crate::actor) fn new(socket: ws::WebSocket) -> Self { + let (sender, receiver) = mpsc::channel(8); + + let actor = WebSocket::new(receiver, socket); + let id = actor.id.clone(); + + tokio::spawn(actor.run()); + + Self { id, sender } + } + + pub(in crate::actor) async fn send_message( + &self, + data: impl Into>, + ) -> Result<(), HandleError> { + let data = data.into(); + self.sender + .send(Message::Send(data)) + .await + .map_err(|_| HandleError::SocketClosed) + } +} diff --git a/delivery-service/src/main.rs b/delivery-service/src/main.rs index cc9069f..181fadc 100644 --- a/delivery-service/src/main.rs +++ b/delivery-service/src/main.rs @@ -24,12 +24,16 @@ use tower_http::{ }; use tracing::Level; +use crate::actor::switchboard; + +mod actor; mod extractor; mod telemetry; #[derive(Clone, Default)] struct AppState { channels: Arc, Arc>>>>, + switchboard: switchboard::Handle, } /// The single page application setup used in production. During development a vite proxy is used to host the app and @@ -111,6 +115,14 @@ async fn main() -> Result<(), Box> { Ok(()) } +struct Droppochino(String); + +impl Drop for Droppochino { + fn drop(&mut self) { + println!("Droppochino {}", self.0); + } +} + /// This endpoint receives messages sent by clients to be delivered to other clients #[tracing::instrument(skip(state))] async fn create_message( @@ -137,6 +149,8 @@ async fn create_message( #[tracing::instrument(skip(socket, state))] async fn handle_socket(mut socket: WebSocket, State(state): State, client_id: Arc) { debug!("[{}] connected", client_id); + + let _droppochino = Droppochino(client_id.to_string()); let span = tracing::span!(Level::INFO, "handling"); let _enter = span.enter(); //TODO keep alive @@ -169,6 +183,7 @@ async fn handle_socket(mut socket: WebSocket, State(state): State, cli tokio::select! { Some(message) = receiver.recv() => { let message = ws::Message::Binary(message.into()); + tracing::debug!("[{}] Sending message through websocket", client_id); if let Err(error) = socket.send(message).await { tracing::error!("[{}] Error sending message through websocket: {}", client_id, error); //TODO remove channel to avoid memory leak. It is the clients responsibility to reestablish a connection @@ -177,8 +192,13 @@ async fn handle_socket(mut socket: WebSocket, State(state): State, cli }, // We only use the socket unidirectional for now // but we want to know when the client closes the socket - Some(Ok(_)) = socket.recv() => {}, - else => break, + Some(Ok(message)) = socket.recv() => { + tracing::debug!("[{}] Received message through websocket: {:?}", client_id, message); + }, + else => { + tracing::debug!("[{}] Handling else", client_id); + break; + }, } } -- 2.51.2