From e9c2e584b4ee11986d4a872ffe01dfd7024261ea Mon Sep 17 00:00:00 2001 From: yyjeqhc <1772413353@qq.com> Date: Mon, 4 Aug 2025 07:39:38 +0800 Subject: [PATCH] feat: Add local testing API and enhance orion/orion-server with heartbeat and task scheduling 1.Added a new API in mono to enable local testing without login requirement 2.Implemented heartbeat mechanism in orion-server and orion for persistent connection and reconnection 3.Added task scheduling and status checking features in orion-server and orion --- mono/src/api/mr/mr_router.rs | 37 +++ orion-server/.env | 4 +- orion-server/Cargo.toml | 1 + orion-server/src/api.rs | 548 ++++++++++++++++++++----------- orion-server/src/main.rs | 11 +- orion-server/src/model/builds.rs | 18 +- orion-server/src/server.rs | 123 +++++-- orion/.env | 4 +- orion/src/api.rs | 107 +++--- orion/src/buck_controller.rs | 198 ++++++----- orion/src/main.rs | 35 +- orion/src/ws.rs | 308 ++++++++++------- 12 files changed, 906 insertions(+), 488 deletions(-) diff --git a/mono/src/api/mr/mr_router.rs b/mono/src/api/mr/mr_router.rs index a94a00291..0b3656315 100644 --- a/mono/src/api/mr/mr_router.rs +++ b/mono/src/api/mr/mr_router.rs @@ -34,6 +34,7 @@ pub fn routers() -> OpenApiRouter { .routes(routes!(fetch_mr_list)) .routes(routes!(mr_detail)) .routes(routes!(merge)) + .routes(routes!(merge_no_auth)) .routes(routes!(close_mr)) .routes(routes!(reopen_mr)) .routes(routes!(mr_files_changed)) @@ -170,6 +171,42 @@ async fn merge( Ok(Json(CommonResult::success(None))) } +/// Merge Request without authentication +/// It's for local testing purposes. +#[utoipa::path( + post, + params( + ("link", description = "MR link"), + ), + path = "/{link}/merge-no-auth", + responses( + (status = 200, body = CommonResult, content_type = "application/json") + ), + tag = MR_TAG +)] +async fn merge_no_auth( + Path(link): Path, + state: State, +) -> Result>, ApiError> { + let res = state.mr_stg().get_mr(&link).await?; + let model = res.ok_or(MegaError::with_message("MR Not Found"))?; + + if model.status != MergeStatusEnum::Open { + return Err(ApiError::from(MegaError::with_message(format!( + "MR is not in Open status, current status: {:?}", + model.status + )))); + } + + // No authentication required - using default system user + let default_username = "system"; + state.monorepo().merge_mr(default_username, model).await?; + + Ok(Json(CommonResult::success(Some( + "Merge completed successfully".to_string(), + )))) +} + /// Fetch MR list #[utoipa::path( post, diff --git a/orion-server/.env b/orion-server/.env index 24ced463c..e2ccd1d86 100644 --- a/orion-server/.env +++ b/orion-server/.env @@ -1,3 +1,3 @@ -BUILD_LOG_DIR="/tmp/buck2ctl" +BUILD_LOG_DIR="/tmp/megadir/buck2ctl" DATABASE_URL="postgres://postgres:postgres@localhost/orion" -PORT=8004 \ No newline at end of file +PORT=80 \ No newline at end of file diff --git a/orion-server/Cargo.toml b/orion-server/Cargo.toml index c80bdd567..e6f312a43 100644 --- a/orion-server/Cargo.toml +++ b/orion-server/Cargo.toml @@ -29,3 +29,4 @@ dashmap = { workspace = true } utoipa-axum.workspace = true utoipa.workspace = true utoipa-swagger-ui = { workspace = true, features = ["axum"] } +chrono = { version = "0.4", features = ["serde"] } \ No newline at end of file diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index a8866db61..c73190cd5 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,23 +1,26 @@ use crate::model::builds; -use axum::Json; -use axum::body::Bytes; -use axum::extract::ws::{Message, Utf8Bytes, WebSocket}; -use axum::extract::{ConnectInfo, Path, State, WebSocketUpgrade}; -use axum::http::StatusCode; -use axum::response::sse::{Event, KeepAlive}; -use axum::response::{IntoResponse, Sse}; -use axum::routing::{any, get}; -use axum_extra::json; +use axum::{ + Json, Router, + extract::{ + ConnectInfo, Path, State, WebSocketUpgrade, + ws::{Message, Utf8Bytes, WebSocket}, + }, + http::StatusCode, + response::{IntoResponse, Sse, sse::Event, sse::KeepAlive}, + routing::{any, get}, +}; use dashmap::DashMap; use futures_util::{SinkExt, Stream, StreamExt, stream}; use once_cell::sync::Lazy; use orion::ws::WSMessage; -use rand::seq::IndexedRandom; -use scopeguard::defer; -use sea_orm::ActiveValue::Set; -use sea_orm::prelude::DateTimeUtc; -use sea_orm::sqlx::types::chrono; -use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter as _}; +use rand::Rng; +use sea_orm::{ + prelude::DateTimeUtc, + { + ActiveModelTrait, ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, + QueryFilter as _, + }, +}; use serde::{Deserialize, Serialize}; use std::convert::Infallible; use std::io::Write; @@ -26,15 +29,35 @@ use std::ops::ControlFlow; use std::sync::Arc; use std::time::Duration; use tokio::io::AsyncReadExt; +use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::mpsc::UnboundedSender; use utoipa::ToSchema; -use utoipa_axum::{router::OpenApiRouter, routes}; use uuid::Uuid; +// Global configuration for build log directory static BUILD_LOG_DIR: Lazy = Lazy::new(|| std::env::var("BUILD_LOG_DIR").expect("BUILD_LOG_DIR must be set")); +/// Creates a log file for a specific task ID +/// Ensures parent directories exist before creating the file +fn create_log_file(task_id: &str) -> Result { + let log_path = format!("{}/{}", *BUILD_LOG_DIR, task_id); + let path = std::path::Path::new(&log_path); + + // Ensure parent directory exists + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + // Create or open the log file in append mode + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) +} + +/// Request payload for creating a new build task #[derive(Debug, Deserialize, ToSchema)] pub struct BuildRequest { repo: String, @@ -43,24 +66,45 @@ pub struct BuildRequest { mr: Option, } +/// Information about an active build task +#[derive(Clone)] pub struct BuildInfo { - repo: String, - target: String, - args: Option>, - start_at: DateTimeUtc, - mr: Option, + pub repo: String, + pub target: String, + pub args: Option>, + pub start_at: DateTimeUtc, + pub mr: Option, + pub _worker_id: String, + pub log_file: Arc>, +} + +/// Status of a worker node +#[derive(Debug, Clone)] +pub enum WorkerStatus { + Idle, + Busy(String), // Contains task ID when busy } +/// Information about a connected worker +#[derive(Debug)] +pub struct WorkerInfo { + pub sender: UnboundedSender, + pub status: WorkerStatus, + pub last_heartbeat: DateTimeUtc, +} + +/// Enumeration of possible task statuses #[derive(Debug, Serialize, Default, ToSchema)] pub enum TaskStatusEnum { Building, - Interrupted, // exit code is None + Interrupted, // Task was interrupted, exit code is None Failed, Completed, #[default] NotFound, } +/// Response structure for task status queries #[derive(Debug, Serialize, Default, ToSchema)] pub struct TaskStatus { status: TaskStatusEnum, @@ -70,20 +114,22 @@ pub struct TaskStatus { message: Option, } +/// Shared application state containing worker connections, database, and active builds #[derive(Clone)] pub struct AppState { - pub clients: Arc>>, + pub workers: Arc>, pub conn: DatabaseConnection, - pub building: Arc>, + pub active_builds: Arc>, } -pub fn routers() -> OpenApiRouter { - OpenApiRouter::new() +/// Creates and configures all API routes +pub fn routers() -> Router { + Router::new() .route("/ws", any(ws_handler)) - .routes(routes!(task_handler)) - .routes(routes!(task_status_handler)) + .route("/task", axum::routing::post(task_handler)) + .route("/task-status/{id}", get(task_status_handler)) .route("/task-output/{id}", get(task_output_handler)) - .routes(routes!(task_query_by_mr)) + .route("/mr-task/{mr}", get(task_query_by_mr)) } #[utoipa::path( @@ -96,11 +142,14 @@ pub fn routers() -> OpenApiRouter { (status = 200, description = "Task status", body = TaskStatus) ) )] -async fn task_status_handler( +/// Retrieves the current status of a build task by its ID +/// Returns status information including exit code and current state +pub async fn task_status_handler( Path(id): Path, State(state): State, ) -> impl IntoResponse { - let (code, status) = if state.building.contains_key(&id) { + let (code, status) = if state.active_builds.contains_key(&id) { + // Task is currently active/building ( StatusCode::OK, TaskStatus { @@ -109,12 +158,16 @@ async fn task_status_handler( }, ) } else { + // Check database for completed/historical tasks match Uuid::parse_str(&id) { - Ok(id) => { - let output = builds::Model::get_by_build_id(id, state.conn).await; + Ok(id_uuid) => { + let output = builds::Model::get_by_build_id(id_uuid, &state.conn).await; match output { Some(model) => { - let status = if model.exit_code.is_none() { + // Determine task status based on database fields + let status = if model.end_at.is_none() { + TaskStatusEnum::Building + } else if model.exit_code.is_none() { TaskStatusEnum::Interrupted } else if model.exit_code.unwrap() == 0 { TaskStatusEnum::Completed @@ -153,56 +206,57 @@ async fn task_status_handler( (code, Json(status)) } -/// SSE +/// Streams build output logs in real-time using Server-Sent Events (SSE) +/// Continuously monitors the log file and streams new content as it becomes available async fn task_output_handler( State(state): State, Path(id): Path, -) -> Sse>> // impl IntoResponse -{ - let path = format!("{}/{}", *BUILD_LOG_DIR, id); - if !std::path::Path::new(&path).exists() { - // 2 return types must same, which is hard without `.boxed()` - // `Sse, ..., ...>>` != Sse> != Sse> +) -> Sse>> { + let log_path_str = format!("{}/{}", *BUILD_LOG_DIR, id); + let log_path = std::path::Path::new(&log_path_str); + + // Return error message if log file doesn't exist + if !log_path.exists() { return Sse::new( stream::once(async { Ok(Event::default().data("Task output file not found")) }).boxed(), ); } - let file = tokio::fs::File::open(&path).await.unwrap(); // read-only mode + let file = tokio::fs::File::open(log_path).await.unwrap(); let reader = tokio::io::BufReader::new(file); + // Create a stream that continuously reads from the log file let stream = stream::unfold(reader, move |mut reader| { - let id_c = id.clone(); // must, or err - let building = state.building.clone(); + let id_c = id.clone(); + let active_builds = state.active_builds.clone(); async move { let mut buf = String::new(); - let is_building = building.contains_key(&id_c); // MUST check before reading + let is_building = active_builds.contains_key(&id_c); let size = reader.read_to_string(&mut buf).await.unwrap(); + if size == 0 { if is_building { - // wait for new content + // Wait for more content if build is still active tokio::time::sleep(Duration::from_secs(1)).await; let size = reader.read_to_string(&mut buf).await.unwrap(); if size > 0 { - tracing::debug!("Read: {}", buf); // little duplicate code, but more efficient Some((Ok(Event::default().data(buf)), reader)) } else { - tracing::debug!("Not Modified, waiting..."); - // return control to `axum`, or it can't auto-detect client disconnect & close + // Send keep-alive comment Some((Ok(Event::default().comment("")), reader)) } } else { - // build end & no more content + // Build finished and no more content None } } else { - tracing::debug!("Read: {}", buf); + // Send new content Some((Ok(Event::default().data(buf)), reader)) } } }); - Sse::new(stream.boxed()).keep_alive(KeepAlive::new()) // empty comment to keep alive + Sse::new(stream.boxed()).keep_alive(KeepAlive::new()) } #[utoipa::path( @@ -210,231 +264,333 @@ async fn task_output_handler( path = "/task", request_body = BuildRequest, responses( - (status = 200, description = "Task created", body = inline(json::Value)) + (status = 200, description = "Task created", body = serde_json::Value) ) )] -async fn task_handler( +/// Creates a new build task and assigns it to an available worker +/// Returns task ID and assigned worker information upon successful creation +pub async fn task_handler( State(state): State, Json(req): Json, -) -> impl IntoResponse { - let id = Uuid::now_v7().to_string(); - state.building.insert( - id.clone(), - BuildInfo { - repo: req.repo.clone(), - target: req.target.clone(), - args: req.args.clone(), - start_at: chrono::Utc::now(), - mr: req.mr.clone(), - }, - ); - - let client_ids: Vec<_> = state - .clients +) -> (StatusCode, Json) { + // Find all idle workers + let idle_workers: Vec = state + .workers .iter() + .filter(|entry| matches!(entry.value().status, WorkerStatus::Idle)) .map(|entry| entry.key().clone()) .collect(); - if client_ids.is_empty() { - return json!({"message": "No clients connected"}); + // Return error if no workers are available + if idle_workers.is_empty() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({"message": "No available workers at the moment"})), + ); } - let mut rng = rand::rng(); - let chosen_id = client_ids.choose(&mut rng).unwrap(); + // Randomly select an idle worker + let chosen_index = { + let mut rng = rand::rng(); + rng.random_range(0..idle_workers.len()) + }; + let chosen_id = idle_workers[chosen_index].clone(); + let task_id = Uuid::now_v7(); + + // Create log file for the task + let log_file = match create_log_file(&task_id.to_string()) { + Ok(file) => Arc::new(Mutex::new(file)), + Err(e) => { + tracing::error!("Failed to create log file for task {}: {}", task_id, e); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"message": "Failed to create log file"})), + ); + } + }; + + // Create build information structure + let build_info = BuildInfo { + repo: req.repo.clone(), + target: req.target.clone(), + args: req.args.clone(), + start_at: chrono::Utc::now(), + mr: req.mr.clone(), + _worker_id: chosen_id.clone(), + log_file, + }; + + // Save task to database + let model = builds::ActiveModel { + build_id: Set(task_id), + output_file: Set(format!("{}/{}", *BUILD_LOG_DIR, task_id)), + exit_code: Set(None), + start_at: Set(build_info.start_at), + end_at: Set(None), + repo_name: Set(build_info.repo.clone()), + target: Set(build_info.target.clone()), + arguments: Set(build_info.args.clone().unwrap_or_default().join(" ")), + mr: Set(build_info.mr.clone().unwrap_or_default()), + }; + if let Err(e) = model.insert(&state.conn).await { + tracing::error!("Failed to insert new build task into DB: {}", e); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"message": "Failed to create task in database"})), + ); + } + // Create WebSocket message for the worker let msg = WSMessage::Task { - id: id.clone(), + id: task_id.to_string(), repo: req.repo, target: req.target, args: req.args, - mr: req.mr.clone().unwrap_or_default(), + mr: req.mr.unwrap_or_default(), }; - state.clients.get(chosen_id).unwrap().send(msg).unwrap(); // TODO client maybe disconnected - - json!({"task_id": id, "client_id": chosen_id}) + // Send task to the selected worker + if let Some(mut worker) = state.workers.get_mut(&chosen_id) { + if worker.sender.send(msg).is_ok() { + worker.status = WorkerStatus::Busy(task_id.to_string()); + state.active_builds.insert(task_id.to_string(), build_info); + tracing::info!("Task {} dispatched to worker {}", task_id, chosen_id); + ( + StatusCode::OK, + Json(serde_json::json!({"task_id": task_id.to_string(), "client_id": chosen_id})), + ) + } else { + tracing::error!( + "Failed to send task to supposedly idle worker {}. It might have just disconnected.", + chosen_id + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json( + serde_json::json!({"message": "Failed to dispatch task to worker. Please try again."}), + ), + ) + } + } else { + tracing::error!( + "Chosen idle worker {} not found in map. This should not happen.", + chosen_id + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"message": "Internal scheduler error."})), + ) + } } + +/// Handles WebSocket upgrade requests from workers +/// Establishes bidirectional communication channel with worker nodes async fn ws_handler( ws: WebSocketUpgrade, ConnectInfo(addr): ConnectInfo, State(state): State, ) -> impl IntoResponse { - println!("{addr} connected."); + tracing::info!("{addr} connected. Waiting for registration..."); ws.on_upgrade(move |socket| handle_socket(socket, addr, state)) } -/// Actual websocket statemachine (one will be spawned per connection) -async fn handle_socket(mut socket: WebSocket, who: SocketAddr, state: AppState) { - let client_id = Uuid::now_v7().to_string(); +/// Manages WebSocket connection lifecycle for worker communication +/// Handles message sending/receiving and connection cleanup +async fn handle_socket(socket: WebSocket, who: SocketAddr, state: AppState) { let (tx, mut rx) = mpsc::unbounded_channel::(); - let clients = state.clients.clone(); - clients.insert(client_id.clone(), tx); - defer! { - println!("clean Client {who}."); - clients.remove(&client_id); - } - - // send a ping (unsupported by some browsers) just to kick things off and get a response - if socket - .send(Message::Ping(Bytes::from_static(b"Server hello"))) - .await - .is_ok() - { - println!("Pinged {who}..."); - } else { - println!("Could not send ping {who}!"); - return; // exit - } + let mut worker_id: Option = None; - // By splitting socket we can send and receive at the same time. In this example we will send - // unsolicited messages to client based on some sort of server's internal event (i.e .timer). let (mut sender, mut receiver) = socket.split(); - // Spawn a task that will push several messages to the client (does not matter what client does) - let mut send_task = tokio::spawn(async move { + // Task for sending messages to the worker + let send_task = tokio::spawn(async move { while let Some(msg) = rx.recv().await { - let msg = serde_json::to_string(&msg).unwrap(); + let msg_str = serde_json::to_string(&msg).unwrap(); if sender - .send(Message::Text(Utf8Bytes::from(msg))) + .send(Message::Text(Utf8Bytes::from(msg_str))) .await .is_err() { - println!("Error sending message to {who}"); + tracing::warn!("Failed to send message to {who}, client disconnected."); break; } } }); - // This second task will receive messages from client and print them on server console - let mut recv_task = tokio::spawn(async move { + let state_clone = state.clone(); + let tx_clone = tx.clone(); + + // Task for receiving messages from the worker + let recv_task = tokio::spawn(async move { + let mut worker_id_inner: Option = None; while let Some(Ok(msg)) = receiver.next().await { - if process_message(msg, who, state.clone()).await.is_break() { + if process_message(msg, who, &state_clone, &mut worker_id_inner, &tx_clone) + .await + .is_break() + { break; } } + worker_id_inner }); - // If any one of the tasks exit, abort the other. tokio::select! { - rv_a = (&mut send_task) => { - match rv_a { - Ok(_) => println!("send_task to {who} over"), - Err(a) => println!("Error sending messages {a:?}") + _ = send_task => { }, + result = recv_task => { + if let Ok(final_worker_id) = result { + worker_id = final_worker_id; } - recv_task.abort(); }, - rv_b = (&mut recv_task) => { - match rv_b { - Ok(_) => println!("recv_task from {who} over"), - Err(b) => println!("Error receiving messages {b:?}") - } - send_task.abort(); - } } - // returning from the handler closes the websocket connection - println!("Websocket context {who} destroyed"); + // Cleanup worker connection when socket closes + if let Some(id) = &worker_id { + tracing::info!("Cleaning up for worker: {id} from {who}."); + state.workers.remove(id); + } else { + tracing::info!("Cleaning up unregistered connection from {who}."); + } + + tracing::info!("Websocket context {who} destroyed"); } -async fn process_message(msg: Message, who: SocketAddr, state: AppState) -> ControlFlow<(), ()> { +/// Processes individual WebSocket messages from workers +/// Handles registration, heartbeats, build output, and completion messages +async fn process_message( + msg: Message, + who: SocketAddr, + state: &AppState, + worker_id: &mut Option, + tx: &UnboundedSender, +) -> ControlFlow<(), ()> { match msg { - Message::Text(t) => match serde_json::from_str::(t.as_str()) { - Ok(msg) => match msg { - // todo useless ? - WSMessage::TaskAck { - id, - success, - message, - } => { - println!(">>> task ack: id:{id}, success:{success}, msg:{message}"); + Message::Text(t) => { + let ws_msg: Result = serde_json::from_str(&t); + if let Err(e) = ws_msg { + tracing::warn!("Failed to parse message from {who}: {e}"); + return ControlFlow::Continue(()); + } + let ws_msg = ws_msg.unwrap(); + + // Handle worker registration (must be first message) + if worker_id.is_none() { + if let WSMessage::Register { id } = ws_msg { + tracing::info!("Worker from {who} registered as: {id}"); + state.workers.insert( + id.clone(), + WorkerInfo { + sender: tx.clone(), + status: WorkerStatus::Idle, + last_heartbeat: chrono::Utc::now(), + }, + ); + *worker_id = Some(id); + } else { + tracing::error!( + "First message from {who} was not Register. Closing connection." + ); + return ControlFlow::Break(()); + } + return ControlFlow::Continue(()); + } + + // Process messages from registered workers + let current_worker_id = worker_id.as_ref().unwrap(); + match ws_msg { + WSMessage::Register { .. } => { + tracing::warn!( + "Worker {current_worker_id} sent Register message again. Ignoring." + ); + } + WSMessage::Heartbeat => { + if let Some(mut worker) = state.workers.get_mut(current_worker_id) { + worker.last_heartbeat = chrono::Utc::now(); + tracing::debug!("Received heartbeat from {current_worker_id}"); + } } WSMessage::BuildOutput { id, output } => { - println!(">>> build output: id:{id}, output:{output}"); - let mut file = std::fs::OpenOptions::new() // TODO optimize: open & close too many times - .append(true) - .create(true) - .open(format!("{}/{}", *BUILD_LOG_DIR, id)) - .unwrap(); - file.write_all(format!("{output}\n").as_bytes()).unwrap(); + // Write build output to the associated log file + if let Some(build_info) = state.active_builds.get(&id) { + let log_file = build_info.log_file.clone(); + tokio::spawn(async move { + let mut file = log_file.lock().await; + if let Err(e) = writeln!(file, "{output}") { + tracing::error!( + "Failed to write to log file for task {}: {}", + id, + e + ); + } else if let Err(e) = file.flush() { + tracing::error!("Failed to flush log file for task {}: {}", id, e); + } + }); + } else { + tracing::warn!("Received output for unknown task: {}", id); + } } WSMessage::BuildComplete { id, - success, + success: _, exit_code, - message, + message: _, } => { - println!( - ">>> got build complete: id:{id}, success:{success}, exit_code:{exit_code:?}, msg:{message}" + // Handle build completion + tracing::info!( + "Build {id} completed by worker {current_worker_id} with exit code: {exit_code:?}" ); - let info = state.building.get(&id).expect("Build info not found"); - let model = builds::ActiveModel { - build_id: Set(id.parse().unwrap()), - output_file: Set(format!("{}/{}", *BUILD_LOG_DIR, id)), - exit_code: Set(exit_code), - start_at: Set(info.start_at), - end_at: Set(chrono::Utc::now()), - repo_name: Set(info.repo.clone()), - target: Set(info.target.clone()), - arguments: Set(info.args.clone().unwrap_or_default().join(" ")), - mr: Set(info.mr.clone().unwrap_or_default()), - }; - drop(info); // !!release ref or deadlock when insert - model.insert(&state.conn).await.unwrap(); - state.building.remove(&id); // task over + + // Remove from active builds and update database + state.active_builds.remove(&id); + let _ = builds::Entity::update_many() + .set(builds::ActiveModel { + exit_code: Set(exit_code), + end_at: Set(Some(chrono::Utc::now())), + ..Default::default() + }) + .filter(builds::Column::BuildId.eq(id.parse::().unwrap())) + .exec(&state.conn) + .await; + + // Mark worker as idle again + if let Some(mut worker) = state.workers.get_mut(current_worker_id) { + worker.status = WorkerStatus::Idle; + } } - _ => unreachable!(), - }, - Err(e) => { - println!("Error parsing message: {e}"); + _ => {} } - }, - Message::Binary(d) => { - println!(">>> {} sent {} bytes: {:?}", who, d.len(), d); } - Message::Close(c) => { - if let Some(cf) = c { - println!( - ">>> {} sent close with code {} and reason `{}`", - who, cf.code, cf.reason - ); - } else { - println!(">>> {who} somehow sent close message without CloseFrame"); - } + Message::Close(_) => { + tracing::info!("Client {who} sent close message."); return ControlFlow::Break(()); } - Message::Pong(v) => { - println!(">>> {who} sent pong with {v:?}"); - } - // You should never need to manually handle Message::Ping, as axum's websocket library - // will do so for you automagically by replying with Pong and copying the v according to - // spec. But if you need the contents of the pings you can see them here. - Message::Ping(v) => { - println!(">>> {who} sent ping with {v:?}"); - } + _ => {} } ControlFlow::Continue(()) } +/// Data transfer object for build information in API responses #[derive(Debug, Serialize, ToSchema)] pub struct BuildDTO { pub build_id: String, pub output_file: String, pub exit_code: Option, pub start_at: String, - pub end_at: String, + pub end_at: Option, pub repo_name: String, pub target: String, pub arguments: String, pub mr: String, } + impl BuildDTO { + /// Converts a database model to a DTO for API responses pub fn from_model(model: builds::Model) -> Self { Self { build_id: model.build_id.to_string(), output_file: model.output_file, exit_code: model.exit_code, start_at: model.start_at.to_rfc3339(), - end_at: model.end_at.to_rfc3339(), + end_at: model.end_at.map(|dt| dt.to_rfc3339()), repo_name: model.repo_name, target: model.target, arguments: model.arguments, @@ -442,11 +598,7 @@ impl BuildDTO { } } } -/// Query builds by merge request (MR) number -/// This is a new endpoint to query builds by MR number. -/// It returns a list of builds associated with the given MR number. -/// If no builds are found, it returns an empty list. -/// If an error occurs during the query, it returns an empty list with a 500 status code. + #[utoipa::path( get, path = "/mr-task/{mr}", @@ -455,11 +607,13 @@ impl BuildDTO { ), responses( (status = 200, description = "Builds for MR", body = [BuildDTO]), - (status = 404, description = "No builds found for the given MR", body = inline(json::Value)), - (status = 500, description = "Internal server error", body = inline(json::Value)) + (status = 404, description = "No builds found for the given MR", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) ) )] -async fn task_query_by_mr( +/// Retrieves all build tasks associated with a specific merge request +/// Returns a list of builds filtered by MR number +pub async fn task_query_by_mr( State(state): State, Path(mr): Path, ) -> Result>, (StatusCode, Json)> { @@ -486,8 +640,10 @@ async fn task_query_by_mr( } } } + #[cfg(test)] mod tests { + /// Test random number generation for worker selection #[test] fn test_rng() { use rand::seq::IndexedRandom; diff --git a/orion-server/src/main.rs b/orion-server/src/main.rs index 8ce031a8a..b2e429fa3 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -2,16 +2,23 @@ mod api; mod model; mod server; +/// Orion Build Server +/// A distributed build system that manages build tasks and worker nodes #[tokio::main] async fn main() { - tracing_subscriber::fmt() // default is INFO + // Initialize logging with DEBUG level + tracing_subscriber::fmt() .with_max_level(tracing::Level::DEBUG) .init(); - dotenvy::dotenv().ok(); // .env file is optional + // Load environment variables from .env file (optional) + dotenvy::dotenv().ok(); + + // Get server port from environment or use default let port: u16 = std::env::var("PORT") .unwrap_or_else(|_| "8004".to_string()) .parse() .expect("PORT must be a number"); + server::start_server(port).await; } diff --git a/orion-server/src/model/builds.rs b/orion-server/src/model/builds.rs index 9afe99b2f..ce109f835 100644 --- a/orion-server/src/model/builds.rs +++ b/orion-server/src/model/builds.rs @@ -1,18 +1,23 @@ use sea_orm::entity::prelude::*; use serde::Serialize; -#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize)] +/// Database model for build tasks +/// Stores information about build jobs including their status, timing, and metadata +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Default)] #[sea_orm(table_name = "builds")] pub struct Model { #[sea_orm(primary_key, auto_increment = false)] pub build_id: Uuid, pub output_file: String, - pub exit_code: Option, // On Unix, return `None` if the process was terminated by a signal. + /// Exit code of the build process. None if process was terminated by signal on Unix + pub exit_code: Option, pub start_at: DateTimeUtc, - pub end_at: DateTimeUtc, + pub end_at: Option, pub repo_name: String, - pub target: String, // build target, e.g. "//:main" + /// Build target specification (e.g., "//:main") + pub target: String, pub arguments: String, + /// Merge request identifier pub mr: String, } @@ -22,10 +27,11 @@ pub enum Relation {} impl ActiveModelBehavior for ActiveModel {} impl Model { - pub async fn get_by_build_id(build_id: Uuid, conn: DatabaseConnection) -> Option { + /// Retrieves a build record by its UUID from the database + pub async fn get_by_build_id(build_id: Uuid, conn: &DatabaseConnection) -> Option { Entity::find() .filter(Column::BuildId.eq(build_id)) - .one(&conn) + .one(conn) .await .expect("Failed to get by `build_id`") } diff --git a/orion-server/src/server.rs b/orion-server/src/server.rs index 951d3d449..a227c09f4 100644 --- a/orion-server/src/server.rs +++ b/orion-server/src/server.rs @@ -1,74 +1,149 @@ -use crate::api; -use crate::api::AppState; +use crate::api::{self, AppState}; use crate::model::builds; -use api::{BuildDTO, BuildRequest, TaskStatus, TaskStatusEnum}; use axum::Router; use axum::routing::get; use dashmap::DashMap; -use sea_orm::{ConnectionTrait, Database, DatabaseConnection, DbErr, Schema, TransactionTrait}; +use sea_orm::{ + ActiveValue::Set, ColumnTrait, ConnectionTrait, Database, DatabaseConnection, DbErr, + EntityTrait, QueryFilter, Schema, TransactionTrait, +}; use std::net::SocketAddr; use std::sync::Arc; +use std::time::Duration; use tower_http::trace::TraceLayer; use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; + +/// OpenAPI documentation configuration #[derive(OpenApi)] #[openapi( paths( - crate::api::task_handler, - crate::api::task_status_handler, - crate::api::task_query_by_mr, + api::task_handler, + api::task_status_handler, + api::task_query_by_mr, ), components( - schemas(BuildRequest, TaskStatus, TaskStatusEnum, BuildDTO) + schemas(api::BuildRequest, api::TaskStatus, api::TaskStatusEnum, api::BuildDTO) ), tags( (name = "Build", description = "Build related endpoints") ) )] pub struct ApiDoc; + +/// Starts the Orion server with the specified port +/// Initializes database connection, sets up routes, and starts health check tasks pub async fn start_server(port: u16) { let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL is not set in .env file"); - - let conn = Database::connect(db_url) // TODO pool + let conn = Database::connect(db_url) .await .expect("Database connection failed"); setup_tables(&conn).await.expect("Failed to setup tables"); + let state = AppState { + workers: Arc::new(DashMap::new()), + conn, + active_builds: Arc::new(DashMap::new()), + }; + + // Start background health check task + tokio::spawn(start_health_check_task(state.clone())); + let app = Router::new() .route("/", get(|| async { "Hello, World!" })) .merge(api::routers()) .merge(SwaggerUi::new("/swagger-ui").url("/api-doc/openapi.json", ApiDoc::openapi())) - .with_state(AppState { - clients: Arc::new(DashMap::new()), - conn, - building: Arc::new(DashMap::new()), - }) - // logging so we can see what's going on + .with_state(state) .layer(TraceLayer::new_for_http()); tracing::info!("Listening on port {}", port); - let addr = tokio::net::TcpListener::bind(&format!("0.0.0.0:{port}")) .await .unwrap(); axum::serve( addr, - app.into_make_service_with_connect_info::(), // or `ConnectInfo` fail + app.into_make_service_with_connect_info::(), ) .await .unwrap(); } -/// create if not exists +/// Sets up database tables if they don't exist async fn setup_tables(conn: &DatabaseConnection) -> Result<(), DbErr> { let trans = conn.begin().await?; - let builder = conn.get_database_backend(); let schema = Schema::new(builder); - let mut table_statement = schema.create_table_from_entity(builds::Entity); - table_statement.if_not_exists(); - let statement = builder.build(&table_statement); + let statement = builder.build( + schema + .create_table_from_entity(builds::Entity) + .if_not_exists(), + ); trans.execute(statement).await?; - trans.commit().await } + +/// Background task that monitors worker health and handles timeouts +/// Removes dead workers and marks their tasks as interrupted +async fn start_health_check_task(state: AppState) { + let health_check_interval = Duration::from_secs(30); + let worker_timeout = Duration::from_secs(90); + + tracing::info!( + "Health check task started. Interval: {:?}, Worker timeout: {:?}", + health_check_interval, + worker_timeout + ); + + loop { + tokio::time::sleep(health_check_interval).await; + tracing::debug!("Running health check..."); + + let mut dead_workers = Vec::new(); + let now = chrono::Utc::now(); + + // Find workers that haven't sent heartbeat within timeout period + for entry in state.workers.iter() { + if now.signed_duration_since(entry.value().last_heartbeat) + > chrono::Duration::from_std(worker_timeout).unwrap() + { + dead_workers.push(entry.key().clone()); + } + } + + if dead_workers.is_empty() { + continue; + } + + tracing::warn!("Found dead workers: {:?}", dead_workers); + + // Remove dead workers and handle their tasks + for worker_id in dead_workers { + if let Some((_, worker_info)) = state.workers.remove(&worker_id) { + tracing::info!("Removed dead worker: {}", worker_id); + + // If worker was busy, mark task as interrupted + if let api::WorkerStatus::Busy(task_id) = worker_info.status { + tracing::warn!( + "Worker {} was busy with task {}. Marking task as Interrupted.", + worker_id, + task_id + ); + state.active_builds.remove(&task_id); + + let update_res = builds::Entity::update_many() + .set(builds::ActiveModel { + end_at: Set(Some(chrono::Utc::now())), + ..Default::default() + }) + .filter(builds::Column::BuildId.eq(task_id.parse::().unwrap())) + .exec(&state.conn) + .await; + + if let Err(e) = update_res { + tracing::error!("Failed to update orphaned task {} in DB: {}", task_id, e); + } + } + } + } + } +} diff --git a/orion/.env b/orion/.env index ad2700e68..b253a92bd 100644 --- a/orion/.env +++ b/orion/.env @@ -1,4 +1,4 @@ BUCK_PROJECT_ROOT="/tmp/megadir/mount" -SERVER_WS="ws://orion.gitmega.com/ws" +SERVER_WS="ws://127.0.0.1/ws" SELECT_TASK_COUNT="30" -INITIAL_POLL_INTERVAL_SECS="2" \ No newline at end of file +INITIAL_POLL_INTERVAL_SECS="2" diff --git a/orion/src/api.rs b/orion/src/api.rs index 7ec4683aa..e9c9d4a6e 100644 --- a/orion/src/api.rs +++ b/orion/src/api.rs @@ -1,48 +1,65 @@ use crate::buck_controller; use crate::ws::WSMessage; -use axum::Json; -use dashmap::DashSet; -use once_cell::sync::Lazy; -use serde::{Deserialize, Serialize}; +use serde::Serialize; use tokio::sync::mpsc::UnboundedSender; use uuid::Uuid; -#[derive(Debug, Deserialize)] + +/// Parameters required to execute a buck build operation. +#[derive(Debug)] pub struct BuildRequest { + /// Repository path or identifier pub repo: String, + /// Buck build target (e.g., "//path/to:target") pub target: String, + /// Additional command-line arguments for the build pub args: Option>, - pub mr: String, // merge request id - // pub webhook: Option, // post + /// Merge request identifier for context + pub mr: String, } +/// Result of a build operation containing status and metadata. #[derive(Debug, Serialize)] pub struct BuildResult { + /// Whether the build operation was successful pub success: bool, + /// Unique identifier for the build task pub id: String, + /// Process exit code (None if not yet completed) #[serde(skip_serializing_if = "Option::is_none")] pub exit_code: Option, + /// Human-readable status or error message pub message: String, } -static BUILDING: Lazy> = Lazy::new(DashSet::new); -// TODO avoid multi-task in one repo? -// #[debug_handler] // better error msg -// `Json` must be last arg, because it consumes the request body +/// Initiates an asynchronous buck build process. +/// +/// The build executes in a background task, allowing this function to return immediately +/// with an acknowledgment. Build progress and completion are communicated via WebSocket. +/// +/// # Arguments +/// * `id` - Unique identifier for tracking the build task +/// * `req` - Build parameters including repository, target, and arguments +/// * `sender` - Channel for sending WebSocket messages during build execution +/// +/// # Returns +/// Immediate acknowledgment that the build task has been queued and started pub async fn buck_build( id: Uuid, req: BuildRequest, sender: UnboundedSender, -) -> Json { - let id_c = id; - BUILDING.insert(id.to_string()); - tracing::info!("Start build task: {}", id); +) -> BuildResult { + let id_str = id.to_string(); + tracing::info!("[Task {}] Received build request.", id_str); + + // Spawn background task to handle the actual build process tokio::spawn(async move { - let build_resp = match buck_controller::build( - id_c.to_string(), - req.repo.clone(), - req.target.clone(), + // Execute the build operation via buck_controller + let build_result = match buck_controller::build( + id_str.clone(), + req.repo, + req.target, req.args.unwrap_or_default(), - req.mr.clone(), + req.mr, sender.clone(), ) .await @@ -51,45 +68,57 @@ pub async fn buck_build( let message = format!( "Build {}", if status.success() { - "success" + "succeeded" } else { "failed" } ); - tracing::info!("{}; Exit code: {:?}", message, status.code()); + tracing::info!( + "[Task {}] {}; Exit code: {:?}", + id_str, + message, + status.code() + ); BuildResult { success: status.success(), - id: id_c.to_string(), + id: id_str.clone(), exit_code: status.code(), message, } } Err(e) => { - tracing::error!("Run buck2 failed: {}", e); + let error_msg = format!("Build execution failed: {e}"); + tracing::error!("[Task {}] {}", id_str, error_msg); BuildResult { success: false, - id: id_c.to_string(), + id: id_str.clone(), exit_code: None, - message: e.to_string(), + message: error_msg, } } }; - BUILDING.remove(&id_c.to_string()); // MUST after database insert to ensure data accessible - sender - .send(WSMessage::BuildComplete { - id: id_c.to_string(), - success: build_resp.success, - exit_code: build_resp.exit_code, - message: build_resp.message.clone(), - }) - .unwrap(); + // Send build completion notification via WebSocket + let complete_msg = WSMessage::BuildComplete { + id: build_result.id, + success: build_result.success, + exit_code: build_result.exit_code, + message: build_result.message, + }; + + if sender.send(complete_msg).is_err() { + tracing::error!( + "[Task {}] Failed to send BuildComplete message. Connection likely lost.", + id_str + ); + } }); - Json(BuildResult { - success: true, // TODO + // Return immediate acknowledgment of task acceptance + BuildResult { + success: true, id: id.to_string(), exit_code: None, - message: "Build started".to_string(), - }) + message: "Build task has been accepted and started.".to_string(), + } } diff --git a/orion/src/buck_controller.rs b/orion/src/buck_controller.rs index e6bf003e9..16888541f 100644 --- a/orion/src/buck_controller.rs +++ b/orion/src/buck_controller.rs @@ -1,7 +1,8 @@ use crate::ws::WSMessage; use once_cell::sync::Lazy; use serde_json::{json, Value}; -use std::io; +// Import complete Error trait for better error handling +use std::error::Error; use std::process::{ExitStatus, Stdio}; use tokio::io::AsyncBufReadExt; use tokio::process::Command; @@ -11,20 +12,21 @@ use tokio::time::{sleep, Duration}; static PROJECT_ROOT: Lazy = Lazy::new(|| std::env::var("BUCK_PROJECT_ROOT").expect("BUCK_PROJECT_ROOT must be set")); -/// Sends a filesystem mount request to the specified API endpoint and waits for completion -/// Parameters: -/// - repo: Repository path to mount -/// - mr: Merge request number -/// Returns: Result containing success status on completion, or an error on failure -pub async fn mount_fs(repo: &str, mr: &str) -> Result> { - // Create HTTP client +/// Mounts filesystem via remote API for repository access. +/// +/// Initiates mount request and polls for completion with exponential backoff. +/// Required for accessing repository files during build process. +/// +/// # Arguments +/// * `repo` - Repository path to mount +/// * `mr` - Merge request identifier +/// +/// # Returns +/// * `Ok(true)` - Mount operation completed successfully +/// * `Err(_)` - Mount request failed or timed out +pub async fn mount_fs(repo: &str, mr: &str) -> Result> { let client = reqwest::Client::new(); - - // Step 1: Send mount request to get request_id - let mount_payload = json!({ - "path": repo, - "mr": mr, - }); + let mount_payload = json!({ "path": repo, "mr": mr }); let mount_res = client .post("http://localhost:2725/api/fs/mount") @@ -33,53 +35,67 @@ pub async fn mount_fs(repo: &str, mr: &str) -> Result().unwrap_or(10); - for _attempt in 1..=max_attempts { - // Wait before checking status + + let mut poll_interval = initial_poll_interval_secs; + + for attempt in 1..=max_attempts { sleep(Duration::from_secs(poll_interval)).await; - // Exponential backoff: double interval, up to max_poll_interval_secs poll_interval = std::cmp::min(poll_interval * 2, max_poll_interval_secs); let select_url = format!("http://localhost:2725/api/fs/select/{request_id}"); let select_res = client.get(&select_url).send().await?; - let select_body: Value = select_res.json().await?; - println!("Select response: {select_body}"); - // Check overall status + tracing::debug!( + "Polling mount status (attempt {}/{}): {:?}", + attempt, + max_attempts, + select_body + ); + if select_body.get("status").and_then(|v| v.as_str()) != Some("Success") { - return Err("Select request failed".into()); + let err_msg = select_body + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("Select request failed"); + return Err(err_msg.into()); } - // Check task status match select_body.get("task_status").and_then(|v| v.as_str()) { Some("finished") => { - println!("Mount task completed successfully"); + tracing::info!( + "Mount task completed successfully for request_id: {}", + request_id + ); return Ok(true); } Some("error") => { @@ -89,22 +105,31 @@ pub async fn mount_fs(repo: &str, mr: &str) -> Result { - println!("Mount task still in progress (fetching)..."); - continue; - } - Some(other_status) => { - println!("Mount task status: {other_status}"); - continue; - } - None => { - return Err("Missing task_status in select response".into()); - } + _ => continue, } } + Err("Mount operation timed out".into()) } +/// Executes buck build with filesystem mounting and output streaming. +/// +/// Process flow: +/// 1. Mount repository filesystem via remote API +/// 2. Execute buck build command with specified target and arguments +/// 3. Stream build output in real-time via WebSocket +/// 4. Return final build status +/// +/// # Arguments +/// * `id` - Build task identifier for logging and tracking +/// * `repo` - Repository path for filesystem mounting +/// * `target` - Buck build target specification +/// * `args` - Additional command-line arguments for buck +/// * `mr` - Merge request context identifier +/// * `sender` - WebSocket channel for streaming build output +/// +/// # Returns +/// Process exit status indicating build success or failure pub async fn build( id: String, repo: String, @@ -112,71 +137,70 @@ pub async fn build( args: Vec, mr: String, sender: UnboundedSender, -) -> io::Result { - tracing::info!("Building {} in repo {} with target {}", id, repo, target); - // Prepare the command to run - // Note: `args` is a list of additional arguments to pass to the `buck - - // Mount filesystem before building - match mount_fs(&repo, &mr).await { - Ok(true) => { - tracing::info!("Filesystem mounted successfully for repo: {}", repo); - } - Ok(false) => { - tracing::error!("Filesystem mount failed for repo: {}", repo); - return Err(io::Error::other("Filesystem mount failed")); - } - Err(e) => { - tracing::error!("Error mounting filesystem for repo {}: {}", repo, e); - return Err(io::Error::other(format!("Filesystem mount error: {e}"))); - } - } +) -> Result> { + tracing::info!( + "[Task {}] Building target '{}' in repo '{}'", + id, + target, + repo + ); + + mount_fs(&repo, &mr).await?; + tracing::info!("[Task {}] Filesystem mounted successfully.", id); let mut cmd = Command::new("buck2"); let cmd = cmd .arg("build") .args(args) - .arg(target) + .arg(&target) .current_dir(format!("{}/{}", *PROJECT_ROOT, repo)) .stdout(Stdio::piped()) .stderr(Stdio::piped()); - // actually, some info (like: "BUILD SUCCESSFUL") is printed to stderr - tracing::debug!("cmd:{:?}", cmd); + tracing::debug!("[Task {}] Executing command: {:?}", id, cmd); let mut child = cmd.spawn()?; let stdout = child.stdout.take().unwrap(); let stderr = child.stderr.take().unwrap(); let mut stdout_reader = tokio::io::BufReader::new(stdout).lines(); let mut stderr_reader = tokio::io::BufReader::new(stderr).lines(); + loop { tokio::select! { result = stdout_reader.next_line() => { match result { Ok(Some(line)) => { - sender.send(WSMessage::BuildOutput { - id: id.clone(), - output: line.clone(), - }).unwrap(); + if sender.send(WSMessage::BuildOutput { id: id.clone(), output: line }).is_err() { + child.kill().await?; + return Err("WebSocket connection lost during build.".into()); + } }, - Err(_) => break, - _ => (), + Ok(None) => break, + Err(e) => { + tracing::error!("[Task {}] Error reading stdout: {}", id, e); + break; + } } - } + }, result = stderr_reader.next_line() => { match result { Ok(Some(line)) => { - sender.send(WSMessage::BuildOutput { - id: id.clone(), - output: line.clone(), - }).unwrap(); + if sender.send(WSMessage::BuildOutput { id: id.clone(), output: line }).is_err() { + child.kill().await?; + return Err("WebSocket connection lost during build.".into()); + } + }, + Ok(None) => break, + Err(e) => { + tracing::error!("[Task {}] Error reading stderr: {}", id, e); + break; }, - Err(_) => break, - _ => (), } - } - result = child.wait() => { - return result; + }, + status = child.wait() => { + let exit_status = status?; + tracing::info!("[Task {}] Buck2 process finished with status: {}", id, exit_status); + return Ok(exit_status); } } } diff --git a/orion/src/main.rs b/orion/src/main.rs index e8b7f5323..9e71d5e66 100644 --- a/orion/src/main.rs +++ b/orion/src/main.rs @@ -1,18 +1,37 @@ -use crate::ws::spawn_client; - +// Orion worker client modules mod api; mod buck_controller; mod util; mod ws; +use uuid::Uuid; + #[tokio::main] async fn main() { - tracing_subscriber::fmt() // default is INFO - .with_max_level(tracing::Level::DEBUG) + // Initialize structured logging + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .with_target(true) .init(); - tracing::info!("current dir: {:?}", std::env::current_dir().unwrap()); - dotenvy::dotenv().ok(); // .env file is optional - let server_ws = std::env::var("SERVER_WS").expect("SERVER_WS not set"); - spawn_client(&server_ws).await; + // Load environment variables from .env file + dotenvy::dotenv().ok(); + + // Configure WebSocket server address + let server_addr = + std::env::var("SERVER_WS").unwrap_or_else(|_| "ws://127.0.0.1:8004/ws".to_string()); + + // Configure worker identification + let worker_id = std::env::var("ORION_WORKER_ID").unwrap_or_else(|_| { + tracing::warn!("ORION_WORKER_ID not set, generating a random worker ID for this session."); + // Generate time-ordered UUID for better traceability + Uuid::now_v7().to_string() + }); + + tracing::info!("Starting orion worker..."); + tracing::info!(" Worker ID: {}", worker_id); + tracing::info!(" Connecting to server at: {}", server_addr); + + // Start WebSocket client with persistent connection + ws::run_client(server_addr, worker_id).await; } diff --git a/orion/src/ws.rs b/orion/src/ws.rs index 966affcb3..990409807 100644 --- a/orion/src/ws.rs +++ b/orion/src/ws.rs @@ -1,25 +1,25 @@ -use axum::Json; +use crate::api::{buck_build, BuildRequest}; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; use std::ops::ControlFlow; -use tokio::sync::mpsc::UnboundedSender; -use tokio::sync::{mpsc, OnceCell}; -// we will use tungstenite for websocket client impl (same library as what axum is using) -use crate::api::{buck_build, BuildRequest}; -use tokio_tungstenite::{connect_async, tungstenite::protocol::Message}; -use tungstenite::Utf8Bytes; - -static SENDER: OnceCell> = OnceCell::const_new(); +use std::time::Duration; +use tokio::net::TcpStream; +use tokio::sync::mpsc; +use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; +use tokio_tungstenite::{ + connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream, +}; +use uuid::Uuid; -#[derive(Debug, Serialize, Deserialize)] +/// Message protocol for WebSocket communication between worker and server. +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(tag = "type")] pub enum WSMessage { - Task { + // Worker -> Server messages + Register { id: String, - repo: String, - target: String, - args: Option>, - mr: String, }, + Heartbeat, TaskAck { id: String, success: bool, @@ -35,149 +35,213 @@ pub enum WSMessage { exit_code: Option, message: String, }, + // Server -> Worker messages + Task { + id: String, + repo: String, + target: String, + args: Option>, + mr: String, + }, } -const MAX_RETRIES: u32 = 3; - -/// send message and retry -async fn send_with_retry( - sender: &mut (impl SinkExt + Unpin), - msg: &WSMessage, -) -> Result<(), String> { - let msg_str = serde_json::to_string(msg).map_err(|e| format!("seriaz fail: {e}"))?; - let ws_msg = Message::Text(Utf8Bytes::from(msg_str)); +/// Manages persistent WebSocket connection with automatic reconnection. +/// +/// Handles connection establishment, registration, heartbeat, and task processing. +/// Implements exponential backoff for reconnection attempts. +/// +/// # Arguments +/// * `server_addr` - WebSocket server endpoint URL +/// * `worker_id` - Unique identifier for this worker instance +pub async fn run_client(server_addr: String, worker_id: String) { + let mut reconnect_delay = Duration::from_secs(1); + const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(60); - for attempt in 0..=MAX_RETRIES { - match sender.send(ws_msg.clone()).await { - Ok(_) => return Ok(()), - Err(_) => { - if attempt == MAX_RETRIES { - return Err(format!("max retry times:({MAX_RETRIES}), ")); - } - println!("send fail (retry {}/{})", attempt + 1, MAX_RETRIES); - const RETRY_DELAY_MS: u64 = 200; - tokio::time::sleep(std::time::Duration::from_millis(RETRY_DELAY_MS)).await; + loop { + tracing::info!("Attempting to connect to server: {}", server_addr); + match connect_async(&server_addr).await { + Ok((ws_stream, response)) => { + tracing::info!( + "WebSocket handshake successful. Server response: {:?}", + response.status() + ); + // Reset reconnect delay after successful connection + reconnect_delay = Duration::from_secs(1); + // Handle the active connection + handle_connection(ws_stream, worker_id.clone()).await; + tracing::warn!("Disconnected from server."); + } + Err(e) => { + tracing::error!( + "WebSocket handshake failed: {}. Retrying in {:?}...", + e, + reconnect_delay + ); } } + // Wait before attempting to reconnect + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(MAX_RECONNECT_DELAY); } - Err("max retry error".into()) } -pub async fn spawn_client(server: &str) { - let ws_stream = match connect_async(server).await { - Ok((stream, response)) => { - println!("Server response was {response:?}"); - stream - } - Err(e) => { - println!("WebSocket handshake failed with {e}!"); +/// Processes an established WebSocket connection. +/// +/// Coordinates three concurrent tasks: +/// - Heartbeat transmission to maintain connection +/// - Message sending from internal channels +/// - Message receiving and processing from server +/// +/// # Arguments +/// * `ws_stream` - Established WebSocket connection +/// * `worker_id` - Worker identifier for registration +async fn handle_connection( + ws_stream: WebSocketStream>, + worker_id: String, +) { + let (ws_sender, mut ws_receiver) = ws_stream.split(); + let (internal_tx, mut internal_rx): (UnboundedSender, UnboundedReceiver) = + mpsc::unbounded_channel(); + + let worker_id_clone = worker_id.clone(); + let internal_tx_clone = internal_tx.clone(); + let heartbeat_task = tokio::spawn(async move { + tracing::info!("Registering with worker ID: {}", worker_id_clone); + if internal_tx_clone + .send(WSMessage::Register { + id: worker_id_clone, + }) + .is_err() + { + tracing::error!("Failed to queue register message. Internal channel closed."); return; } - }; - - let (mut sender, mut receiver) = ws_stream.split(); - - let (tx, mut rx) = mpsc::unbounded_channel::(); - SENDER.set(tx).unwrap(); - - let mut send_task = tokio::spawn(async move { - while let Some(msg) = rx.recv().await { - //let msg = serde_json::to_string(&msg).unwrap(); - if let Err(e) = send_with_retry(&mut sender, &msg).await { - println!("Error sending message: {e}"); + let heartbeat_interval = Duration::from_secs(30); + loop { + tokio::time::sleep(heartbeat_interval).await; + tracing::debug!("Sending heartbeat..."); + if internal_tx_clone.send(WSMessage::Heartbeat).is_err() { + tracing::warn!("Failed to queue heartbeat message. Internal channel closed."); break; } } }); - let mut recv_task = tokio::spawn(async move { - while let Some(Ok(msg)) = receiver.next().await { - if process_message(msg).await.is_break() { + let mut ws_sender = ws_sender; + let send_task = tokio::spawn(async move { + while let Some(msg) = internal_rx.recv().await { + match serde_json::to_string(&msg) { + Ok(msg_str) => { + if let Err(e) = ws_sender.send(Message::Text(msg_str.into())).await { + tracing::error!( + "Failed to send message to server: {}. Terminating send task.", + e + ); + break; + } + } + Err(e) => { + tracing::error!("Failed to serialize WSMessage: {}", e); + } + } + } + }); + + let internal_tx_clone = internal_tx.clone(); + let recv_task = tokio::spawn(async move { + while let Some(Ok(msg)) = ws_receiver.next().await { + if process_server_message(msg, internal_tx_clone.clone()) + .await + .is_break() + { break; } } }); - //wait for either task to finish and kill the other task + // Wait for any task to complete tokio::select! { - _ = (&mut send_task) => { - recv_task.abort(); - }, - _ = (&mut recv_task) => { - send_task.abort(); - } + _ = heartbeat_task => tracing::info!("Heartbeat task finished."), + _ = send_task => tracing::info!("Send task finished."), + _ = recv_task => tracing::info!("Receive task finished."), } } -async fn process_message(msg: Message) -> ControlFlow<(), ()> { +/// Processes incoming server messages and handles task execution. +/// +/// Handles different message types including Task assignments and connection management. +/// For Task messages, spawns build processes and sends acknowledgments. +/// +/// # Arguments +/// * `msg` - WebSocket message received from server +/// * `tx` - Channel for sending response messages +/// +/// # Returns +/// * `ControlFlow::Continue(())` - Continue message processing +/// * `ControlFlow::Break(())` - Terminate connection +async fn process_server_message( + msg: Message, + sender: UnboundedSender, +) -> ControlFlow<(), ()> { match msg { - Message::Text(t) => match serde_json::from_str::(&t) { - Ok(msg) => match msg { - WSMessage::Task { - id, - repo, - target, - args, - mr, - } => { - println!(">>> got task: id:{id}, repo:{repo}, target:{target}, args:{args:?}, mr:{mr}"); - let Json(res) = buck_build( - id.parse().unwrap(), - BuildRequest { + Message::Text(t) => { + match serde_json::from_str::(&t) { + Ok(ws_msg) => { + match ws_msg { + WSMessage::Task { + id, repo, target, args, mr, - }, - SENDER.get().unwrap().clone(), - ) - .await; - SENDER - .get() - .unwrap() - .send(WSMessage::TaskAck { - id: id.clone(), - success: res.success, - message: res.message, - }) - .unwrap(); + } => { + tracing::info!("Received task: id={}", id); + tokio::spawn(async move { + let task_id_uuid = match Uuid::parse_str(&id) { + Ok(uuid) => uuid, + Err(e) => { + tracing::error!("Failed to parse task id '{}' as Uuid: {}. Aborting task.", id, e); + return; + } + }; + + let build_result = buck_build( + task_id_uuid, + BuildRequest { + repo, + target, + args, + mr, + }, + sender.clone(), + ) + .await; + + if let Err(e) = sender.send(WSMessage::TaskAck { + id, + success: build_result.success, + message: build_result.message.clone(), + }) { + tracing::error!("Failed to send TaskAck: {}", e); + } + }); + } + // Log unexpected message types + _ => { + tracing::warn!("Received unexpected message from server: {:?}", ws_msg); + } + } } - _ => { - unreachable!("Impossible msg to client {msg:?}"); + Err(e) => { + tracing::error!("Error deserializing message from server: {}", e); } - }, - Err(e) => { - println!("Error parsing message: {e}"); } - }, - Message::Binary(d) => { - println!(">>> got {} bytes: {:?}", d.len(), d); } Message::Close(c) => { - if let Some(cf) = c { - println!( - ">>> got close with code {} and reason `{}`", - cf.code, cf.reason - ); - } else { - println!(">>> somehow got close message without CloseFrame"); - } + tracing::warn!("Server sent close frame: {:?}", c); return ControlFlow::Break(()); } - - Message::Pong(v) => { - println!(">>> got pong with {v:?}"); - } - // Just as with axum server, the underlying tungstenite websocket library - // will handle Ping for you automagically by replying with Pong and copying the - // v according to spec. But if you need the contents of the pings you can see them here. - Message::Ping(v) => { - println!(">>> got ping with {v:?}"); - } - - Message::Frame(_) => { - unreachable!("This is never supposed to happen") - } + _ => {} // Ignore Binary, Ping, Pong and other message types } ControlFlow::Continue(()) }