diff --git a/Cargo.lock b/Cargo.lock index 14a66d9e5..6ee97b4aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4738,9 +4738,9 @@ dependencies = [ [[package]] name = "serror" -version = "0.5.1" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bee3f9a0b6dee3c0e8c581dddc560f66f20a1c5c20f6609e827167265ac84eb6" +checksum = "24c14f04906667a9c590474a470a92fa2c76189d4afe186017329e9cb22f1d6a" dependencies = [ "anyhow", "axum", diff --git a/Cargo.toml b/Cargo.toml index 9670191de..72be7a0b6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,7 +37,7 @@ git = { path = "lib/git" } # MOGH run_command = { version = "0.0.6", features = ["async_tokio"] } -serror = { version = "0.5.1", default-features = false } +serror = { version = "0.5.2", default-features = false } slack = { version = "0.4.0", package = "slack_client_rs", default-features = false, features = ["rustls"] } derive_default_builder = "0.1.8" derive_empty_traits = "0.1.0" diff --git a/bin/periphery/src/connection.rs b/bin/periphery/src/connection.rs index 605009356..71b1af2c9 100644 --- a/bin/periphery/src/connection.rs +++ b/bin/periphery/src/connection.rs @@ -8,13 +8,13 @@ use axum::{ use bytes::Bytes; use futures::{SinkExt, StreamExt}; use periphery_client::message::{ - MessageState, from_transport_bytes, to_transport_bytes, + MessageState, parse_from_transport_bytes, to_transport_bytes, }; use resolver_api::Resolve; use response::JsonBytes; use serror::{AddStatusCode, serialize_error_bytes}; use tokio::sync::{Mutex, mpsc::Sender}; -use transport::{BufferedReceiver, buffered_channel}; +use transport::channel::{BufferedReceiver, buffered_channel}; use crate::api::PeripheryRequest; @@ -90,7 +90,7 @@ pub async fn inbound_connection( fn handle_msg(msg: Bytes) { tokio::spawn(async move { let (id, request) = - match from_transport_bytes::(&msg) { + match parse_from_transport_bytes::(&msg) { Ok((id, _, Some(request))) => (id, request), Ok(_) => { // No data, ignore @@ -151,9 +151,9 @@ fn response_sender() -> &'static Sender> { .expect("response_sender accessed before initialized") } -static RESPONSE_RECEIVER: OnceLock> = +static RESPONSE_RECEIVER: OnceLock>>> = OnceLock::new(); -fn response_receiver() -> &'static Mutex { +fn response_receiver() -> &'static Mutex>> { RESPONSE_RECEIVER .get() .expect("response_receiver accessed before initialized") diff --git a/client/periphery/rs/src/connection.rs b/client/periphery/rs/src/connection.rs new file mode 100644 index 000000000..da96697f3 --- /dev/null +++ b/client/periphery/rs/src/connection.rs @@ -0,0 +1,227 @@ +use std::{sync::OnceLock, time::Duration}; + +use anyhow::{Context, anyhow}; +use cache::CloneCache; +use futures_util::{SinkExt, StreamExt}; +use serror::deserialize_error_bytes; +use tokio::sync::mpsc; +use tokio_tungstenite::tungstenite::Message; +use tokio_util::{bytes::Bytes, sync::CancellationToken}; +use tracing::warn; +use transport::channel::{BufferedReceiver, buffered_channel}; +use uuid::Uuid; + +use crate::message::{ + MessageState, from_transport_bytes, id_from_transport_bytes, + to_transport_bytes, +}; + +#[derive(Debug)] +struct Connection { + cancel: CancellationToken, + request_sender: mpsc::Sender>, +} + +impl Connection { + fn new() + -> (Connection, BufferedReceiver>, CancellationToken) { + let (request_sender, request_receiver) = buffered_channel(1000); + let cancel = CancellationToken::new(); + ( + Connection { + cancel: cancel.clone(), + request_sender, + }, + request_receiver, + cancel, + ) + } +} + +impl Clone for Connection { + fn clone(&self) -> Self { + Self { + cancel: self.cancel.clone(), + request_sender: self.request_sender.clone(), + } + } +} + +/// server id => connection +type Connections = CloneCache; +fn connections() -> &'static Connections { + static CONNECTIONS: OnceLock = OnceLock::new(); + CONNECTIONS.get_or_init(Default::default) +} + +type ResponseChannels = CloneCache>>; +fn response_channels() -> &'static ResponseChannels { + static RESPONSE_CHANNELS: OnceLock = + OnceLock::new(); + RESPONSE_CHANNELS.get_or_init(Default::default) +} + +pub async fn send_request( + server_id: &String, + data: &[u8], +) -> anyhow::Result> { + let request_sender = &connections() + .get(server_id) + .await + .with_context(|| { + format!("No connection found for server {server_id}") + })? + .request_sender; + let id = Uuid::new_v4(); + let (response_sender, mut response_receiever) = mpsc::channel(1000); + response_channels().insert(id, response_sender).await; + if let Err(e) = request_sender + .send(to_transport_bytes(id, MessageState::InProgress, data)) + .await + .context("Failed to send request over channel") + { + // cleanup + response_channels().remove(&id).await; + return Err(e); + } + + // Poll for the response + loop { + let bytes = match response_receiever.recv().await { + Some(bytes) => bytes, + None => { + return Err(anyhow!( + "Sender dropped before response was recieved" + )); + } + }; + let (state, data) = match from_transport_bytes(&bytes) { + Ok((_, state, Some(data))) => (state, data), + // TODO: Handle no data cases + Ok(_) => continue, + Err(e) => { + warn!("Received invalid message | {e:#}"); + continue; + } + }; + match state { + // TODO: improve the allocation in .to_vec + MessageState::Successful => return Ok(data.to_vec()), + MessageState::Failed => { + return Err(deserialize_error_bytes(&data)); + } + // TODO: bail if this isn't received frequent enough + MessageState::InProgress => continue, + } + } +} + +// Assumes address already wss formatted +pub async fn spawn_connection( + server_id: String, + address: String, +) -> anyhow::Result<()> { + let response_channels = response_channels(); + let (channel, mut request_receiver, cancel) = Connection::new(); + if let Some(existing) = + connections().insert(server_id, channel).await + { + existing.cancel.cancel(); + } + + tokio::spawn(async move { + let connection = async { + // Outer connection loop + loop { + let socket = + match crate::ws::connect_websocket(&address).await { + Ok(socket) => socket, + Err(e) => { + warn!("{e:#}"); + tokio::time::sleep(Duration::from_secs(5)).await; + continue; + } + }; + + let (mut ws_write, mut ws_read) = socket.split(); + + let forward_requests = async { + loop { + let message = match request_receiver.recv().await { + None => break, + Some(request) => { + Message::Binary(Bytes::copy_from_slice(request)) + } + }; + match ws_write.send(message).await { + Ok(_) => request_receiver.clear_buffer(), + Err(e) => { + warn!("Failed to send request to {address} | {e:#}"); + break; + } + } + } + }; + + let read_responses = async { + loop { + let bytes = match ws_read.next().await { + Some(Ok(Message::Binary(bytes))) => bytes, + Some(Ok(Message::Close(frame))) => { + warn!( + "Connection to {address} broken with frame: {frame:?}" + ); + break; + } + Some(Err(e)) => { + warn!( + "Connection to {address} broken with error: {e:?}" + ); + break; + } + None => { + warn!("Connection to {address} closed"); + break; + } + // Can ignore other message types + Some(Ok(_)) => { + continue; + } + }; + let id = match id_from_transport_bytes(&bytes) { + Ok(res) => res, + Err(e) => { + warn!("Failed to read id from {address} | {e:#}"); + continue; + } + }; + let Some(channel) = response_channels.get(&id).await + else { + warn!( + "Failed to send response for {address} | No response channel found" + ); + continue; + }; + if let Err(e) = channel.send(bytes.into()).await { + warn!( + "Failed to send response for {address} | Channel failure | {e:#}" + ); + } + } + }; + + tokio::select! { + _ = forward_requests => {}, + _ = read_responses => {} + }; + } + }; + + tokio::select! { + _ = connection => {}, + _ = cancel.cancelled() => {} + } + }); + + Ok(()) +} diff --git a/client/periphery/rs/src/connections.rs b/client/periphery/rs/src/connections.rs deleted file mode 100644 index c239ac798..000000000 --- a/client/periphery/rs/src/connections.rs +++ /dev/null @@ -1,91 +0,0 @@ -use std::sync::OnceLock; - -use cache::CloneCache; -use futures_util::StreamExt; -use tokio::sync::{broadcast, mpsc}; -use tokio_util::sync::CancellationToken; -use transport::{BufferedReceiver, buffered_channel}; - -#[derive(Debug)] -struct ConnectionChannel { - request_sender: mpsc::Sender>, - response_receiver: broadcast::Receiver>, - cancel: CancellationToken, -} - -impl ConnectionChannel { - fn new() -> ( - ConnectionChannel, - BufferedReceiver, - broadcast::Sender>, - CancellationToken, - ) { - let (request_sender, request_receiver) = buffered_channel(1000); - let (response_sender, response_receiver) = - broadcast::channel(1000); - let cancel = CancellationToken::new(); - ( - ConnectionChannel { - request_sender, - response_receiver, - cancel: cancel.clone(), - }, - request_receiver, - response_sender, - cancel, - ) - } -} - -impl Clone for ConnectionChannel { - fn clone(&self) -> Self { - Self { - request_sender: self.request_sender.clone(), - response_receiver: self.response_receiver.resubscribe(), - cancel: self.cancel.clone(), - } - } -} - -type ConnectionCache = CloneCache; -fn connections() -> &'static ConnectionCache { - static CONNECTIONS: OnceLock = OnceLock::new(); - CONNECTIONS.get_or_init(Default::default) -} - -// Assumes address already wss formatted -async fn spawn_connection(address: String) -> anyhow::Result<()> { - let (channel, request_receiver, response_sender, cancel) = - ConnectionChannel::new(); - if let Some(existing) = - connections().insert(address.clone(), channel).await - { - existing.cancel.cancel(); - } - - tokio::spawn(async move { - // Outer connection loop - loop { - let socket = match crate::ws::connect_websocket(&address).await - { - Ok(socket) => socket, - Err(e) => { - // TODO: handle connect error - return; - } - }; - let (mut ws_read, mut ws_write) = socket.split(); - - let forward_requests = async { - loop { - // match request_receiver.recv().await { - // None => break, - - // } - } - }; - } - }); - - Ok(()) -} diff --git a/client/periphery/rs/src/lib.rs b/client/periphery/rs/src/lib.rs index b35d22199..13248f0b2 100644 --- a/client/periphery/rs/src/lib.rs +++ b/client/periphery/rs/src/lib.rs @@ -7,7 +7,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::json; pub mod api; -pub mod connections; +pub mod connection; pub mod message; mod terminal; diff --git a/client/periphery/rs/src/message.rs b/client/periphery/rs/src/message.rs index 622fe0627..ab6a89e8c 100644 --- a/client/periphery/rs/src/message.rs +++ b/client/periphery/rs/src/message.rs @@ -19,19 +19,37 @@ pub fn to_transport_bytes( res } +/// Deserializes channel id from +/// incoming transport bytes. +pub fn id_from_transport_bytes(bytes: &[u8]) -> anyhow::Result { + if bytes.len() < 16 { + return Err(anyhow!("Transport bytes too short to include uuid")); + } + Uuid::from_slice(&bytes[..16]).context("Invalid Uuid bytes") +} + /// Deserializes channel id + data from /// incoming transport bytes. -pub fn from_transport_bytes( +pub fn from_transport_bytes( bytes: &[u8], -) -> anyhow::Result<(Uuid, MessageState, Option)> { +) -> anyhow::Result<(Uuid, MessageState, Option<&[u8]>)> { if bytes.len() < 17 { return Err(anyhow!( - "Transport bytes too short to include uuid at state" + "Transport bytes too short to include uuid and state" )); } let (id, state, data) = (&bytes[..16], bytes[16], bytes.get(17..)); let id = Uuid::from_slice(id).context("Invalid Uuid bytes")?; let state = MessageState::from_byte(state); + Ok((id, state, data)) +} + +/// Deserializes channel id + data from +/// incoming transport bytes. +pub fn parse_from_transport_bytes( + bytes: &[u8], +) -> anyhow::Result<(Uuid, MessageState, Option)> { + let (id, state, data) = from_transport_bytes(bytes)?; let data = data .map(|data| { serde_json::from_slice::(data) diff --git a/lib/cache/src/lib.rs b/lib/cache/src/lib.rs index e1a3ea9af..c7476b1ee 100644 --- a/lib/cache/src/lib.rs +++ b/lib/cache/src/lib.rs @@ -59,6 +59,7 @@ fn clone_anyhow_error(e: &anyhow::Error) -> anyhow::Error { e } +#[derive(Debug)] pub struct CloneCache { cache: RwLock>, } diff --git a/lib/transport/src/channel.rs b/lib/transport/src/channel.rs new file mode 100644 index 000000000..44d32ea15 --- /dev/null +++ b/lib/transport/src/channel.rs @@ -0,0 +1,48 @@ +use std::ops::Deref; + +use tokio::sync::mpsc; + +/// Create a buffered channel +pub fn buffered_channel( + buffer: usize, +) -> (mpsc::Sender, BufferedReceiver) { + let (sender, receiver) = mpsc::channel(buffer); + (sender, BufferedReceiver::new(receiver)) +} + +/// Wrapper around channel receiver to control when +/// the latest message is dropped, +/// in case it must be re-transmitted. +#[derive(Debug)] +pub struct BufferedReceiver { + receiver: mpsc::Receiver, + buffer: Option, +} + +impl BufferedReceiver { + pub fn new(receiver: mpsc::Receiver) -> BufferedReceiver { + BufferedReceiver { + receiver, + buffer: None, + } + } + + /// If 'next: Some(bytes)': + /// - Immediately returns borrow of next. + /// Else: + /// - Wait for next item + /// - store in 'next' + /// - return borrow of next. + pub async fn recv(&mut self) -> Option<&::Target> { + if self.buffer.is_none() { + self.buffer = Some(self.receiver.recv().await?); + } + self.buffer.as_deref() + } + + /// Clears buffer. + /// Should be called after transmission confirmed. + pub fn clear_buffer(&mut self) { + self.buffer = None; + } +} diff --git a/lib/transport/src/lib.rs b/lib/transport/src/lib.rs index 77b1d97e6..ff4a7365e 100644 --- a/lib/transport/src/lib.rs +++ b/lib/transport/src/lib.rs @@ -1,46 +1 @@ -use tokio::sync::mpsc; - -/// Create a buffered channel -pub fn buffered_channel( - buffer: usize, -) -> (mpsc::Sender>, BufferedReceiver) { - let (sender, receiver) = mpsc::channel(buffer); - (sender, BufferedReceiver::new(receiver)) -} - -/// Wrapper around channel receiver to control when -/// the latest message is dropped, -/// in case it must be re-transmitted. -#[derive(Debug)] -pub struct BufferedReceiver { - receiver: mpsc::Receiver>, - buffer: Option>, -} - -impl BufferedReceiver { - pub fn new(receiver: mpsc::Receiver>) -> BufferedReceiver { - BufferedReceiver { - receiver, - buffer: None, - } - } - - /// If 'next: Some(bytes)': - /// - Immediately returns borrow of next. - /// Else: - /// - Wait for next item - /// - store in 'next' - /// - return borrow of next. - pub async fn recv(&mut self) -> Option<&[u8]> { - if self.buffer.is_none() { - self.buffer = Some(self.receiver.recv().await?); - } - self.buffer.as_deref() - } - - /// Clears buffer. - /// Should be called after transmission confirmed. - pub fn clear_buffer(&mut self) { - self.buffer = None; - } -} +pub mod channel; \ No newline at end of file