feat: Remove hyperdecks, bundle client

This commit is contained in:
2024-05-22 01:20:03 +01:00
parent f9cac35a29
commit 8cce21193a
19 changed files with 3789 additions and 2575 deletions
+212
View File
@@ -0,0 +1,212 @@
use axum::extract::ws::Message;
use axum::extract::{State, WebSocketUpgrade};
use axum::response::Html;
use axum::Json;
use axum::{
body::Bytes,
extract::Path,
http::{header, HeaderValue, Method},
response::IntoResponse,
routing::get,
Router,
};
use message::{ClientRequest, HyperdeckMonitorState, ServerEvent};
use serde::{Deserialize, Serialize};
use std::{
collections::HashMap,
net::{Ipv4Addr, SocketAddr},
sync::Arc,
time::Duration,
};
use tokio::sync::{Mutex, RwLock};
use tower::ServiceBuilder;
use tower_http::timeout::TimeoutLayer;
use tower_http::ServiceBuilderExt;
use tower_http::{
cors::{Any, CorsLayer},
trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
LatencyUnit,
};
use tracing::info;
use uuid::Uuid;
pub mod message;
mod ws;
const FILE_NAME_INDEX: &str = env!("FILE_NAME_INDEX");
const FILE_NAME_WASM: &str = env!("FILE_NAME_WASM");
const FILE_NAME_JS: &str = env!("FILE_NAME_JS");
const FILE_NAME_MANIFEST: &str = env!("FILE_NAME_MANIFEST");
const FILE_NAME_SERVICE_WORKER: &str = env!("FILE_NAME_SERVICE_WORKER");
#[derive(Debug, Clone)]
pub struct Client {
pub sender: Option<tokio::sync::broadcast::Sender<Message>>,
}
type Clients = Arc<Mutex<HashMap<Uuid, Client>>>;
pub async fn initialize_api(
mut state_rx: tokio::sync::broadcast::Receiver<HyperdeckMonitorState>,
client_request_tx: tokio::sync::mpsc::UnboundedSender<ClientRequest>,
) {
info!("Initializing API");
let clients: Clients = Default::default();
let state = Arc::new(RwLock::new(state_rx.recv().await.unwrap()));
let state_clients = clients.clone();
let state_loop = state.clone();
tokio::spawn(async move {
loop {
if let Ok(hyperdeck_monitor_state) = state_rx.recv().await {
let mut state = state_loop.write().await;
*state = hyperdeck_monitor_state.clone();
let clients = state_clients.lock().await;
let state_json = serde_json::to_string(&ServerEvent::HyperdeckMonitorState(
hyperdeck_monitor_state.into(),
))
.unwrap();
for (_, client) in clients.iter() {
if let Some(sender) = &client.sender {
let message: Message = Message::Text(state_json.clone());
let _ = sender.send(message);
}
}
}
}
});
let app_state = AppState {
state,
client_request_tx,
clients,
port: 9681,
};
let addr = SocketAddr::from((Ipv4Addr::UNSPECIFIED, app_state.port));
info!("Listening on {}", addr);
// TODO: This could fail, need to figure out how to get a result from this
let _ = axum::Server::bind(&addr)
.serve(app(app_state).into_make_service())
.await;
}
#[derive(Clone)]
struct AppState {
state: Arc<RwLock<HyperdeckMonitorState>>,
client_request_tx: tokio::sync::mpsc::UnboundedSender<ClientRequest>,
clients: Clients,
port: u16,
}
fn app(state: AppState) -> Router {
let sensitive_headers: Arc<[_]> = vec![header::AUTHORIZATION, header::COOKIE].into();
let middleware = ServiceBuilder::new()
// Mark the `Authorization` and `Cookie` headers as sensitive so it doesn't show in logs
.sensitive_request_headers(sensitive_headers.clone())
// Add high level tracing/logging to all requests
.layer(
TraceLayer::new_for_http()
.on_body_chunk(|chunk: &Bytes, latency: Duration, _: &tracing::Span| {
tracing::trace!(size_bytes = chunk.len(), latency = ?latency, "sending body chunk")
})
.make_span_with(DefaultMakeSpan::new().include_headers(true))
.on_response(DefaultOnResponse::new().include_headers(true).latency_unit(LatencyUnit::Micros)),
)
.sensitive_response_headers(sensitive_headers)
// Set a timeout
.layer(TimeoutLayer::new(Duration::from_secs(10)))
// Box the response body so it implements `Default` which is required by axum
.map_response_body(axum::body::boxed)
// Compress responses
.compression()
// Set a `Content-Type` if there isn't one already.
.insert_response_header_if_not_present(
header::CONTENT_TYPE,
HeaderValue::from_static("application/octet-stream"),
);
let cors = CorsLayer::new()
.allow_methods(vec![
Method::GET,
Method::POST,
Method::PUT,
Method::DELETE,
Method::OPTIONS,
])
.allow_headers(Any)
.allow_origin(Any)
.allow_credentials(false);
Router::new()
.route("/", get(get_index))
.route(&format!("/{}", FILE_NAME_INDEX), get(get_index))
.route(&format!("/{}", FILE_NAME_WASM), get(get_wasm))
.route(&format!("/{}", FILE_NAME_JS), get(get_js))
.route(&format!("/{}", FILE_NAME_MANIFEST), get(get_manifest))
.route(
&format!("/{}", FILE_NAME_SERVICE_WORKER),
get(get_service_worker),
)
.route("/ws", get(upgrade_ws))
.layer(middleware)
.layer(cors)
.with_state(state)
}
#[derive(Debug, Serialize, Deserialize)]
pub struct WebSocketUpgradeRequest {}
async fn get_index() -> Html<String> {
Html(include_str!(env!("INCLUDE_PATH_INDEX")).to_string())
}
async fn get_wasm() -> impl IntoResponse {
(
[(header::CONTENT_TYPE, "application/wasm")],
include_bytes!(env!("INCLUDE_PATH_WASM")),
)
}
async fn get_js() -> impl IntoResponse {
(
[(header::CONTENT_TYPE, "text/javascript")],
include_str!(env!("INCLUDE_PATH_JS")),
)
}
async fn get_manifest() -> Json<String> {
Json(include_str!(env!("INCLUDE_PATH_MANIFEST")).to_string())
}
async fn get_service_worker() -> impl IntoResponse {
(
[(header::CONTENT_TYPE, "text/javascript")],
include_str!(env!("INCLUDE_PATH_SERVICE_WORKER")),
)
}
#[axum::debug_handler]
async fn upgrade_ws(state: State<AppState>, ws: WebSocketUpgrade) -> impl IntoResponse {
info!("New client websocket connection");
let client_id = uuid::Uuid::new_v4();
state
.clients
.lock()
.await
.insert(client_id, Client { sender: None });
let client = state.clients.lock().await.get(&client_id).cloned().unwrap();
ws.on_upgrade(move |socket| {
ws::client_connection(
state.client_request_tx.clone(),
socket,
client_id,
state.state.clone(),
state.clients.clone(),
client,
)
})
}
+49
View File
@@ -0,0 +1,49 @@
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[serde(tag = "type")]
pub enum ClientRequest {
AddHyperdeck(AddHyperdeckRequest),
RemoveHyperdeck(RemoveHyperdeckRequest),
}
#[derive(Debug, Serialize, Deserialize)]
pub struct AddHyperdeckRequest {
pub name: String,
pub ip: String,
pub port: u16,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct RemoveHyperdeckRequest {
pub id: String,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[serde(tag = "type")]
pub enum ServerEvent {
HyperdeckMonitorState(HyperdeckMonitorState),
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct HyperdeckMonitorState {
pub hyperdecks: HashMap<String, HyperdeckState>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HyperdeckState {
pub name: String,
pub ip: String,
pub port: u16,
pub connection_state: HyperdeckConnectionState,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum HyperdeckConnectionState {
Connected,
Disconnected,
}
+84
View File
@@ -0,0 +1,84 @@
use std::{future, sync::Arc};
use super::message::{ClientRequest, HyperdeckMonitorState};
use crate::api::ServerEvent;
use axum::extract::ws::{Message, WebSocket};
use futures::StreamExt;
use tokio::sync::RwLock;
use tokio_stream::wrappers::BroadcastStream;
use tracing::{debug, error, log::info};
use uuid::Uuid;
use super::{Client, Clients};
pub async fn client_connection(
client_request_tx: tokio::sync::mpsc::UnboundedSender<ClientRequest>,
ws: WebSocket,
id: Uuid,
state: Arc<RwLock<HyperdeckMonitorState>>,
clients: Clients,
mut client: Client,
) {
let (client_ws_sender, mut client_ws_rcv) = ws.split();
let (client_sender, client_rcv) = tokio::sync::broadcast::channel::<Message>(10);
let client_rcv = BroadcastStream::new(client_rcv);
tokio::task::spawn(
client_rcv
.filter(|msg| future::ready(msg.is_ok()))
.map(|msg| Ok(msg.unwrap()))
.forward(client_ws_sender),
);
let current_state = state.read().await.clone();
let state_json =
serde_json::to_string(&ServerEvent::HyperdeckMonitorState(current_state.into())).unwrap();
client_sender.send(Message::Text(state_json.clone())).ok();
client.sender = Some(client_sender);
clients.lock().await.insert(id, client);
info!("{} connected", id);
while let Some(result) = client_ws_rcv.next().await {
let msg = match result {
Ok(msg) => msg,
Err(e) => {
error!("error resolving ws message for id: {}: {}", id.clone(), e);
break;
}
};
client_msg(client_request_tx.clone(), &id, msg).await;
}
clients.lock().await.remove(&id);
info!("{} disconnected", id);
}
async fn client_msg(
client_request_tx: tokio::sync::mpsc::UnboundedSender<ClientRequest>,
id: &Uuid,
msg: Message,
) {
debug!("received message from {}: {:?}", id, msg);
let message = match msg.into_text() {
Ok(v) => v,
Err(err) => {
error!("error: {:?}", err);
return;
}
};
if message == "ping" || message == "ping\n" {
return;
}
let client_request: super::message::ClientRequest = match serde_json::from_str(&message) {
Ok(v) => v,
Err(_) => {
return;
}
};
let _ = client_request_tx.send(client_request);
}
+337
View File
@@ -0,0 +1,337 @@
use std::{process::Stdio, time::Duration};
use api::message::{
AddHyperdeckRequest, ClientRequest, HyperdeckConnectionState, HyperdeckMonitorState,
HyperdeckState, RemoveHyperdeckRequest,
};
use color_eyre::Report;
use futures_util::{
pin_mut, select,
stream::{SplitSink, SplitStream},
FutureExt, SinkExt, StreamExt,
};
use serde::{Deserialize, Serialize};
use tokio::net::TcpStream;
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
use tokio_util::{
codec::{FramedRead, LinesCodec},
sync::CancellationToken,
};
use tracing_subscriber::EnvFilter;
mod api;
#[tokio::main]
async fn main() {
setup_logging().expect("Failed to setup logging");
tracing::info!("Hello, world!");
let cancel = CancellationToken::new();
let node_process = run_node_process(cancel.clone()).fuse();
let (node_ws_message_tx, node_ws_message_rx) = tokio::sync::mpsc::unbounded_channel();
let (node_commands_tx, node_commands_rx) = tokio::sync::mpsc::unbounded_channel();
let state = AppState::default();
let node_ws_communication =
talk_to_node_ws(state, node_ws_message_tx, node_commands_rx, cancel.clone()).fuse();
let (state_tx, state_rx) = tokio::sync::broadcast::channel(1);
let (client_request_tx, client_request_rx) = tokio::sync::mpsc::unbounded_channel();
let api = api::initialize_api(state_rx, client_request_tx).fuse();
let hyperdeck_monitor = run(
node_commands_tx,
node_ws_message_rx,
state_tx,
client_request_rx,
cancel.clone(),
)
.fuse();
pin_mut!(node_process);
pin_mut!(node_ws_communication);
pin_mut!(api);
pin_mut!(hyperdeck_monitor);
select! {
_ = node_process => {},
_ = node_ws_communication => {},
_ = api => {},
_ = hyperdeck_monitor => {},
_ = cancel.cancelled().fuse() => {}
};
cancel.cancel();
}
async fn run(
mut node_commands_tx: tokio::sync::mpsc::UnboundedSender<NodeWsCommand>,
mut node_ws_message_rx: tokio::sync::mpsc::UnboundedReceiver<NodeWsMessageReceived>,
mut state_tx: tokio::sync::broadcast::Sender<HyperdeckMonitorState>,
mut client_request_rx: tokio::sync::mpsc::UnboundedReceiver<ClientRequest>,
cancel: CancellationToken,
) {
let mut state = HyperdeckMonitorState::default();
let _ = state_tx.send(state.clone());
while !cancel.is_cancelled() {
let state_modified = select! {
message_from_node = node_ws_message_rx.recv().fuse() => {
if let Some(msg) = message_from_node {
handle_message_from_node(msg, &mut node_commands_tx, &mut state).await
} else {
false
}
},
message_from_client = client_request_rx.recv().fuse() => {
if let Some(msg) = message_from_client {
handle_message_from_client(msg, &mut node_commands_tx, &mut state).await
} else {
false
}
}
};
if state_modified {
let _ = state_tx.send(state.clone());
}
}
}
async fn handle_message_from_node(
msg: NodeWsMessageReceived,
node_commands_tx: &mut tokio::sync::mpsc::UnboundedSender<NodeWsCommand>,
state: &mut HyperdeckMonitorState,
) -> bool {
match msg {
NodeWsMessageReceived::Log { message } => {
tracing::info!("[NODE] {message}");
false
}
NodeWsMessageReceived::HyperdeckConnected { id } => {
state.hyperdecks.entry(id).and_modify(|hyperdeck| {
hyperdeck.connection_state = HyperdeckConnectionState::Connected
});
true
}
NodeWsMessageReceived::HypderdeckDisconnected { id } => {
state.hyperdecks.entry(id).and_modify(|hyperdeck| {
hyperdeck.connection_state = HyperdeckConnectionState::Disconnected
});
true
}
}
}
async fn handle_message_from_client(
msg: ClientRequest,
node_commands_tx: &mut tokio::sync::mpsc::UnboundedSender<NodeWsCommand>,
state: &mut HyperdeckMonitorState,
) -> bool {
match msg {
ClientRequest::AddHyperdeck(AddHyperdeckRequest { name, ip, port }) => {
tracing::info!("Adding hyperdeck");
let id = uuid::Uuid::new_v4();
state.hyperdecks.insert(
id.to_string(),
HyperdeckState {
name,
ip: ip.clone(),
port,
connection_state: api::message::HyperdeckConnectionState::Disconnected,
},
);
let _ = node_commands_tx.send(NodeWsCommand::AddHyperdeck(AddHyperdeckCommand {
id: id.to_string(),
ip,
port,
}));
true
}
ClientRequest::RemoveHyperdeck(RemoveHyperdeckRequest { id }) => {
let _ = state.hyperdecks.remove(&id);
let _ = node_commands_tx.send(NodeWsCommand::RemoveHyperdeck(RemoveHyperdeckCommand {
id,
}));
true
}
}
}
async fn run_node_process(cancel: CancellationToken) {
while !cancel.is_cancelled() {
// Back-off in case we are immediately crashing in a loop.
tokio::time::sleep(Duration::from_secs(1)).await;
let result = tokio::process::Command::new("node")
.arg("./index.js")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn();
match result {
Ok(mut child_process) => {
let Some(raw_stdout) = child_process.stdout.take() else {
let _ = child_process.kill().await;
continue;
};
let Some(raw_stderr) = child_process.stderr.take() else {
let _ = child_process.kill().await;
continue;
};
let mut stdout = FramedRead::new(raw_stdout, LinesCodec::new())
.map(|data| data.expect("Could not read stdout"));
let mut stderr = FramedRead::new(raw_stderr, LinesCodec::new())
.map(|data| data.expect("Could not read stderr"));
while !cancel.is_cancelled() {
select! {
line = stdout.next().fuse() => {
if let Some(line) = line {
tracing::info!("[NODE] {line}");
}
}
line = stderr.next().fuse() => {
if let Some(line) = line {
tracing::error!("[NODE] {line}");
}
}
}
}
let _ = child_process.kill().await;
}
Err(err) => {
tracing::error!("Error running Node child process: {err}");
}
}
}
}
#[derive(Default)]
struct AppState {}
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
enum NodeWsCommand {
#[serde(rename = "add_hyperdeck")]
AddHyperdeck(AddHyperdeckCommand),
#[serde(rename = "remove_hyperdeck")]
RemoveHyperdeck(RemoveHyperdeckCommand),
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(tag = "event")]
enum NodeWsMessageReceived {
Log { message: String },
HyperdeckConnected { id: String },
HypderdeckDisconnected { id: String },
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct AddHyperdeckCommand {
id: String,
ip: String,
port: u16,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct RemoveHyperdeckCommand {
id: String,
}
async fn talk_to_node_ws(
state: AppState,
ws_message_tx: tokio::sync::mpsc::UnboundedSender<NodeWsMessageReceived>,
commands_rx: tokio::sync::mpsc::UnboundedReceiver<NodeWsCommand>,
cancel: CancellationToken,
) {
let ws_stream = wait_for_connection().await;
let (write, read) = ws_stream.split();
let outgoing = handle_outbound_messages(commands_rx, write).fuse();
let incoming = handle_inbound_messages(read, ws_message_tx).fuse();
pin_mut!(outgoing);
pin_mut!(incoming);
select! {
_ = outgoing => {},
_ = incoming => {},
_ = cancel.cancelled().fuse() => {},
}
}
async fn wait_for_connection() -> WebSocketStream<MaybeTlsStream<TcpStream>> {
loop {
// Wait for Node to wake up...
tokio::time::sleep(Duration::from_secs(1)).await;
let ws_url = url::Url::parse("ws://127.0.0.1:7867").expect("Invalid websocket URL");
match tokio_tungstenite::connect_async(ws_url.clone()).await {
Ok((ws_stream, _)) => {
tracing::info!("Connected to Node process on {ws_url}");
return ws_stream;
}
Err(err) => {
tracing::error!("Error connecting to Node process: {:?}", err)
}
}
}
}
async fn handle_outbound_messages(
mut commands_rx: tokio::sync::mpsc::UnboundedReceiver<NodeWsCommand>,
mut socket_tx: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
) {
while let Some(command) = commands_rx.recv().await {
if let Err(err) = socket_tx
.send(tokio_tungstenite::tungstenite::Message::Text(
serde_json::to_string(&command).expect("Could not serialize command"),
))
.await
{
tracing::error!("Error sending command to Node proccess: {err}");
}
}
}
async fn handle_inbound_messages(
socket_rx: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
ws_message_tx: tokio::sync::mpsc::UnboundedSender<NodeWsMessageReceived>,
) {
socket_rx
.for_each(|message| async {
match message {
Ok(tokio_tungstenite::tungstenite::Message::Text(text)) => {
if let Ok(received) = serde_json::from_str::<NodeWsMessageReceived>(&text) {
let _ = ws_message_tx.send(received);
}
}
Ok(tokio_tungstenite::tungstenite::Message::Pong(_)) => {}
_ => {}
}
})
.await;
}
fn setup_logging() -> Result<(), Report> {
if std::env::var("RUST_LIB_BACKTRACE").is_err() {
std::env::set_var("RUST_LIB_BACKTRACE", "1");
}
color_eyre::install()?;
if std::env::var("RUST_LOG").is_err() {
std::env::set_var("RUST_LOG", "debug");
}
tracing_subscriber::fmt::fmt()
.with_env_filter(EnvFilter::from_default_env())
.init();
Ok(())
}