From ecf10f9679263d2d6139cb5f5f0185cf33ac180a Mon Sep 17 00:00:00 2001 From: yyjeqhc <1772413353@qq.com> Date: Mon, 11 Aug 2025 15:15:19 +0800 Subject: [PATCH] fix: buck2 build use a fix target --- orion-server/src/api.rs | 106 ++++++++++++++++++---------------- orion-server/src/buck2.rs | 80 ++++++++++++------------- orion-server/src/main.rs | 5 +- orion-server/src/scheduler.rs | 89 ++++++++++++++++------------ orion-server/src/server.rs | 2 +- 5 files changed, 153 insertions(+), 129 deletions(-) diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 543fdf022..0301efc19 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,8 +1,7 @@ -use crate::buck2::download_and_get_buck2_targets; use crate::model::builds; use crate::scheduler::{ - BuildInfo, BuildRequest, TaskScheduler, WorkerInfo, WorkerStatus, - create_log_file, get_build_log_dir, TaskQueueStats + BuildInfo, BuildRequest, TaskQueueStats, TaskScheduler, WorkerInfo, WorkerStatus, + create_log_file, get_build_log_dir, }; use axum::{ Json, Router, @@ -19,10 +18,8 @@ use futures_util::{SinkExt, Stream, StreamExt, stream}; use orion::ws::WSMessage; use rand::Rng; use sea_orm::{ - { - ActiveModelTrait, ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, - QueryFilter as _, - }, + ActiveModelTrait, ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, + QueryFilter as _, }; use serde::Serialize; use std::convert::Infallible; @@ -38,7 +35,6 @@ use tokio::sync::mpsc::UnboundedSender; use utoipa::ToSchema; use uuid::Uuid; - /// Enumeration of possible task statuses #[derive(Debug, Serialize, Default, ToSchema)] pub enum TaskStatusEnum { @@ -69,20 +65,15 @@ pub struct AppState { impl AppState { /// Create new AppState instance - pub fn new(conn: DatabaseConnection, queue_config: Option) -> Self { + pub fn new( + conn: DatabaseConnection, + queue_config: Option, + ) -> Self { let workers = Arc::new(DashMap::new()); let active_builds = Arc::new(DashMap::new()); - let scheduler = TaskScheduler::new( - conn.clone(), - workers, - active_builds, - queue_config, - ); - - Self { - scheduler, - conn, - } + let scheduler = TaskScheduler::new(conn.clone(), workers, active_builds, queue_config); + + Self { scheduler, conn } } } @@ -111,9 +102,7 @@ pub async fn start_queue_manager(state: AppState) { (status = 200, description = "Queue statistics", body = TaskQueueStats) ) )] -pub async fn queue_stats_handler( - State(state): State, -) -> impl IntoResponse { +pub async fn queue_stats_handler(State(state): State) -> impl IntoResponse { let stats = state.scheduler.get_queue_stats().await; (StatusCode::OK, Json(stats)) } @@ -270,17 +259,18 @@ pub async fn task_handler( Json(req): Json, ) -> impl IntoResponse { // Download and get buck2 targets first - let target = match download_and_get_buck2_targets(&req.buck_hash, &req.buckconfig_hash).await { - Ok(target) => target, - Err(e) => { - tracing::error!("Failed to download buck2 targets: {}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ "message": format!("Failed to download buck2 targets: {}", e) })), - ).into_response(); - } - }; - + // let target = match download_and_get_buck2_targets(&req.buck_hash, &req.buckconfig_hash).await { + // Ok(target) => target, + // Err(e) => { + // tracing::error!("Failed to download buck2 targets: {}", e); + // return ( + // StatusCode::INTERNAL_SERVER_ERROR, + // Json(serde_json::json!({ "message": format!("Failed to download buck2 targets: {}", e) })), + // ).into_response(); + // } + // }; + // for now we do not extract from file, just use the fixed build target. + let target = "//...".to_string(); // Check if there are idle workers available if state.scheduler.has_idle_workers() { @@ -288,10 +278,14 @@ pub async fn task_handler( handle_immediate_task_dispatch(state, req, target).await } else { // No idle workers, add task to queue - match state.scheduler.enqueue_task(req.clone(), target.clone()).await { + match state + .scheduler + .enqueue_task(req.clone(), target.clone()) + .await + { Ok(task_id) => { tracing::info!("Task {} queued for later processing", task_id); - + // Save to database (mark as Pending status) let model = builds::ActiveModel { build_id: Set(task_id), @@ -304,7 +298,7 @@ pub async fn task_handler( arguments: Set(req.args.clone().unwrap_or_default().join(" ")), mr: Set(req.mr.clone().unwrap_or_default()), }; - + if let Err(e) = model.insert(&state.conn).await { tracing::error!("Failed to insert queued task into DB: {}", e); } @@ -312,11 +306,12 @@ pub async fn task_handler( ( StatusCode::OK, Json(serde_json::json!({ - "task_id": task_id.to_string(), + "task_id": task_id.to_string(), "status": "queued", "message": "Task queued for processing when workers become available" })), - ).into_response() + ) + .into_response() } Err(e) => { tracing::warn!("Failed to queue task: {}", e); @@ -325,7 +320,8 @@ pub async fn task_handler( Json(serde_json::json!({ "message": format!("Unable to queue task: {}", e) })), - ).into_response() + ) + .into_response() } } } @@ -345,7 +341,8 @@ async fn handle_immediate_task_dispatch( return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({"message": "No available workers at the moment"})), - ).into_response(); + ) + .into_response(); } // Randomly select an idle worker @@ -364,7 +361,8 @@ async fn handle_immediate_task_dispatch( return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Failed to create log file"})), - ).into_response(); + ) + .into_response(); } }; @@ -396,7 +394,8 @@ async fn handle_immediate_task_dispatch( return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Failed to create task in database"})), - ).into_response(); + ) + .into_response(); } // Create WebSocket message for the worker @@ -412,16 +411,24 @@ async fn handle_immediate_task_dispatch( if let Some(mut worker) = state.scheduler.workers.get_mut(&chosen_id) { if worker.sender.send(msg).is_ok() { worker.status = WorkerStatus::Busy(task_id.to_string()); - state.scheduler.active_builds.insert(task_id.to_string(), build_info); - tracing::info!("Task {} dispatched immediately to worker {}", task_id, chosen_id); + state + .scheduler + .active_builds + .insert(task_id.to_string(), build_info); + tracing::info!( + "Task {} dispatched immediately to worker {}", + task_id, + chosen_id + ); ( StatusCode::OK, Json(serde_json::json!({ - "task_id": task_id.to_string(), + "task_id": task_id.to_string(), "client_id": chosen_id, "status": "dispatched" })), - ).into_response() + ) + .into_response() } else { tracing::error!( "Failed to send task to supposedly idle worker {}. It might have just disconnected.", @@ -442,7 +449,8 @@ async fn handle_immediate_task_dispatch( ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Internal scheduler error."})), - ).into_response() + ) + .into_response() } } @@ -548,7 +556,7 @@ async fn process_message( }, ); *worker_id = Some(id); - + // After new worker registration, notify to process queued tasks state.scheduler.notify_task_available(); } else { diff --git a/orion-server/src/buck2.rs b/orion-server/src/buck2.rs index e3f3a6dd8..8f6741dc1 100644 --- a/orion-server/src/buck2.rs +++ b/orion-server/src/buck2.rs @@ -1,14 +1,15 @@ -use std::process::Command; use std::fs; use std::path::Path; -use tokio_retry::{strategy::ExponentialBackoff, Retry}; -use uuid::Uuid; +use std::process::Command; use std::time::Duration; - - +use tokio_retry::{Retry, strategy::ExponentialBackoff}; +use uuid::Uuid; /// Download files from file blob API using two hash values and save them to a new folder in tmp directory -async fn download_files_to_tmp(hash1: &str, hash2: &str) -> Result> { +async fn download_files_to_tmp( + hash1: &str, + hash2: &str, +) -> Result> { // Create tmp directory path // Generate a unique folder name using UUID let folder_name = Uuid::now_v7().to_string(); @@ -16,25 +17,25 @@ async fn download_files_to_tmp(hash1: &str, hash2: &str) -> Result Result<(), Box> { let api_endpoint = file_blob_endpoint(); let url = format!("{api_endpoint}/{hash}"); @@ -43,25 +44,26 @@ async fn download_file_with_retry( let retry_strategy = ExponentialBackoff::from_millis(100) .max_delay(Duration::from_secs(2)) .take(max_retries); - - Retry::spawn(retry_strategy, || { - download_single_file(&url, file_path) - }) - .await - .map_err(|e| format!("Failed to download {hash} after {max_retries} attempts: {e}").into()) + + Retry::spawn(retry_strategy, || download_single_file(&url, file_path)) + .await + .map_err(|e| format!("Failed to download {hash} after {max_retries} attempts: {e}").into()) } /// Download a single file from URL and save to specified path -async fn download_single_file(url: &str, file_path: &Path) -> Result<(), Box> { +async fn download_single_file( + url: &str, + file_path: &Path, +) -> Result<(), Box> { let response = reqwest::get(url).await?; - + if !response.status().is_success() { return Err(format!("HTTP error: {}", response.status()).into()); } - + let content = response.bytes().await?; fs::write(file_path, &content)?; - + Ok(()) } @@ -76,6 +78,7 @@ pub fn file_blob_endpoint() -> String { } /// Execute buck2 targets //... command in the given directory and return the last line string +#[allow(dead_code)] pub fn get_buck2_targets_last_line(directory: &str) -> Result> { let output = Command::new("buck2") .args(["targets", "//..."]) @@ -84,36 +87,33 @@ pub fn get_buck2_targets_last_line(directory: &str) -> Result Result> { // First, download the files (BUCK and .buckconfig) to tmp directory let tmp_dir_path = download_files_to_tmp(hash1, hash2).await?; - + // Then, execute buck2 targets command in the downloaded directory let last_line = get_buck2_targets_last_line(&tmp_dir_path)?; - + // Clean up the temporary directory after getting the result //TODO: Cleanup should happen in a finally block or using RAII to ensure temporary files are removed even if buck2 command fails. Consider using a guard or defer pattern. TempDirGuard , for example. if Path::new(&tmp_dir_path).exists() { fs::remove_dir_all(&tmp_dir_path)?; } - + Ok(last_line) } @@ -130,24 +130,24 @@ mod tests { if Path::new(test_dir).exists() { // Create a temporary directory for testing let tmp_test_dir = std::env::temp_dir().join("buck2_test"); - + // Remove existing tmp directory if it exists if tmp_test_dir.exists() { fs::remove_dir_all(&tmp_test_dir).expect("Failed to remove existing tmp directory"); } - + // Copy test directory contents to tmp directory copy_dir_all(test_dir, &tmp_test_dir).expect("Failed to copy test directory to tmp"); - + // Run buck2 targets command in the tmp directory match get_buck2_targets_last_line(tmp_test_dir.to_str().unwrap()) { Ok(last_line) => { println!("Last line: {last_line}"); assert!(!last_line.is_empty()); - }, + } Err(e) => println!("Warning: {e}"), } - + // Clean up: remove the tmp directory after test if tmp_test_dir.exists() { fs::remove_dir_all(&tmp_test_dir).expect("Failed to clean up tmp directory"); @@ -171,4 +171,4 @@ mod tests { } Ok(()) } -} \ No newline at end of file +} diff --git a/orion-server/src/main.rs b/orion-server/src/main.rs index ad1c37485..7ef9bdbec 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -1,9 +1,8 @@ mod api; -mod model; -mod server; mod buck2; +mod model; mod scheduler; - +mod server; /// Orion Build Server /// A distributed build system that manages build tasks and worker nodes diff --git a/orion-server/src/scheduler.rs b/orion-server/src/scheduler.rs index 948fa9c9a..864c2b94c 100644 --- a/orion-server/src/scheduler.rs +++ b/orion-server/src/scheduler.rs @@ -2,20 +2,18 @@ use crate::model::builds; use dashmap::DashMap; use orion::ws::WSMessage; use rand::Rng; -use sea_orm::{ - prelude::DateTimeUtc, - ActiveModelTrait, ActiveValue::Set, DatabaseConnection, -}; +use sea_orm::{ActiveModelTrait, ActiveValue::Set, DatabaseConnection, prelude::DateTimeUtc}; use serde::{Deserialize, Serialize}; use std::collections::VecDeque; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::sync::{mpsc::UnboundedSender, Mutex, Notify}; +use tokio::sync::{Mutex, Notify, mpsc::UnboundedSender}; use utoipa::ToSchema; use uuid::Uuid; /// Request payload for creating a new build task #[derive(Debug, Clone, Deserialize, ToSchema)] +#[allow(dead_code)] pub struct BuildRequest { pub repo: String, pub buck_hash: String, @@ -91,7 +89,7 @@ impl TaskQueue { pub fn cleanup_expired(&mut self) -> Vec { let now = Instant::now(); let mut expired_tasks = Vec::new(); - + self.queue.retain(|task| { if now.duration_since(task.created_at) > self.config.max_wait_time { expired_tasks.push(task.clone()); @@ -100,7 +98,7 @@ impl TaskQueue { true } }); - + expired_tasks } @@ -108,7 +106,9 @@ impl TaskQueue { pub fn get_stats(&self) -> TaskQueueStats { TaskQueueStats { total_queued: self.queue.len(), - oldest_task_age_seconds: self.queue.front() + oldest_task_age_seconds: self + .queue + .front() .map(|task| Instant::now().duration_since(task.created_at).as_secs()), } } @@ -183,9 +183,13 @@ impl TaskScheduler { } /// Add task to queue - pub async fn enqueue_task(&self, request: BuildRequest, target: String) -> Result { + pub async fn enqueue_task( + &self, + request: BuildRequest, + target: String, + ) -> Result { let task_id = Uuid::now_v7(); - + let pending_task = PendingTask { task_id, request, @@ -197,7 +201,7 @@ impl TaskScheduler { let mut queue = self.pending_tasks.lock().await; queue.enqueue(pending_task)?; } - + // Notify that there's a new task to process self.task_notifier.notify_one(); Ok(task_id) @@ -242,7 +246,7 @@ impl TaskScheduler { // Process tasks in batches, up to the number of idle workers let max_tasks = idle_workers.len(); let mut tasks_to_dispatch = Vec::with_capacity(max_tasks); - + // Batch dequeue tasks { let mut queue = self.pending_tasks.lock().await; @@ -257,14 +261,17 @@ impl TaskScheduler { // Dispatch tasks concurrently if !tasks_to_dispatch.is_empty() { - let dispatch_futures: Vec<_> = tasks_to_dispatch.into_iter().map(|task| { - let scheduler = self.clone(); - tokio::spawn(async move { - if let Err(e) = scheduler.dispatch_task(task).await { - tracing::error!("Failed to dispatch queued task: {}", e); - } + let dispatch_futures: Vec<_> = tasks_to_dispatch + .into_iter() + .map(|task| { + let scheduler = self.clone(); + tokio::spawn(async move { + if let Err(e) = scheduler.dispatch_task(task).await { + tracing::error!("Failed to dispatch queued task: {}", e); + } + }) }) - }).collect(); + .collect(); // Wait for all dispatch tasks to complete for future in dispatch_futures { @@ -291,7 +298,11 @@ impl TaskScheduler { let log_file = match create_log_file(&pending_task.task_id.to_string()) { Ok(file) => Arc::new(Mutex::new(file)), Err(e) => { - tracing::error!("Failed to create log file for task {}: {}", pending_task.task_id, e); + tracing::error!( + "Failed to create log file for task {}: {}", + pending_task.task_id, + e + ); return Err(format!("Failed to create log file: {e}")); } }; @@ -338,8 +349,13 @@ impl TaskScheduler { if let Some(mut worker) = self.workers.get_mut(&chosen_id) { if worker.sender.send(msg).is_ok() { worker.status = WorkerStatus::Busy(pending_task.task_id.to_string()); - self.active_builds.insert(pending_task.task_id.to_string(), build_info); - tracing::info!("Queued task {} dispatched to worker {}", pending_task.task_id, chosen_id); + self.active_builds + .insert(pending_task.task_id.to_string(), build_info); + tracing::info!( + "Queued task {} dispatched to worker {}", + pending_task.task_id, + chosen_id + ); Ok(()) } else { Err(format!("Failed to send task to worker {chosen_id}")) @@ -383,15 +399,18 @@ impl TaskScheduler { let cleanup_scheduler = self.clone(); let cleanup_task = tokio::spawn(async move { let mut interval = tokio::time::interval(cleanup_interval); - + loop { interval.tick().await; - + // Clean up expired tasks let expired_tasks = cleanup_scheduler.cleanup_expired_tasks().await; if !expired_tasks.is_empty() { - tracing::warn!("Cleaned up {} expired tasks from queue", expired_tasks.len()); - + tracing::warn!( + "Cleaned up {} expired tasks from queue", + expired_tasks.len() + ); + // Log expired task information for task in expired_tasks { tracing::debug!("Expired task: {} ({})", task.task_id, task.request.repo); @@ -446,7 +465,7 @@ mod tests { fn test_task_queue_fifo() { let config = TaskQueueConfig::default(); let mut queue = TaskQueue::new(config); - + // Create test tasks let task1 = PendingTask { task_id: Uuid::now_v7(), @@ -460,7 +479,7 @@ mod tests { target: "target1".to_string(), created_at: Instant::now(), }; - + let task2 = PendingTask { task_id: Uuid::now_v7(), request: BuildRequest { @@ -473,20 +492,18 @@ mod tests { target: "target2".to_string(), created_at: Instant::now(), }; - + // Test FIFO behavior assert!(queue.enqueue(task1.clone()).is_ok()); assert!(queue.enqueue(task2.clone()).is_ok()); - - + let dequeued1 = queue.dequeue().unwrap(); assert_eq!(dequeued1.task_id, task1.task_id); assert_eq!(dequeued1.request.repo, "test1"); - + let dequeued2 = queue.dequeue().unwrap(); assert_eq!(dequeued2.task_id, task2.task_id); assert_eq!(dequeued2.request.repo, "test2"); - } /// Test queue capacity limit @@ -498,7 +515,7 @@ mod tests { cleanup_interval: Duration::from_secs(30), }; let mut queue = TaskQueue::new(config); - + let task = PendingTask { task_id: Uuid::now_v7(), request: BuildRequest { @@ -511,11 +528,11 @@ mod tests { target: "target".to_string(), created_at: Instant::now(), }; - + // Fill queue to capacity assert!(queue.enqueue(task.clone()).is_ok()); assert!(queue.enqueue(task.clone()).is_ok()); - + // Should fail when full assert!(queue.enqueue(task).is_err()); } diff --git a/orion-server/src/server.rs b/orion-server/src/server.rs index 144866dac..4bf1f124a 100644 --- a/orion-server/src/server.rs +++ b/orion-server/src/server.rs @@ -43,7 +43,7 @@ pub async fn start_server(port: u16) { // Start background health check task tokio::spawn(start_health_check_task(state.clone())); - + // Start queue manager tokio::spawn(api::start_queue_manager(state.clone()));