Merge pull request #10 from AOSC-Dev/buildit-monitor

feat(worker): stream compile log to buildit-monitor

server/Cargo.toml

@@ -24,7 +24,7 @@ buildit-utils = { path = "../buildit-utils" }
 jsonwebtoken = "9.2.0"
 size = "0.4.1"
 dickens = { git = "https://github.com/AOSC-Dev/dickens.git", version = "0.1.0" }
-axum = "0.7.4"
+axum = { version = "0.7.4", features = ["ws"] }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 tracing = "0.1.40"
 tower-http = { version = "0.5.2", features = ["trace", "fs", "cors"] }

server/src/main.rs

@@ -15,13 +15,15 @@ use server::bot::{answer, Command};
 use server::recycler::recycler_worker;
 use server::routes::{
     dashboard_status, job_info, job_list, job_restart, ping, pipeline_info, pipeline_list,
-    pipeline_new_pr, worker_info, worker_job_update, worker_list, worker_poll, AppState,
+    pipeline_new_pr, worker_info, worker_job_update, worker_list, worker_poll, ws_handler,
+    AppState, PeerMap,
 };
 use server::routes::{pipeline_new, worker_heartbeat};
 use server::routes::{pipeline_status, worker_status};
 use server::{DbPool, ARGS};
+use std::collections::HashMap;
 use std::os::unix::fs::PermissionsExt;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use teloxide::prelude::*;
 use tokio::net::unix::UCred;
 use tokio::net::UnixStream;
@@ -99,7 +101,9 @@ async fn main() -> anyhow::Result<()> {
     let state = AppState {
         pool: pool.clone(),
         bot,
+        ws_peer_map: PeerMap::new(RwLock::new(HashMap::new())),
     };
+
     let mut app = Router::new()
         .route("/api/ping", get(ping))
         .route("/api/pipeline/new", post(pipeline_new))
@@ -117,6 +121,7 @@ async fn main() -> anyhow::Result<()> {
         .route("/api/worker/list", get(worker_list))
         .route("/api/worker/info", get(worker_info))
         .route("/api/dashboard/status", get(dashboard_status))
+        .route("/api/ws/:hostname", get(ws_handler))
         .nest_service("/assets", ServeDir::new("frontend/dist/assets"))
         .route_service("/favicon.ico", ServeFile::new("frontend/dist/favicon.ico"))
         .fallback_service(ServeFile::new("frontend/dist/index.html"))

server/src/routes/mod.rs

@@ -13,28 +13,39 @@ use diesel::dsl::{count, sum};
 
 use diesel::{Connection, ExpressionMethods, QueryDsl, RunQueryDsl};
 
+use futures::channel::mpsc::UnboundedSender;
 use serde::Serialize;
-use std::collections::BTreeMap;
+use std::{
+    collections::{BTreeMap, HashMap},
+    net::SocketAddr,
+    sync::{Arc, RwLock},
+};
 
 use teloxide::prelude::*;
 use tracing::info;
 
 pub mod job;
 pub mod pipeline;
+pub mod websocket;
 pub mod worker;
 
 pub use job::*;
 pub use pipeline::*;
+pub use websocket::*;
 pub use worker::*;
 
 pub async fn ping() -> &'static str {
     "PONG"
 }
 
+type Tx = (UnboundedSender<axum::extract::ws::Message>, String);
+pub type PeerMap = Arc<RwLock<HashMap<SocketAddr, Tx>>>;
+
 #[derive(Clone)]
 pub struct AppState {
     pub pool: DbPool,
     pub bot: Option<Bot>,
+    pub ws_peer_map: PeerMap,
 }
 
 // learned from https://github.com/tokio-rs/axum/blob/main/examples/anyhow-error-response/src/main.rs

server/src/routes/websocket.rs

@@ -0,0 +1,67 @@
+use std::net::SocketAddr;
+
+use axum::{
+    extract::{ws::WebSocket, ConnectInfo, Path, State, WebSocketUpgrade},
+    response::IntoResponse,
+};
+use futures::{channel::mpsc::unbounded, future, pin_mut, StreamExt, TryStreamExt};
+use tracing::info;
+
+use super::{AppState, PeerMap};
+
+/// The handler for the HTTP request (this gets called when the HTTP GET lands at the start
+/// of websocket negotiation). After this completes, the actual switching from HTTP to
+/// websocket protocol will occur.
+/// This is the last point where we can extract TCP/IP metadata such as IP address of the client
+/// as well as things from HTTP headers such as user-agent of the browser etc.
+pub async fn ws_handler(
+    Path(hostname): Path<String>,
+    ws: WebSocketUpgrade,
+    ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    State(state): State<AppState>,
+) -> impl IntoResponse {
+    // finalize the upgrade process by returning upgrade callback.
+    // we can customize the callback by sending additional info such as address.
+    ws.on_upgrade(move |socket| handle_socket(socket, addr, hostname, state.ws_peer_map))
+}
+
+/// Actual websocket statemachine (one will be spawned per connection)
+async fn handle_socket(socket: WebSocket, who: SocketAddr, hostname: String, peer_map: PeerMap) {
+    let (tx, rx) = unbounded();
+    peer_map
+        .write()
+        .unwrap()
+        .insert(who, (tx, hostname.clone()));
+
+    let (outgoing, incoming) = socket.split();
+
+    let broadcast_incoming = incoming.try_for_each(|msg| {
+        info!("Received a message from {}: {:?}", who, msg);
+
+        let peers = peer_map.read().unwrap();
+
+        // We want to broadcast the message to everyone except ourselves.
+        let broadcast_recipients = peers
+            .iter()
+            .filter(|(peer_addr, _)| peer_addr != &&who)
+            .map(|(_, (ws_sink, port))| (ws_sink, port));
+
+        for recp in broadcast_recipients {
+            let recp_path = recp.1;
+
+            if *recp_path == hostname {
+                recp.0.unbounded_send(msg.clone()).unwrap();
+            }
+        }
+
+        future::ok(())
+    });
+
+    let receive_from_others = rx.map(Ok).forward(outgoing);
+
+    pin_mut!(broadcast_incoming, receive_from_others);
+    future::select(broadcast_incoming, receive_from_others).await;
+
+    info!("{} disconnected", &who);
+    peer_map.write().unwrap().remove(&who);
+}

server/src/routes/worker.rs

@@ -309,7 +309,7 @@ pub async fn worker_poll(
 }
 
 pub async fn worker_job_update(
-    State(AppState { pool, bot }): State<AppState>,
+    State(AppState { pool, bot, .. }): State<AppState>,
     Json(payload): Json<WorkerJobUpdateRequest>,
 ) -> Result<(), AnyhowError> {
     if payload.worker_secret != ARGS.worker_secret {

worker/Cargo.toml

@@ -20,6 +20,9 @@ num_cpus = "1.16.0"
 reqwest = { version = "0.11.24", features = ["json"] }
 sysinfo = "0.30.5"
 tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread", "process", "sync", "fs"] }
+tokio-tungstenite = "0.23.1"
+futures-util = "0.3.30"
+flume = "0.11.0"
 
 [build-dependencies]
 vergen = { version = "8.3.1", features = ["build", "cargo", "git", "gitcl", "rustc", "si"] }