From 7e54f711f0090da2f9d8b1a239cfe91d755819e1 Mon Sep 17 00:00:00 2001 From: Xiaoyang Han Date: Wed, 6 Aug 2025 17:08:42 +0800 Subject: [PATCH 1/5] [orion-server]: auto targets generation for buck2 Signed-off-by: Xiaoyang Han --- orion-server/.env | 3 +- orion-server/Cargo.toml | 5 +- orion-server/src/api.rs | 34 +++++--- orion-server/src/buck2.rs | 173 ++++++++++++++++++++++++++++++++++++++ orion-server/src/main.rs | 1 + 5 files changed, 204 insertions(+), 12 deletions(-) create mode 100644 orion-server/src/buck2.rs diff --git a/orion-server/.env b/orion-server/.env index e2ccd1d86..0f281d582 100644 --- a/orion-server/.env +++ b/orion-server/.env @@ -1,3 +1,4 @@ BUILD_LOG_DIR="/tmp/megadir/buck2ctl" DATABASE_URL="postgres://postgres:postgres@localhost/orion" -PORT=80 \ No newline at end of file +PORT=80 +MONOBASE_URL="http://git.gitmega.com" \ No newline at end of file diff --git a/orion-server/Cargo.toml b/orion-server/Cargo.toml index e6f312a43..142495dd3 100644 --- a/orion-server/Cargo.toml +++ b/orion-server/Cargo.toml @@ -9,6 +9,7 @@ orion = { workspace = true} axum = { workspace = true, features = ["macros", "ws"] } axum-extra = { workspace = true, features = ["erased-json"] } tokio = { workspace = true, features = ["rt-multi-thread", "fs", "process"] } +tokio-retry = "0.3" tracing = { workspace = true } tracing-subscriber = { workspace = true } serde = { workspace = true, features = ["derive"] } @@ -29,4 +30,6 @@ 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 +chrono = { version = "0.4", features = ["serde"] } +reqwest.workspace = true + diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 855d18bcc..35ebaec8a 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,3 +1,4 @@ +use crate::buck2::download_and_get_buck2_targets; use crate::model::builds; use axum::{ Json, Router, @@ -61,7 +62,8 @@ fn create_log_file(task_id: &str) -> Result { #[derive(Debug, Deserialize, ToSchema)] pub struct BuildRequest { repo: String, - target: String, + buck_hash: String, + buckconfig_hash: String, args: Option>, mr: Option, } @@ -281,7 +283,19 @@ pub async fn task_output_handler( pub async fn task_handler( State(state): State, Json(req): Json, -) -> (StatusCode, 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(); + } + }; + // Find all idle workers let idle_workers: Vec = state .workers @@ -295,7 +309,7 @@ pub async fn task_handler( return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({"message": "No available workers at the moment"})), - ); + ).into_response(); } // Randomly select an idle worker @@ -314,14 +328,14 @@ pub async fn task_handler( return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Failed to create log file"})), - ); + ).into_response(); } }; // Create build information structure let build_info = BuildInfo { repo: req.repo.clone(), - target: req.target.clone(), + target: target.clone(), args: req.args.clone(), start_at: chrono::Utc::now(), mr: req.mr.clone(), @@ -346,14 +360,14 @@ pub async fn task_handler( return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Failed to create task in database"})), - ); + ).into_response(); } // Create WebSocket message for the worker let msg = WSMessage::Task { id: task_id.to_string(), repo: req.repo, - target: req.target, + target, args: req.args, mr: req.mr.unwrap_or_default(), }; @@ -367,7 +381,7 @@ pub async fn task_handler( ( StatusCode::OK, Json(serde_json::json!({"task_id": task_id.to_string(), "client_id": chosen_id})), - ) + ).into_response() } else { tracing::error!( "Failed to send task to supposedly idle worker {}. It might have just disconnected.", @@ -378,7 +392,7 @@ pub async fn task_handler( Json( serde_json::json!({"message": "Failed to dispatch task to worker. Please try again."}), ), - ) + ).into_response() } } else { tracing::error!( @@ -388,7 +402,7 @@ pub async fn task_handler( ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"message": "Internal scheduler error."})), - ) + ).into_response() } } diff --git a/orion-server/src/buck2.rs b/orion-server/src/buck2.rs new file mode 100644 index 000000000..1cb8a8275 --- /dev/null +++ b/orion-server/src/buck2.rs @@ -0,0 +1,173 @@ +use std::process::Command; +use std::fs; +use std::path::Path; +use tokio_retry::{strategy::ExponentialBackoff, Retry}; +use uuid::Uuid; +use std::time::Duration; + + + +/// 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> { + // Create tmp directory path + // Generate a unique folder name using UUID + let folder_name = Uuid::now_v7().to_string(); + println!("Generated folder name: {folder_name},buck:{hash1},.buckconfig:{hash2}"); + + // Create tmp directory path + let tmp_dir = std::env::temp_dir().join(folder_name); + + fs::create_dir_all(&tmp_dir)?; + + // Download first file as BUCK + let buck_path = tmp_dir.join("BUCK"); + download_file_with_retry(hash1, &buck_path, 3).await?; + + // Download second file as .buckconfig + let buckconfig_path = tmp_dir.join(".buckconfig"); + download_file_with_retry(hash2, &buckconfig_path, 3).await?; + + Ok(tmp_dir.to_string_lossy().to_string()) +} + +/// Download a single file with retry mechanism using tokio-retry +async fn download_file_with_retry( + hash: &str, + file_path: &Path, + max_retries: usize +) -> Result<(), Box> { + let api_endpoint = file_blob_endpoint(); + let url = format!("{api_endpoint}/{hash}"); + + // Create retry strategy: exponential backoff starting from 100ms, max 3 attempts + 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()) +} + +/// Download a single file from URL and save to specified path +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(()) +} + +/// Get the base URL for API requests +fn base_url() -> String { + std::env::var("MONOBASE_URL").unwrap_or_else(|_| "http://localhost:8000".to_string()) +} + +/// Get the file blob API endpoint +pub fn file_blob_endpoint() -> String { + format!("{}/api/v1/file/blob", base_url()) +} + +/// Execute buck2 targets //... command in the given directory and return the last line string +pub fn get_buck2_targets_last_line(directory: &str) -> Result> { + let output = Command::new("buck2") + .args(["targets", "//..."]) + .current_dir(directory) + .output()?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(format!("buck2 command failed: {stderr}" ).into()); + } + + let stdout = String::from_utf8_lossy(&output.stdout); + let last_line = stdout + .lines() + .last() + .unwrap_or("") + .to_string(); + + Ok(last_line) +} + +/// Download files and execute buck2 targets command to get the last line output +pub async fn download_and_get_buck2_targets( + hash1: &str, + hash2: &str, +) -> 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 + if Path::new(&tmp_dir_path).exists() { + fs::remove_dir_all(&tmp_dir_path)?; + } + + Ok(last_line) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::path::Path; + + #[test] + fn test_get_buck2_targets_last_line() { + // Use the test directory within the project + let test_dir = "./test"; + 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"); + } + } else { + println!("Test directory '{}' does not exist. Skipping test.", test_dir); + } + } + + /// Helper function to recursively copy directory contents + fn copy_dir_all(src: impl AsRef, dst: impl AsRef) -> std::io::Result<()> { + fs::create_dir_all(&dst)?; + for entry in fs::read_dir(src)? { + let entry = entry?; + let ty = entry.file_type()?; + if ty.is_dir() { + copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name()))?; + } else { + fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?; + } + } + Ok(()) + } +} \ No newline at end of file diff --git a/orion-server/src/main.rs b/orion-server/src/main.rs index b2e429fa3..2db33d7cb 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -1,6 +1,7 @@ mod api; mod model; mod server; +mod buck2; /// Orion Build Server /// A distributed build system that manages build tasks and worker nodes From 445a12dd0faa32d16c646e1b58162c8492904873 Mon Sep 17 00:00:00 2001 From: Han Xiaoyang Date: Wed, 6 Aug 2025 23:33:03 +0800 Subject: [PATCH 2/5] [orion-server]: fix typo Signed-off-by: Han Xiaoyang --- orion-server/src/buck2.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/orion-server/src/buck2.rs b/orion-server/src/buck2.rs index 1cb8a8275..e3f3a6dd8 100644 --- a/orion-server/src/buck2.rs +++ b/orion-server/src/buck2.rs @@ -109,6 +109,7 @@ pub async fn download_and_get_buck2_targets( 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)?; } @@ -141,10 +142,10 @@ mod tests { // 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); + println!("Last line: {last_line}"); assert!(!last_line.is_empty()); }, - Err(e) => println!("Warning: {}", e), + Err(e) => println!("Warning: {e}"), } // Clean up: remove the tmp directory after test @@ -152,7 +153,7 @@ mod tests { fs::remove_dir_all(&tmp_test_dir).expect("Failed to clean up tmp directory"); } } else { - println!("Test directory '{}' does not exist. Skipping test.", test_dir); + println!("Test directory '{test_dir}' does not exist. Skipping test."); } } From b04f05b67b149dc3fbe4bcd6dfae030b79f65bce Mon Sep 17 00:00:00 2001 From: Xiaoyang Han Date: Thu, 7 Aug 2025 16:29:05 +0800 Subject: [PATCH 3/5] temp change by copilot Signed-off-by: Xiaoyang Han --- orion-server/Cargo.toml | 1 + orion-server/SCHEDULER_README.md | 247 ++++++++ orion-server/src/api.rs | 12 +- orion-server/src/api_scheduler.rs | 296 ++++++++++ orion-server/src/integration_example.rs | 325 +++++++++++ orion-server/src/main.rs | 1 + orion-server/src/scheduler.rs | 724 ++++++++++++++++++++++++ 7 files changed, 1600 insertions(+), 6 deletions(-) create mode 100644 orion-server/SCHEDULER_README.md create mode 100644 orion-server/src/api_scheduler.rs create mode 100644 orion-server/src/integration_example.rs create mode 100644 orion-server/src/scheduler.rs diff --git a/orion-server/Cargo.toml b/orion-server/Cargo.toml index 142495dd3..cce1e4a45 100644 --- a/orion-server/Cargo.toml +++ b/orion-server/Cargo.toml @@ -32,4 +32,5 @@ utoipa.workspace = true utoipa-swagger-ui = { workspace = true, features = ["axum"] } chrono = { version = "0.4", features = ["serde"] } reqwest.workspace = true +thiserror = "2.0" diff --git a/orion-server/SCHEDULER_README.md b/orion-server/SCHEDULER_README.md new file mode 100644 index 000000000..8e7f82f68 --- /dev/null +++ b/orion-server/SCHEDULER_README.md @@ -0,0 +1,247 @@ +# Task Scheduler Refactoring + +This document describes the refactored task scheduling system that replaces the simple direct worker assignment with a sophisticated queue-based scheduler. + +## Overview + +The original `task_handler` function had several limitations: +- Direct worker assignment without queuing +- No retry mechanism for failed tasks +- No task prioritization +- No timeout handling +- Poor resource utilization when workers are busy + +The new scheduler addresses these issues with: +- **Queue-based task management** with priority support +- **Automatic timeout detection and retry logic** +- **Worker load balancing** across multiple nodes +- **Task state tracking** and monitoring +- **Graceful degradation** when resources are constrained + +## Architecture + +### Components + +1. **TaskScheduler**: Core scheduler managing task queues and worker assignment +2. **QueuedTask**: Task representation with state, priority, and retry information +3. **SchedulerHandle**: Thread-safe handle for interacting with the scheduler +4. **SchedulerConfig**: Configuration for timeouts, retries, and queue limits + +### Task States + +Tasks progress through the following states: +- `Pending`: Waiting in queue for assignment +- `Assigned`: Sent to a worker but not yet running +- `Running`: Currently executing on a worker +- `Completed`: Successfully finished +- `Failed`: Execution failed +- `Timeout`: Exceeded time limit +- `Cancelled`: Manually cancelled + +### Priority Levels + +Tasks can have different priority levels: +- `Critical`: Hotfixes, urgent production issues +- `High`: Important features, releases, core repositories +- `Normal`: Standard development tasks +- `Low`: Documentation, refactoring, non-critical work + +## Key Features + +### 1. Queue Management + +```rust +// Tasks are queued with priority ordering +let task = QueuedTask::new(build_request, target, &config, Some(TaskPriority::High)); +scheduler_handle.add_task(task).await?; +``` + +### 2. Automatic Retry Logic + +```rust +// Tasks are automatically retried on failure/timeout +if task.retry_count < task.max_retries { + task.reset_for_retry(); + // Re-queue for execution +} +``` + +### 3. Worker Load Balancing + +The scheduler distributes tasks across available workers using round-robin assignment, ensuring even workload distribution. + +### 4. Timeout Handling + +Tasks that exceed their timeout are automatically: +- Removed from workers +- Marked as timed out +- Retried if retry count allows +- Cleaned up if max retries exceeded + +### 5. Statistics and Monitoring + +```rust +let stats = scheduler_handle.get_stats().await?; +// Returns: pending_tasks, running_tasks, completed_tasks, etc. +``` + +## Usage Examples + +### Basic Task Submission + +```rust +// Create and submit a task +let config = SchedulerConfig::default(); +let task = QueuedTask::new(build_request, target, &config, Some(TaskPriority::Normal)); +let task_id = scheduler_handle.add_task(task).await?; +``` + +### Task Cancellation + +```rust +// Cancel a running or queued task +scheduler_handle.cancel_task(task_id).await?; +``` + +### Getting Task Status + +```rust +// Get detailed task information +let task = scheduler_handle.get_task(task_id).await?; +``` + +### System Monitoring + +```rust +// Monitor scheduler performance +let stats = scheduler_handle.get_stats().await?; +println!("Pending: {}, Running: {}, Workers: {}", + stats.pending_tasks, stats.running_tasks, stats.total_workers); +``` + +## Configuration + +### Scheduler Configuration + +```rust +let config = SchedulerConfig { + max_queue_length: 500, // Maximum queued tasks + default_task_timeout: Duration::from_secs(7200), // 2 hours + default_max_retries: 2, // Retry failed tasks twice + scheduler_interval: Duration::from_millis(500), // Process queue every 500ms + cleanup_interval: Duration::from_secs(300), // Clean up every 5 minutes + completed_task_retention: Duration::from_secs(7200), // Keep completed tasks for 2 hours +}; +``` + +### Priority Determination + +The system automatically determines task priority based on: +- **Merge Request labels**: `hotfix`, `urgent`, `critical` → Critical priority +- **Repository importance**: `core`, `main`, `production` → High priority +- **Build complexity**: Fewer arguments → Higher priority (faster feedback) +- **System load**: High load → Lower priority for complex tasks + +## API Endpoints + +### New Scheduler Endpoints + +- `POST /scheduler/task` - Submit task to scheduler +- `GET /scheduler/task/{id}` - Get task status and details +- `DELETE /scheduler/task/{id}` - Cancel a task +- `GET /scheduler/stats` - Get scheduler statistics + +### Enhanced Endpoints + +- `POST /v2/task` - Enhanced task submission with scheduler integration +- `GET /v2/health` - Health check including scheduler status + +## Migration Guide + +### From Original System + +1. **Replace direct worker assignment**: + ```rust + // Old: Direct assignment + let chosen_worker = select_random_worker(); + send_task_to_worker(task, chosen_worker); + + // New: Queue-based scheduling + let task = QueuedTask::new(build_request, target, &config, priority); + scheduler_handle.add_task(task).await?; + ``` + +2. **Update error handling**: + ```rust + // Handle scheduler-specific errors + match scheduler_handle.add_task(task).await { + Ok(task_id) => { /* success */ }, + Err(SchedulerError::QueueFull) => { /* queue full */ }, + Err(SchedulerError::NoWorkers) => { /* no workers */ }, + Err(e) => { /* other errors */ }, + } + ``` + +3. **Initialize scheduler at startup**: + ```rust + // Initialize scheduler before starting server + let scheduler_handle = initialize_scheduler(app_state).await?; + ``` + +## Benefits + +### Performance Improvements + +- **Better resource utilization**: Tasks wait in queue instead of being rejected +- **Load balancing**: Even distribution across workers +- **Priority handling**: Critical tasks get processed first + +### Reliability Improvements + +- **Automatic retries**: Failed tasks are retried automatically +- **Timeout detection**: Hung tasks are detected and recovered +- **Graceful degradation**: System remains functional under high load + +### Operational Improvements + +- **Monitoring**: Detailed statistics for system health +- **Task tracking**: Complete task lifecycle visibility +- **Capacity planning**: Queue metrics help with scaling decisions + +## Monitoring and Alerting + +### Key Metrics to Monitor + +- `pending_tasks`: Queue depth (alert if > threshold) +- `timeout_tasks`: Failed tasks due to timeout (investigate if increasing) +- `idle_workers`: Available capacity (scale up if consistently 0) +- `queue_full_errors`: Rejected tasks (increase queue size or workers) + +### Health Check Integration + +The scheduler is integrated into health checks: +- System is "healthy" if workers are available and queue is manageable +- System is "degraded" if queue is backing up +- System is "unhealthy" if scheduler is not operational + +## Future Enhancements + +1. **Persistent Queue**: Store queue state in database for crash recovery +2. **Worker Affinity**: Assign tasks to specific worker types +3. **Resource Constraints**: Consider CPU/memory when scheduling +4. **Task Dependencies**: Support for task execution ordering +5. **Batch Processing**: Group related tasks for efficiency +6. **Auto-scaling**: Automatically request more workers based on queue depth + +## Error Handling + +The scheduler provides comprehensive error handling: +- `QueueFull`: Queue has reached maximum capacity +- `TaskNotFound`: Requested task doesn't exist +- `TaskExists`: Attempting to add duplicate task +- `NoWorkers`: No workers available for assignment +- `WorkerNotFound`: Assigned worker is no longer available +- `InvalidStateTransition`: Illegal task state change +- `Internal`: Other internal errors + +Each error provides detailed context for debugging and user feedback. diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 35ebaec8a..2dffdce84 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -59,13 +59,13 @@ fn create_log_file(task_id: &str) -> Result { } /// Request payload for creating a new build task -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Clone, Deserialize, ToSchema)] pub struct BuildRequest { - repo: String, - buck_hash: String, - buckconfig_hash: String, - args: Option>, - mr: Option, + pub repo: String, + pub buck_hash: String, + pub buckconfig_hash: String, + pub args: Option>, + pub mr: Option, } /// Information about an active build task diff --git a/orion-server/src/api_scheduler.rs b/orion-server/src/api_scheduler.rs new file mode 100644 index 000000000..4c501f92b --- /dev/null +++ b/orion-server/src/api_scheduler.rs @@ -0,0 +1,296 @@ +use axum::{ + extract::State, + http::StatusCode, + response::IntoResponse, + Json +}; +use uuid::Uuid; +use crate::api::{BuildRequest, AppState}; +use crate::scheduler::{QueuedTask, SchedulerHandle, SchedulerConfig, TaskPriority, SchedulerError}; +use crate::buck2::download_and_get_buck2_targets; + +/// New task handler using the scheduler +#[utoipa::path( + post, + path = "/scheduler/task", + request_body = BuildRequest, + responses( + (status = 200, description = "Task queued successfully", body = serde_json::Value), + (status = 503, description = "Queue is full", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) + ) +)] +pub async fn scheduled_task_handler( + State(state): State, + scheduler_handle: State, + 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!({ + "error": "Failed to download buck2 targets", + "message": format!("{}", e) + })), + ).into_response(); + } + }; + + // Create scheduler config (this should ideally be shared/cached) + let config = SchedulerConfig::default(); + + // Determine task priority based on request parameters + let priority = determine_task_priority(&req); + + // Create a new queued task + let task = QueuedTask::new(req, target, &config, Some(priority)); + let task_id = task.task_id; + + // Add task to scheduler queue + match scheduler_handle.add_task(task).await { + Ok(queued_task_id) => { + tracing::info!("Task {} successfully queued", queued_task_id); + ( + StatusCode::OK, + Json(serde_json::json!({ + "task_id": queued_task_id.to_string(), + "status": "queued", + "message": "Task has been queued for execution" + })), + ).into_response() + } + Err(SchedulerError::QueueFull) => { + tracing::warn!("Queue is full, rejecting task {}", task_id); + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "error": "Queue is full", + "message": "The task queue is currently full. Please try again later." + })), + ).into_response() + } + Err(SchedulerError::TaskExists(existing_id)) => { + tracing::warn!("Task {} already exists", existing_id); + ( + StatusCode::CONFLICT, + Json(serde_json::json!({ + "error": "Task already exists", + "task_id": existing_id.to_string(), + "message": "A task with this ID already exists" + })), + ).into_response() + } + Err(e) => { + tracing::error!("Failed to queue task {}: {}", task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to queue task", + "message": format!("Internal error: {}", e) + })), + ).into_response() + } + } +} + +/// Get task status using scheduler +#[utoipa::path( + get, + path = "/scheduler/task/{id}", + params( + ("id" = String, Path, description = "Task ID") + ), + responses( + (status = 200, description = "Task information", body = serde_json::Value), + (status = 404, description = "Task not found", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) + ) +)] +pub async fn get_scheduled_task( + scheduler_handle: State, + axum::extract::Path(task_id): axum::extract::Path, +) -> impl IntoResponse { + let task_uuid = match Uuid::parse_str(&task_id) { + Ok(uuid) => uuid, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid task ID format", + "message": "Task ID must be a valid UUID" + })), + ).into_response(); + } + }; + + match scheduler_handle.get_task(task_uuid).await { + Ok(Some(task)) => { + ( + StatusCode::OK, + Json(serde_json::json!({ + "task_id": task.task_id.to_string(), + "state": format!("{:?}", task.state), + "priority": format!("{:?}", task.priority), + "retry_count": task.retry_count, + "max_retries": task.max_retries, + "created_at": task.created_at.elapsed().as_secs(), + "assigned_worker": task.assigned_worker, + "repo": task.build_request.repo, + "target": task.target + })), + ).into_response() + } + Ok(None) => { + ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ + "error": "Task not found", + "task_id": task_id, + "message": "No task found with the specified ID" + })), + ).into_response() + } + Err(e) => { + tracing::error!("Failed to get task {}: {}", task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to retrieve task", + "message": format!("Internal error: {}", e) + })), + ).into_response() + } + } +} + +/// Cancel a scheduled task +#[utoipa::path( + delete, + path = "/scheduler/task/{id}", + params( + ("id" = String, Path, description = "Task ID") + ), + responses( + (status = 200, description = "Task cancelled successfully", body = serde_json::Value), + (status = 404, description = "Task not found", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) + ) +)] +pub async fn cancel_scheduled_task( + scheduler_handle: State, + axum::extract::Path(task_id): axum::extract::Path, +) -> impl IntoResponse { + let task_uuid = match Uuid::parse_str(&task_id) { + Ok(uuid) => uuid, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid task ID format", + "message": "Task ID must be a valid UUID" + })), + ).into_response(); + } + }; + + match scheduler_handle.cancel_task(task_uuid).await { + Ok(()) => { + tracing::info!("Task {} cancelled successfully", task_id); + ( + StatusCode::OK, + Json(serde_json::json!({ + "task_id": task_id, + "status": "cancelled", + "message": "Task has been cancelled successfully" + })), + ).into_response() + } + Err(SchedulerError::TaskNotFound(_)) => { + ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ + "error": "Task not found", + "task_id": task_id, + "message": "No task found with the specified ID" + })), + ).into_response() + } + Err(e) => { + tracing::error!("Failed to cancel task {}: {}", task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to cancel task", + "message": format!("Internal error: {}", e) + })), + ).into_response() + } + } +} + +/// Get scheduler statistics +#[utoipa::path( + get, + path = "/scheduler/stats", + responses( + (status = 200, description = "Scheduler statistics", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) + ) +)] +pub async fn get_scheduler_stats( + scheduler_handle: State, +) -> impl IntoResponse { + match scheduler_handle.get_stats().await { + Ok(stats) => { + ( + StatusCode::OK, + Json(serde_json::json!(stats)), + ).into_response() + } + Err(e) => { + tracing::error!("Failed to get scheduler stats: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to get statistics", + "message": format!("Internal error: {}", e) + })), + ).into_response() + } + } +} + +/// Determine task priority based on build request parameters +fn determine_task_priority(req: &BuildRequest) -> TaskPriority { + // Priority logic can be customized based on requirements + if let Some(mr) = &req.mr { + if mr.starts_with("hotfix") || mr.contains("urgent") { + return TaskPriority::Critical; + } + if mr.starts_with("feature") { + return TaskPriority::High; + } + } + + // Check if this is a quick build (fewer args typically means simpler build) + if let Some(args) = &req.args { + if args.len() <= 2 { + return TaskPriority::High; + } + } + + TaskPriority::Normal +} + +/// Helper function to create scheduler routes +pub fn scheduler_routes() -> axum::Router { + axum::Router::new() + .route("/scheduler/task", axum::routing::post(scheduled_task_handler)) + .route("/scheduler/task/:id", axum::routing::get(get_scheduled_task)) + .route("/scheduler/task/:id", axum::routing::delete(cancel_scheduled_task)) + .route("/scheduler/stats", axum::routing::get(get_scheduler_stats)) +} diff --git a/orion-server/src/integration_example.rs b/orion-server/src/integration_example.rs new file mode 100644 index 000000000..d045abb29 --- /dev/null +++ b/orion-server/src/integration_example.rs @@ -0,0 +1,325 @@ +use std::sync::Arc; +use axum::{extract::State, http::StatusCode, response::IntoResponse, Json}; +use tokio::sync::Mutex; +use tracing::{info, error}; + +use crate::api::{AppState, BuildRequest}; +use crate::scheduler::{TaskScheduler, SchedulerConfig, SchedulerHandle, QueuedTask, TaskPriority}; +use crate::buck2::download_and_get_buck2_targets; + +/// Integration example showing how to set up and use the scheduler +/// This demonstrates the complete workflow from setup to task execution + +/// Global scheduler handle - in a real application, this would be managed differently +static SCHEDULER_HANDLE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); + +/// Initialize the scheduler system +pub async fn initialize_scheduler(app_state: AppState) -> Result> { + info!("Initializing task scheduler..."); + + // Create scheduler configuration with production-ready settings + let config = SchedulerConfig { + max_queue_length: 500, + default_task_timeout: std::time::Duration::from_secs(7200), // 2 hours + default_max_retries: 2, + scheduler_interval: std::time::Duration::from_millis(500), + cleanup_interval: std::time::Duration::from_secs(300), // 5 minutes + completed_task_retention: std::time::Duration::from_secs(7200), // 2 hours + }; + + // Create and start the scheduler + let scheduler = TaskScheduler::new(config, app_state); + let handle = scheduler.get_handle(); + + // Start the scheduler in the background + let scheduler_arc = Arc::new(scheduler); + let scheduler_clone = scheduler_arc.clone(); + + tokio::spawn(async move { + scheduler_clone.run().await; + }); + + // Initialize the global handle + SCHEDULER_HANDLE.set(handle.clone()) + .map_err(|_| "Failed to set global scheduler handle")?; + + info!("Task scheduler initialized successfully"); + Ok(handle) +} + +/// Get the global scheduler handle +pub fn get_scheduler_handle() -> Option<&'static SchedulerHandle> { + SCHEDULER_HANDLE.get() +} + +/// Enhanced task handler that demonstrates the complete integration +#[utoipa::path( + post, + path = "/v2/task", + request_body = BuildRequest, + responses( + (status = 200, description = "Task queued successfully", body = serde_json::Value), + (status = 503, description = "Service unavailable", body = serde_json::Value), + (status = 500, description = "Internal server error", body = serde_json::Value) + ) +)] +pub async fn enhanced_task_handler( + State(app_state): State, + Json(req): Json, +) -> impl IntoResponse { + info!("Received task request for repo: {}", req.repo); + + // Get the scheduler handle + let scheduler_handle = match get_scheduler_handle() { + Some(handle) => handle, + None => { + error!("Scheduler not initialized"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Scheduler not available", + "message": "Task scheduler is not initialized" + })), + ).into_response(); + } + }; + + // Check if we have any available workers before proceeding + match scheduler_handle.get_stats().await { + Ok(stats) => { + if stats.total_workers == 0 { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "error": "No workers available", + "message": "No worker nodes are currently connected" + })), + ).into_response(); + } + + if stats.idle_workers == 0 && stats.pending_tasks > 10 { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "error": "System busy", + "message": "All workers are busy and queue is getting full", + "stats": { + "pending_tasks": stats.pending_tasks, + "running_tasks": stats.running_tasks, + "idle_workers": stats.idle_workers, + "busy_workers": stats.busy_workers + } + })), + ).into_response(); + } + } + Err(e) => { + error!("Failed to get scheduler stats: {}", e); + } + } + + // Download and get buck2 targets + let target = match download_and_get_buck2_targets(&req.buck_hash, &req.buckconfig_hash).await { + Ok(target) => target, + Err(e) => { + error!("Failed to download buck2 targets: {}", e); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to prepare build environment", + "message": format!("Buck2 target resolution failed: {}", e) + })), + ).into_response(); + } + }; + + // Determine task priority based on various factors + let priority = determine_priority(&req, &app_state).await; + + // Create scheduler configuration (this could be cached) + let config = SchedulerConfig::default(); + + // Create queued task + let task = QueuedTask::new(req.clone(), target, &config, Some(priority)); + let task_id = task.task_id; + + // Submit task to scheduler + match scheduler_handle.add_task(task).await { + Ok(queued_task_id) => { + info!("Task {} successfully queued with priority {:?}", queued_task_id, priority); + + // Return enhanced response with useful information + ( + StatusCode::OK, + Json(serde_json::json!({ + "task_id": queued_task_id.to_string(), + "status": "queued", + "priority": format!("{:?}", priority), + "message": "Task has been queued for execution", + "endpoints": { + "status": format!("/v2/task/{}/status", queued_task_id), + "cancel": format!("/v2/task/{}/cancel", queued_task_id), + "output": format!("/task-output/{}", queued_task_id) + } + })), + ).into_response() + } + Err(crate::scheduler::SchedulerError::QueueFull) => { + error!("Queue is full, rejecting task for repo: {}", req.repo); + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "error": "Queue is full", + "message": "The task queue is currently at capacity. Please try again later.", + "retry_after": 60 + })), + ).into_response() + } + Err(crate::scheduler::SchedulerError::TaskExists(existing_id)) => { + error!("Duplicate task ID {}", existing_id); + ( + StatusCode::CONFLICT, + Json(serde_json::json!({ + "error": "Task already exists", + "task_id": existing_id.to_string(), + "message": "A task with this ID already exists in the system" + })), + ).into_response() + } + Err(e) => { + error!("Failed to queue task {}: {}", task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": "Failed to queue task", + "message": "An internal error occurred while queueing the task" + })), + ).into_response() + } + } +} + +/// Advanced priority determination logic +async fn determine_priority(req: &BuildRequest, app_state: &AppState) -> TaskPriority { + // Check MR-based priority + if let Some(mr) = &req.mr { + let mr_lower = mr.to_lowercase(); + + // Critical priority for hotfixes and urgent changes + if mr_lower.contains("hotfix") || mr_lower.contains("urgent") || mr_lower.contains("critical") { + return TaskPriority::Critical; + } + + // High priority for releases and important features + if mr_lower.contains("release") || mr_lower.contains("feature") { + return TaskPriority::High; + } + + // Low priority for documentation and minor changes + if mr_lower.contains("docs") || mr_lower.contains("refactor") || mr_lower.contains("minor") { + return TaskPriority::Low; + } + } + + // Check repository importance (you could maintain a config for this) + if req.repo.contains("core") || req.repo.contains("main") || req.repo.contains("production") { + return TaskPriority::High; + } + + // Check build complexity - simpler builds get higher priority for faster feedback + if let Some(args) = &req.args { + if args.len() <= 2 { + return TaskPriority::High; + } + if args.len() > 5 { + return TaskPriority::Low; + } + } + + // Check current system load + if let Some(handle) = get_scheduler_handle() { + if let Ok(stats) = handle.get_stats().await { + // If system is under load, prioritize smaller/faster tasks + if stats.pending_tasks > 20 && stats.idle_workers < 2 { + return TaskPriority::Low; + } + } + } + + TaskPriority::Normal +} + +/// Health check endpoint that includes scheduler status +#[utoipa::path( + get, + path = "/v2/health", + responses( + (status = 200, description = "System health status", body = serde_json::Value), + (status = 503, description = "System unhealthy", body = serde_json::Value) + ) +)] +pub async fn health_check_with_scheduler() -> impl IntoResponse { + let mut health_status = serde_json::json!({ + "status": "healthy", + "timestamp": chrono::Utc::now().to_rfc3339(), + "version": env!("CARGO_PKG_VERSION") + }); + + // Add scheduler health information + if let Some(handle) = get_scheduler_handle() { + match handle.get_stats().await { + Ok(stats) => { + health_status["scheduler"] = serde_json::json!({ + "status": "operational", + "stats": stats + }); + + // Determine if system is healthy based on scheduler stats + let is_healthy = stats.total_workers > 0 && + stats.pending_tasks < stats.total_workers * 10; + + if !is_healthy { + health_status["status"] = serde_json::Value::String("degraded".to_string()); + } + } + Err(e) => { + error!("Failed to get scheduler stats for health check: {}", e); + health_status["scheduler"] = serde_json::json!({ + "status": "error", + "error": format!("{}", e) + }); + health_status["status"] = serde_json::Value::String("unhealthy".to_string()); + } + } + } else { + health_status["scheduler"] = serde_json::json!({ + "status": "not_initialized" + }); + health_status["status"] = serde_json::Value::String("unhealthy".to_string()); + } + + let status_code = match health_status["status"].as_str() { + Some("healthy") => StatusCode::OK, + _ => StatusCode::SERVICE_UNAVAILABLE, + }; + + (status_code, Json(health_status)).into_response() +} + +/// Example of how to integrate scheduler into your main application +pub async fn setup_application_with_scheduler(app_state: AppState) -> Result> { + // Initialize the scheduler + let _scheduler_handle = initialize_scheduler(app_state.clone()).await?; + + // Create router with enhanced endpoints + let app = axum::Router::new() + .route("/v2/task", axum::routing::post(enhanced_task_handler)) + .route("/v2/health", axum::routing::get(health_check_with_scheduler)) + // Add scheduler management routes + .merge(crate::api_scheduler::scheduler_routes()) + // Include original routes for backward compatibility + .merge(crate::api::routers()) + .with_state(app_state); + + Ok(app) +} diff --git a/orion-server/src/main.rs b/orion-server/src/main.rs index 2db33d7cb..d8d6c6384 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -2,6 +2,7 @@ mod api; mod model; mod server; mod buck2; +mod scheduler; /// 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 new file mode 100644 index 000000000..e8af443ba --- /dev/null +++ b/orion-server/src/scheduler.rs @@ -0,0 +1,724 @@ +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::{mpsc, Mutex, RwLock}; +use tokio::time; +use uuid::Uuid; +use serde::Serialize; + +use crate::api::{BuildRequest, WorkerStatus, AppState}; +use orion::ws::WSMessage; + +/// State of queued tasks +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TaskState { + /// Waiting for scheduling + Pending, + /// Assigned to a worker + Assigned(String), // worker_id + /// Currently running + Running(String), // worker_id + /// Completed successfully + Completed, + /// Failed with error + Failed, + /// Timed out + Timeout, + /// Cancelled + Cancelled, +} + +/// Task item in the queue +#[derive(Debug, Clone)] +pub struct QueuedTask { + pub task_id: Uuid, + pub build_request: BuildRequest, + pub target: String, // buck2 targets + pub state: TaskState, + pub created_at: Instant, + pub timeout_duration: Duration, + pub retry_count: u32, + pub max_retries: u32, + pub assigned_worker: Option, + pub priority: TaskPriority, +} + +/// Task priority levels +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum TaskPriority { + Low = 0, + Normal = 1, + High = 2, + Critical = 3, +} + +impl Default for TaskPriority { + fn default() -> Self { + TaskPriority::Normal + } +} + +impl QueuedTask { + /// Create a new QueuedTask from BuildRequest + pub fn new( + build_request: BuildRequest, + target: String, + config: &SchedulerConfig, + priority: Option, + ) -> Self { + Self { + task_id: Uuid::now_v7(), + build_request, + target, + state: TaskState::Pending, + created_at: Instant::now(), + timeout_duration: config.default_task_timeout, + retry_count: 0, + max_retries: config.default_max_retries, + assigned_worker: None, + priority: priority.unwrap_or_default(), + } + } + + /// Check if task has timed out + pub fn is_timed_out(&self) -> bool { + Instant::now().duration_since(self.created_at) > self.timeout_duration + } + + /// Check if task can be retried + pub fn can_retry(&self) -> bool { + self.retry_count < self.max_retries + } + + /// Reset task for retry + pub fn reset_for_retry(&mut self) { + self.retry_count += 1; + self.state = TaskState::Pending; + self.assigned_worker = None; + self.created_at = Instant::now(); + } +} + +/// Scheduler configuration +#[derive(Debug, Clone)] +pub struct SchedulerConfig { + /// Maximum queue length + pub max_queue_length: usize, + /// Default task timeout duration + pub default_task_timeout: Duration, + /// Default maximum retry count + pub default_max_retries: u32, + /// Scheduler check interval + pub scheduler_interval: Duration, + /// Task cleanup interval + pub cleanup_interval: Duration, + /// Retention time for completed tasks + pub completed_task_retention: Duration, +} + +impl Default for SchedulerConfig { + fn default() -> Self { + Self { + max_queue_length: 1000, + default_task_timeout: Duration::from_secs(3600), // 1 hour + default_max_retries: 3, + scheduler_interval: Duration::from_millis(100), + cleanup_interval: Duration::from_secs(60), + completed_task_retention: Duration::from_secs(3600), // 1 hour + } + } +} + +/// Scheduler statistics +#[derive(Debug, Clone, Serialize)] +pub struct SchedulerStats { + pub pending_tasks: usize, + pub running_tasks: usize, + pub completed_tasks: usize, + pub failed_tasks: usize, + pub timeout_tasks: usize, + pub total_workers: usize, + pub idle_workers: usize, + pub busy_workers: usize, +} + +/// Task scheduler +pub struct TaskScheduler { + /// Task queue sorted by priority + task_queue: Arc>>, + /// Registry of all task states + task_registry: Arc>>, + /// Scheduler configuration + config: SchedulerConfig, + /// Application state reference + app_state: AppState, + /// Scheduler control channel + control_tx: mpsc::UnboundedSender, + control_rx: Arc>>, +} + +/// Scheduler commands +#[derive(Debug)] +pub enum SchedulerCommand { + /// Add new task + AddTask { + task: QueuedTask, + response_tx: tokio::sync::oneshot::Sender>, + }, + /// Cancel task + CancelTask { + task_id: Uuid, + response_tx: tokio::sync::oneshot::Sender>, + }, + /// Task completion notification + TaskCompleted { + task_id: Uuid, + success: bool, + worker_id: String, + }, + /// Worker status update + WorkerStatusUpdate { + worker_id: String, + status: TaskState, + }, + /// Get statistics + GetStats { + response_tx: tokio::sync::oneshot::Sender, + }, + /// Get task by ID + GetTask { + task_id: Uuid, + response_tx: tokio::sync::oneshot::Sender>, + }, + /// Stop scheduler + Stop, +} + +/// Scheduler error types +#[derive(Debug, thiserror::Error)] +pub enum SchedulerError { + #[error("Queue is full")] + QueueFull, + #[error("Task not found: {0}")] + TaskNotFound(Uuid), + #[error("Task already exists: {0}")] + TaskExists(Uuid), + #[error("No available workers")] + NoWorkers, + #[error("Worker not found: {0}")] + WorkerNotFound(String), + #[error("Invalid task state transition from {from:?} to {to:?}")] + InvalidStateTransition { from: TaskState, to: TaskState }, + #[error("Internal error: {0}")] + Internal(String), +} + +impl TaskScheduler { + /// Create new task scheduler + pub fn new(config: SchedulerConfig, app_state: AppState) -> Self { + let (control_tx, control_rx) = mpsc::unbounded_channel(); + + Self { + task_queue: Arc::new(Mutex::new(VecDeque::new())), + task_registry: Arc::new(RwLock::new(HashMap::new())), + config, + app_state, + control_tx, + control_rx: Arc::new(Mutex::new(control_rx)), + } + } + + /// Get scheduler control handle + pub fn get_handle(&self) -> SchedulerHandle { + SchedulerHandle { + control_tx: self.control_tx.clone(), + } + } + + /// Start scheduler + pub async fn run(&self) { + let scheduler_task = self.run_scheduler_loop(); + let cleanup_task = self.run_cleanup_loop(); + + tokio::select! { + _ = scheduler_task => { + tracing::info!("Scheduler loop ended"); + } + _ = cleanup_task => { + tracing::info!("Cleanup loop ended"); + } + } + } + + /// Main scheduler loop + async fn run_scheduler_loop(&self) { + let mut interval = time::interval(self.config.scheduler_interval); + let mut control_rx = self.control_rx.lock().await; + + loop { + tokio::select! { + _ = interval.tick() => { + if let Err(e) = self.process_queue().await { + tracing::error!("Error processing queue: {}", e); + } + + if let Err(e) = self.check_timeouts().await { + tracing::error!("Error checking timeouts: {}", e); + } + } + + command = control_rx.recv() => { + match command { + Some(cmd) => { + if let SchedulerCommand::Stop = cmd { + tracing::info!("Scheduler stopping"); + break; + } + self.handle_command(cmd).await; + } + None => { + tracing::info!("Control channel closed, stopping scheduler"); + break; + } + } + } + } + } + } + + /// Process tasks in the queue + async fn process_queue(&self) -> Result<(), SchedulerError> { + let mut queue = self.task_queue.lock().await; + + // Get available workers + let idle_workers: Vec = self.app_state + .workers + .iter() + .filter(|entry| matches!(entry.value().status, WorkerStatus::Idle)) + .map(|entry| entry.key().clone()) + .collect(); + + if idle_workers.is_empty() { + return Ok(()); + } + + // Try to assign tasks to available workers + let mut assigned_count = 0; + let mut worker_iter = idle_workers.iter().cycle(); + + while assigned_count < idle_workers.len() && !queue.is_empty() { + if let Some(mut task) = queue.pop_front() { + if matches!(task.state, TaskState::Pending) { + if let Some(worker_id) = worker_iter.next() { + // Assign task to worker + match self.assign_task_to_worker(&mut task, worker_id.clone()).await { + Ok(()) => { + // Update task state + let mut registry = self.task_registry.write().await; + registry.insert(task.task_id, task); + assigned_count += 1; + } + Err(e) => { + tracing::error!("Failed to assign task {} to worker {}: {}", + task.task_id, worker_id, e); + // Put task back in queue + queue.push_front(task); + break; + } + } + } + } else { + // Task state is not Pending, put back in queue + queue.push_back(task); + } + } + } + + Ok(()) + } + + /// Assign task to specified worker + async fn assign_task_to_worker(&self, task: &mut QueuedTask, worker_id: String) -> Result<(), SchedulerError> { + // Create WebSocket message + let msg = WSMessage::Task { + id: task.task_id.to_string(), + repo: task.build_request.repo.clone(), + target: task.target.clone(), + args: task.build_request.args.clone(), + mr: task.build_request.mr.clone().unwrap_or_default(), + }; + + // Send task to worker + if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { + if worker.sender.send(msg).is_ok() { + // Update worker status + worker.status = WorkerStatus::Busy(task.task_id.to_string()); + + // Update task state + task.state = TaskState::Assigned(worker_id.clone()); + task.assigned_worker = Some(worker_id.clone()); + + tracing::info!("Task {} assigned to worker {}", task.task_id, worker_id); + Ok(()) + } else { + Err(SchedulerError::Internal(format!("Failed to send task to worker {}", worker_id))) + } + } else { + Err(SchedulerError::Internal(format!("Worker {} not found", worker_id))) + } + } + + /// Check for timed out tasks + async fn check_timeouts(&self) -> Result<(), SchedulerError> { + let now = Instant::now(); + let mut registry = self.task_registry.write().await; + let mut timed_out_tasks = Vec::new(); + + for (task_id, task) in registry.iter() { + if matches!(task.state, TaskState::Assigned(_) | TaskState::Running(_)) { + if now.duration_since(task.created_at) > task.timeout_duration { + timed_out_tasks.push(*task_id); + } + } + } + + for task_id in timed_out_tasks { + if let Some(task) = registry.get_mut(&task_id) { + tracing::warn!("Task {} timed out", task_id); + + // Release worker + if let Some(worker_id) = &task.assigned_worker { + if let Some(mut worker) = self.app_state.workers.get_mut(worker_id) { + worker.status = WorkerStatus::Idle; + } + } + + // Check if task can be retried + if task.retry_count < task.max_retries { + task.retry_count += 1; + task.state = TaskState::Pending; + task.assigned_worker = None; + task.created_at = now; // Reset creation time + + // Re-add to queue + let mut queue = self.task_queue.lock().await; + self.insert_task_by_priority(&mut queue, task.clone()); + + tracing::info!("Task {} requeued for retry ({}/{})", + task_id, task.retry_count, task.max_retries); + } else { + task.state = TaskState::Timeout; + tracing::error!("Task {} exceeded max retries and timed out", task_id); + } + } + } + + Ok(()) + } + + /// Insert task to queue by priority + fn insert_task_by_priority(&self, queue: &mut VecDeque, task: QueuedTask) { + let position = queue + .iter() + .position(|existing_task| existing_task.priority < task.priority) + .unwrap_or(queue.len()); + + queue.insert(position, task); + } + + /// Handle scheduler commands + async fn handle_command(&self, command: SchedulerCommand) { + match command { + SchedulerCommand::AddTask { task, response_tx } => { + let result = self.add_task_internal(task).await; + let _ = response_tx.send(result); + } + SchedulerCommand::CancelTask { task_id, response_tx } => { + let result = self.cancel_task_internal(task_id).await; + let _ = response_tx.send(result); + } + SchedulerCommand::TaskCompleted { task_id, success, worker_id } => { + self.handle_task_completion(task_id, success, worker_id).await; + } + SchedulerCommand::WorkerStatusUpdate { worker_id, status } => { + self.handle_worker_status_update(worker_id, status).await; + } + SchedulerCommand::GetStats { response_tx } => { + let stats = self.get_stats().await; + let _ = response_tx.send(stats); + } + SchedulerCommand::GetTask { task_id, response_tx } => { + let task = self.get_task_internal(task_id).await; + let _ = response_tx.send(task); + } + SchedulerCommand::Stop => { + // Already handled in main loop + } + } + } + + /// Internal add task logic + async fn add_task_internal(&self, task: QueuedTask) -> Result { + let mut queue = self.task_queue.lock().await; + + // Check if queue is full + if queue.len() >= self.config.max_queue_length { + return Err(SchedulerError::QueueFull); + } + + // Check if task already exists + let registry = self.task_registry.read().await; + if registry.contains_key(&task.task_id) { + return Err(SchedulerError::TaskExists(task.task_id)); + } + drop(registry); + + let task_id = task.task_id; + + // Add to registry + let mut registry = self.task_registry.write().await; + registry.insert(task_id, task.clone()); + drop(registry); + + // Insert by priority + self.insert_task_by_priority(&mut queue, task); + + tracing::info!("Task {} added to queue", task_id); + Ok(task_id) + } + + /// Internal cancel task logic + async fn cancel_task_internal(&self, task_id: Uuid) -> Result<(), SchedulerError> { + let mut registry = self.task_registry.write().await; + + if let Some(task) = registry.get_mut(&task_id) { + match &task.state { + TaskState::Pending => { + // Remove from queue + let mut queue = self.task_queue.lock().await; + queue.retain(|t| t.task_id != task_id); + task.state = TaskState::Cancelled; + } + TaskState::Assigned(worker_id) | TaskState::Running(worker_id) => { + // Notify worker to cancel task (if needed) + if let Some(mut worker) = self.app_state.workers.get_mut(worker_id) { + worker.status = WorkerStatus::Idle; + } + task.state = TaskState::Cancelled; + } + _ => { + return Err(SchedulerError::Internal("Cannot cancel completed task".to_string())); + } + } + + tracing::info!("Task {} cancelled", task_id); + Ok(()) + } else { + Err(SchedulerError::TaskNotFound(task_id)) + } + } + + /// Handle task completion + async fn handle_task_completion(&self, task_id: Uuid, success: bool, worker_id: String) { + let mut registry = self.task_registry.write().await; + + if let Some(task) = registry.get_mut(&task_id) { + task.state = if success { + TaskState::Completed + } else { + TaskState::Failed + }; + + // Release worker + if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { + worker.status = WorkerStatus::Idle; + } + + tracing::info!("Task {} completed with success: {}", task_id, success); + } + } + + /// Handle worker status update + async fn handle_worker_status_update(&self, worker_id: String, status: TaskState) { + if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { + match status { + TaskState::Running(task_id) => { + worker.status = WorkerStatus::Busy(task_id.clone()); + // Update task state in registry + let mut registry = self.task_registry.write().await; + if let Some(task) = registry.get_mut(&task_id.parse::().unwrap_or_default()) { + task.state = TaskState::Running(worker_id); + } + } + _ => { + // Other status updates can be handled here + tracing::debug!("Worker {} status updated to {:?}", worker_id, status); + } + } + } + } + + /// Get task by ID (internal method) + async fn get_task_internal(&self, task_id: Uuid) -> Option { + let registry = self.task_registry.read().await; + registry.get(&task_id).cloned() + } + + /// Get statistics + async fn get_stats(&self) -> SchedulerStats { + let registry = self.task_registry.read().await; + + let mut stats = SchedulerStats { + pending_tasks: 0, + running_tasks: 0, + completed_tasks: 0, + failed_tasks: 0, + timeout_tasks: 0, + total_workers: self.app_state.workers.len(), + idle_workers: 0, + busy_workers: 0, + }; + + // Count task states + for task in registry.values() { + match task.state { + TaskState::Pending => stats.pending_tasks += 1, + TaskState::Assigned(_) | TaskState::Running(_) => stats.running_tasks += 1, + TaskState::Completed => stats.completed_tasks += 1, + TaskState::Failed => stats.failed_tasks += 1, + TaskState::Timeout => stats.timeout_tasks += 1, + TaskState::Cancelled => {} // Not counted in statistics + } + } + + // Count worker states + for worker in self.app_state.workers.iter() { + match worker.value().status { + WorkerStatus::Idle => stats.idle_workers += 1, + WorkerStatus::Busy(_) => stats.busy_workers += 1, + } + } + + stats + } + + /// Cleanup loop + async fn run_cleanup_loop(&self) { + let mut interval = time::interval(self.config.cleanup_interval); + + loop { + interval.tick().await; + self.cleanup_completed_tasks().await; + } + } + + /// Clean up completed tasks + async fn cleanup_completed_tasks(&self) { + let now = Instant::now(); + let retention = self.config.completed_task_retention; + + let mut registry = self.task_registry.write().await; + let mut to_remove = Vec::new(); + + for (task_id, task) in registry.iter() { + if matches!(task.state, TaskState::Completed | TaskState::Failed | TaskState::Timeout | TaskState::Cancelled) { + if now.duration_since(task.created_at) > retention { + to_remove.push(*task_id); + } + } + } + + for task_id in to_remove { + registry.remove(&task_id); + tracing::debug!("Cleaned up completed task {}", task_id); + } + } +} + +/// Scheduler control handle +#[derive(Clone)] +pub struct SchedulerHandle { + control_tx: mpsc::UnboundedSender, +} + +impl SchedulerHandle { + /// Add task to queue + pub async fn add_task(&self, task: QueuedTask) -> Result { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + + self.control_tx.send(SchedulerCommand::AddTask { + task, + response_tx, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + response_rx.await + .map_err(|_| SchedulerError::Internal("Response channel closed".to_string()))? + } + + /// Cancel task + pub async fn cancel_task(&self, task_id: Uuid) -> Result<(), SchedulerError> { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + + self.control_tx.send(SchedulerCommand::CancelTask { + task_id, + response_tx, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + response_rx.await + .map_err(|_| SchedulerError::Internal("Response channel closed".to_string()))? + } + + /// Notify task completed + pub async fn notify_task_completed(&self, task_id: Uuid, success: bool, worker_id: String) -> Result<(), SchedulerError> { + self.control_tx.send(SchedulerCommand::TaskCompleted { + task_id, + success, + worker_id, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + Ok(()) + } + + /// Update worker status + pub async fn update_worker_status(&self, worker_id: String, status: TaskState) -> Result<(), SchedulerError> { + self.control_tx.send(SchedulerCommand::WorkerStatusUpdate { + worker_id, + status, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + Ok(()) + } + + /// Get task by ID + pub async fn get_task(&self, task_id: Uuid) -> Result, SchedulerError> { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + + self.control_tx.send(SchedulerCommand::GetTask { + task_id, + response_tx, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + response_rx.await + .map_err(|_| SchedulerError::Internal("Response channel closed".to_string())) + } + + /// Get statistics + pub async fn get_stats(&self) -> Result { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + + self.control_tx.send(SchedulerCommand::GetStats { + response_tx, + }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + response_rx.await + .map_err(|_| SchedulerError::Internal("Response channel closed".to_string())) + } + + /// Stop scheduler + pub async fn stop(&self) -> Result<(), SchedulerError> { + self.control_tx.send(SchedulerCommand::Stop) + .map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + + Ok(()) + } +} From 50de1473c2aec79d85c608cb2d41b5ec78c8f043 Mon Sep 17 00:00:00 2001 From: Xiaoyang Han Date: Fri, 8 Aug 2025 17:48:16 +0800 Subject: [PATCH 4/5] [orion-server]:Add a buffer queue for task dispatch Signed-off-by: Xiaoyang Han --- orion-server/src/api.rs | 222 +++-- orion-server/src/api_scheduler.rs | 296 ------- orion-server/src/integration_example.rs | 325 -------- orion-server/src/scheduler.rs | 1018 +++++++++-------------- orion-server/src/server.rs | 21 +- 5 files changed, 550 insertions(+), 1332 deletions(-) delete mode 100644 orion-server/src/api_scheduler.rs delete mode 100644 orion-server/src/integration_example.rs diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 2dffdce84..4971093e8 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,5 +1,9 @@ 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 +}; use axum::{ Json, Router, extract::{ @@ -12,17 +16,15 @@ use axum::{ }; use dashmap::DashMap; use futures_util::{SinkExt, Stream, StreamExt, stream}; -use once_cell::sync::Lazy; use orion::ws::WSMessage; use rand::Rng; use sea_orm::{ - prelude::DateTimeUtc, { ActiveModelTrait, ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter as _, }, }; -use serde::{Deserialize, Serialize}; +use serde::Serialize; use std::convert::Infallible; use std::io::Write; use std::net::SocketAddr; @@ -36,65 +38,6 @@ use tokio::sync::mpsc::UnboundedSender; use utoipa::ToSchema; 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, Clone, Deserialize, ToSchema)] -pub struct BuildRequest { - pub repo: String, - pub buck_hash: String, - pub buckconfig_hash: String, - pub args: Option>, - pub mr: Option, -} - -/// Information about an active build task -#[derive(Clone)] -pub struct BuildInfo { - 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 { @@ -119,9 +62,27 @@ pub struct TaskStatus { /// Shared application state containing worker connections, database, and active builds #[derive(Clone)] pub struct AppState { - pub workers: Arc>, + pub scheduler: TaskScheduler, pub conn: DatabaseConnection, - pub active_builds: Arc>, +} + +impl AppState { + /// Create new AppState instance + 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, + } + } } /// Creates and configures all API routes @@ -132,6 +93,28 @@ pub fn routers() -> Router { .route("/task-status/{id}", get(task_status_handler)) .route("/task-output/{id}", get(task_output_handler)) .route("/mr-task/{mr}", get(task_query_by_mr)) + .route("/queue-stats", get(queue_stats_handler)) +} + +/// Start queue management background task (event-driven + periodic cleanup) +pub async fn start_queue_manager(state: AppState) { + // Start the scheduler's queue manager + state.scheduler.start_queue_manager().await; +} + +/// API endpoint for getting queue statistics +#[utoipa::path( + get, + path = "/queue-stats", + responses( + (status = 200, description = "Queue statistics", body = TaskQueueStats) + ) +)] +pub async fn queue_stats_handler( + State(state): State, +) -> impl IntoResponse { + let stats = state.scheduler.get_queue_stats().await; + (StatusCode::OK, Json(stats)) } #[utoipa::path( @@ -150,7 +133,7 @@ pub async fn task_status_handler( Path(id): Path, State(state): State, ) -> impl IntoResponse { - let (code, status) = if state.active_builds.contains_key(&id) { + let (code, status) = if state.scheduler.active_builds.contains_key(&id) { // Task is currently active/building ( StatusCode::OK, @@ -225,7 +208,7 @@ pub async fn task_output_handler( State(state): State, Path(id): Path, ) -> Result>>, StatusCode> { - let log_path_str = format!("{}/{}", *BUILD_LOG_DIR, id); + let log_path_str = format!("{}/{}", get_build_log_dir(), id); let log_path = std::path::Path::new(&log_path_str); // Return error message if log file doesn't exist @@ -239,7 +222,7 @@ pub async fn task_output_handler( // Create a stream that continuously reads from the log file let stream = stream::unfold(reader, move |mut reader| { let id_c = id.clone(); - let active_builds = state.active_builds.clone(); + let active_builds = state.scheduler.active_builds.clone(); async move { let mut buf = String::new(); let is_building = active_builds.contains_key(&id_c); @@ -275,11 +258,12 @@ pub async fn task_output_handler( path = "/task", request_body = BuildRequest, responses( - (status = 200, description = "Task created", body = serde_json::Value) + (status = 200, description = "Task created", body = serde_json::Value), + (status = 503, description = "Queue is full", body = serde_json::Value) ) )] -/// Creates a new build task and assigns it to an available worker -/// Returns task ID and assigned worker information upon successful creation +/// Creates a new build task and either assigns it immediately or queues it for later processing +/// Returns task ID and status information upon successful creation pub async fn task_handler( State(state): State, Json(req): Json, @@ -296,15 +280,65 @@ pub async fn task_handler( } }; + // Check if there are idle workers available + if state.scheduler.has_idle_workers() { + // Have idle workers, directly dispatch task (keep original logic) + 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 { + 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), + output_file: Set(format!("{}/{}", get_build_log_dir(), task_id)), + exit_code: Set(None), + start_at: Set(chrono::Utc::now()), + end_at: Set(None), + repo_name: Set(req.repo.clone()), + target: Set(target.clone()), + 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); + } + + ( + StatusCode::OK, + Json(serde_json::json!({ + "task_id": task_id.to_string(), + "status": "queued", + "message": "Task queued for processing when workers become available" + })), + ).into_response() + } + Err(e) => { + tracing::warn!("Failed to queue task: {}", e); + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "message": format!("Unable to queue task: {}", e) + })), + ).into_response() + } + } + } +} + +/// Handle immediate task dispatch logic (original task_handler logic) +async fn handle_immediate_task_dispatch( + state: AppState, + req: BuildRequest, + target: String, +) -> axum::response::Response { // 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(); - - // Return error if no workers are available + let idle_workers = state.scheduler.get_idle_workers(); + + // Return error if no workers are available (this shouldn't happen theoretically since we already checked) if idle_workers.is_empty() { return ( StatusCode::SERVICE_UNAVAILABLE, @@ -346,7 +380,7 @@ pub async fn task_handler( // Save task to database let model = builds::ActiveModel { build_id: Set(task_id), - output_file: Set(format!("{}/{}", *BUILD_LOG_DIR, task_id)), + output_file: Set(format!("{}/{}", get_build_log_dir(), task_id)), exit_code: Set(None), start_at: Set(build_info.start_at), end_at: Set(None), @@ -373,14 +407,18 @@ pub async fn task_handler( }; // Send task to the selected worker - if let Some(mut worker) = state.workers.get_mut(&chosen_id) { + 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.active_builds.insert(task_id.to_string(), build_info); - tracing::info!("Task {} dispatched 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(), "client_id": chosen_id})), + Json(serde_json::json!({ + "task_id": task_id.to_string(), + "client_id": chosen_id, + "status": "dispatched" + })), ).into_response() } else { tracing::error!( @@ -469,7 +507,7 @@ async fn handle_socket(socket: WebSocket, who: SocketAddr, state: AppState) { // 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); + state.scheduler.workers.remove(id); } else { tracing::info!("Cleaning up unregistered connection from {who}."); } @@ -499,7 +537,7 @@ async fn process_message( if worker_id.is_none() { if let WSMessage::Register { id } = ws_msg { tracing::info!("Worker from {who} registered as: {id}"); - state.workers.insert( + state.scheduler.workers.insert( id.clone(), WorkerInfo { sender: tx.clone(), @@ -508,6 +546,9 @@ async fn process_message( }, ); *worker_id = Some(id); + + // After new worker registration, notify to process queued tasks + state.scheduler.notify_task_available(); } else { tracing::error!( "First message from {who} was not Register. Closing connection." @@ -526,14 +567,14 @@ async fn process_message( ); } WSMessage::Heartbeat => { - if let Some(mut worker) = state.workers.get_mut(current_worker_id) { + if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { worker.last_heartbeat = chrono::Utc::now(); tracing::debug!("Received heartbeat from {current_worker_id}"); } } WSMessage::BuildOutput { id, output } => { // Write build output to the associated log file - if let Some(build_info) = state.active_builds.get(&id) { + if let Some(build_info) = state.scheduler.active_builds.get(&id) { let log_file = build_info.log_file.clone(); tokio::spawn(async move { let mut file = log_file.lock().await; @@ -563,7 +604,7 @@ async fn process_message( ); // Remove from active builds and update database - state.active_builds.remove(&id); + state.scheduler.active_builds.remove(&id); let _ = builds::Entity::update_many() .set(builds::ActiveModel { exit_code: Set(exit_code), @@ -575,9 +616,12 @@ async fn process_message( .await; // Mark worker as idle again - if let Some(mut worker) = state.workers.get_mut(current_worker_id) { + if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { worker.status = WorkerStatus::Idle; } + + // After worker becomes idle, notify to process queued tasks + state.scheduler.notify_task_available(); } _ => {} } diff --git a/orion-server/src/api_scheduler.rs b/orion-server/src/api_scheduler.rs deleted file mode 100644 index 4c501f92b..000000000 --- a/orion-server/src/api_scheduler.rs +++ /dev/null @@ -1,296 +0,0 @@ -use axum::{ - extract::State, - http::StatusCode, - response::IntoResponse, - Json -}; -use uuid::Uuid; -use crate::api::{BuildRequest, AppState}; -use crate::scheduler::{QueuedTask, SchedulerHandle, SchedulerConfig, TaskPriority, SchedulerError}; -use crate::buck2::download_and_get_buck2_targets; - -/// New task handler using the scheduler -#[utoipa::path( - post, - path = "/scheduler/task", - request_body = BuildRequest, - responses( - (status = 200, description = "Task queued successfully", body = serde_json::Value), - (status = 503, description = "Queue is full", body = serde_json::Value), - (status = 500, description = "Internal server error", body = serde_json::Value) - ) -)] -pub async fn scheduled_task_handler( - State(state): State, - scheduler_handle: State, - 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!({ - "error": "Failed to download buck2 targets", - "message": format!("{}", e) - })), - ).into_response(); - } - }; - - // Create scheduler config (this should ideally be shared/cached) - let config = SchedulerConfig::default(); - - // Determine task priority based on request parameters - let priority = determine_task_priority(&req); - - // Create a new queued task - let task = QueuedTask::new(req, target, &config, Some(priority)); - let task_id = task.task_id; - - // Add task to scheduler queue - match scheduler_handle.add_task(task).await { - Ok(queued_task_id) => { - tracing::info!("Task {} successfully queued", queued_task_id); - ( - StatusCode::OK, - Json(serde_json::json!({ - "task_id": queued_task_id.to_string(), - "status": "queued", - "message": "Task has been queued for execution" - })), - ).into_response() - } - Err(SchedulerError::QueueFull) => { - tracing::warn!("Queue is full, rejecting task {}", task_id); - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "Queue is full", - "message": "The task queue is currently full. Please try again later." - })), - ).into_response() - } - Err(SchedulerError::TaskExists(existing_id)) => { - tracing::warn!("Task {} already exists", existing_id); - ( - StatusCode::CONFLICT, - Json(serde_json::json!({ - "error": "Task already exists", - "task_id": existing_id.to_string(), - "message": "A task with this ID already exists" - })), - ).into_response() - } - Err(e) => { - tracing::error!("Failed to queue task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to queue task", - "message": format!("Internal error: {}", e) - })), - ).into_response() - } - } -} - -/// Get task status using scheduler -#[utoipa::path( - get, - path = "/scheduler/task/{id}", - params( - ("id" = String, Path, description = "Task ID") - ), - responses( - (status = 200, description = "Task information", body = serde_json::Value), - (status = 404, description = "Task not found", body = serde_json::Value), - (status = 500, description = "Internal server error", body = serde_json::Value) - ) -)] -pub async fn get_scheduled_task( - scheduler_handle: State, - axum::extract::Path(task_id): axum::extract::Path, -) -> impl IntoResponse { - let task_uuid = match Uuid::parse_str(&task_id) { - Ok(uuid) => uuid, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "Invalid task ID format", - "message": "Task ID must be a valid UUID" - })), - ).into_response(); - } - }; - - match scheduler_handle.get_task(task_uuid).await { - Ok(Some(task)) => { - ( - StatusCode::OK, - Json(serde_json::json!({ - "task_id": task.task_id.to_string(), - "state": format!("{:?}", task.state), - "priority": format!("{:?}", task.priority), - "retry_count": task.retry_count, - "max_retries": task.max_retries, - "created_at": task.created_at.elapsed().as_secs(), - "assigned_worker": task.assigned_worker, - "repo": task.build_request.repo, - "target": task.target - })), - ).into_response() - } - Ok(None) => { - ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": "Task not found", - "task_id": task_id, - "message": "No task found with the specified ID" - })), - ).into_response() - } - Err(e) => { - tracing::error!("Failed to get task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to retrieve task", - "message": format!("Internal error: {}", e) - })), - ).into_response() - } - } -} - -/// Cancel a scheduled task -#[utoipa::path( - delete, - path = "/scheduler/task/{id}", - params( - ("id" = String, Path, description = "Task ID") - ), - responses( - (status = 200, description = "Task cancelled successfully", body = serde_json::Value), - (status = 404, description = "Task not found", body = serde_json::Value), - (status = 500, description = "Internal server error", body = serde_json::Value) - ) -)] -pub async fn cancel_scheduled_task( - scheduler_handle: State, - axum::extract::Path(task_id): axum::extract::Path, -) -> impl IntoResponse { - let task_uuid = match Uuid::parse_str(&task_id) { - Ok(uuid) => uuid, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "Invalid task ID format", - "message": "Task ID must be a valid UUID" - })), - ).into_response(); - } - }; - - match scheduler_handle.cancel_task(task_uuid).await { - Ok(()) => { - tracing::info!("Task {} cancelled successfully", task_id); - ( - StatusCode::OK, - Json(serde_json::json!({ - "task_id": task_id, - "status": "cancelled", - "message": "Task has been cancelled successfully" - })), - ).into_response() - } - Err(SchedulerError::TaskNotFound(_)) => { - ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": "Task not found", - "task_id": task_id, - "message": "No task found with the specified ID" - })), - ).into_response() - } - Err(e) => { - tracing::error!("Failed to cancel task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to cancel task", - "message": format!("Internal error: {}", e) - })), - ).into_response() - } - } -} - -/// Get scheduler statistics -#[utoipa::path( - get, - path = "/scheduler/stats", - responses( - (status = 200, description = "Scheduler statistics", body = serde_json::Value), - (status = 500, description = "Internal server error", body = serde_json::Value) - ) -)] -pub async fn get_scheduler_stats( - scheduler_handle: State, -) -> impl IntoResponse { - match scheduler_handle.get_stats().await { - Ok(stats) => { - ( - StatusCode::OK, - Json(serde_json::json!(stats)), - ).into_response() - } - Err(e) => { - tracing::error!("Failed to get scheduler stats: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to get statistics", - "message": format!("Internal error: {}", e) - })), - ).into_response() - } - } -} - -/// Determine task priority based on build request parameters -fn determine_task_priority(req: &BuildRequest) -> TaskPriority { - // Priority logic can be customized based on requirements - if let Some(mr) = &req.mr { - if mr.starts_with("hotfix") || mr.contains("urgent") { - return TaskPriority::Critical; - } - if mr.starts_with("feature") { - return TaskPriority::High; - } - } - - // Check if this is a quick build (fewer args typically means simpler build) - if let Some(args) = &req.args { - if args.len() <= 2 { - return TaskPriority::High; - } - } - - TaskPriority::Normal -} - -/// Helper function to create scheduler routes -pub fn scheduler_routes() -> axum::Router { - axum::Router::new() - .route("/scheduler/task", axum::routing::post(scheduled_task_handler)) - .route("/scheduler/task/:id", axum::routing::get(get_scheduled_task)) - .route("/scheduler/task/:id", axum::routing::delete(cancel_scheduled_task)) - .route("/scheduler/stats", axum::routing::get(get_scheduler_stats)) -} diff --git a/orion-server/src/integration_example.rs b/orion-server/src/integration_example.rs deleted file mode 100644 index d045abb29..000000000 --- a/orion-server/src/integration_example.rs +++ /dev/null @@ -1,325 +0,0 @@ -use std::sync::Arc; -use axum::{extract::State, http::StatusCode, response::IntoResponse, Json}; -use tokio::sync::Mutex; -use tracing::{info, error}; - -use crate::api::{AppState, BuildRequest}; -use crate::scheduler::{TaskScheduler, SchedulerConfig, SchedulerHandle, QueuedTask, TaskPriority}; -use crate::buck2::download_and_get_buck2_targets; - -/// Integration example showing how to set up and use the scheduler -/// This demonstrates the complete workflow from setup to task execution - -/// Global scheduler handle - in a real application, this would be managed differently -static SCHEDULER_HANDLE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); - -/// Initialize the scheduler system -pub async fn initialize_scheduler(app_state: AppState) -> Result> { - info!("Initializing task scheduler..."); - - // Create scheduler configuration with production-ready settings - let config = SchedulerConfig { - max_queue_length: 500, - default_task_timeout: std::time::Duration::from_secs(7200), // 2 hours - default_max_retries: 2, - scheduler_interval: std::time::Duration::from_millis(500), - cleanup_interval: std::time::Duration::from_secs(300), // 5 minutes - completed_task_retention: std::time::Duration::from_secs(7200), // 2 hours - }; - - // Create and start the scheduler - let scheduler = TaskScheduler::new(config, app_state); - let handle = scheduler.get_handle(); - - // Start the scheduler in the background - let scheduler_arc = Arc::new(scheduler); - let scheduler_clone = scheduler_arc.clone(); - - tokio::spawn(async move { - scheduler_clone.run().await; - }); - - // Initialize the global handle - SCHEDULER_HANDLE.set(handle.clone()) - .map_err(|_| "Failed to set global scheduler handle")?; - - info!("Task scheduler initialized successfully"); - Ok(handle) -} - -/// Get the global scheduler handle -pub fn get_scheduler_handle() -> Option<&'static SchedulerHandle> { - SCHEDULER_HANDLE.get() -} - -/// Enhanced task handler that demonstrates the complete integration -#[utoipa::path( - post, - path = "/v2/task", - request_body = BuildRequest, - responses( - (status = 200, description = "Task queued successfully", body = serde_json::Value), - (status = 503, description = "Service unavailable", body = serde_json::Value), - (status = 500, description = "Internal server error", body = serde_json::Value) - ) -)] -pub async fn enhanced_task_handler( - State(app_state): State, - Json(req): Json, -) -> impl IntoResponse { - info!("Received task request for repo: {}", req.repo); - - // Get the scheduler handle - let scheduler_handle = match get_scheduler_handle() { - Some(handle) => handle, - None => { - error!("Scheduler not initialized"); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Scheduler not available", - "message": "Task scheduler is not initialized" - })), - ).into_response(); - } - }; - - // Check if we have any available workers before proceeding - match scheduler_handle.get_stats().await { - Ok(stats) => { - if stats.total_workers == 0 { - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "No workers available", - "message": "No worker nodes are currently connected" - })), - ).into_response(); - } - - if stats.idle_workers == 0 && stats.pending_tasks > 10 { - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "System busy", - "message": "All workers are busy and queue is getting full", - "stats": { - "pending_tasks": stats.pending_tasks, - "running_tasks": stats.running_tasks, - "idle_workers": stats.idle_workers, - "busy_workers": stats.busy_workers - } - })), - ).into_response(); - } - } - Err(e) => { - error!("Failed to get scheduler stats: {}", e); - } - } - - // Download and get buck2 targets - let target = match download_and_get_buck2_targets(&req.buck_hash, &req.buckconfig_hash).await { - Ok(target) => target, - Err(e) => { - error!("Failed to download buck2 targets: {}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to prepare build environment", - "message": format!("Buck2 target resolution failed: {}", e) - })), - ).into_response(); - } - }; - - // Determine task priority based on various factors - let priority = determine_priority(&req, &app_state).await; - - // Create scheduler configuration (this could be cached) - let config = SchedulerConfig::default(); - - // Create queued task - let task = QueuedTask::new(req.clone(), target, &config, Some(priority)); - let task_id = task.task_id; - - // Submit task to scheduler - match scheduler_handle.add_task(task).await { - Ok(queued_task_id) => { - info!("Task {} successfully queued with priority {:?}", queued_task_id, priority); - - // Return enhanced response with useful information - ( - StatusCode::OK, - Json(serde_json::json!({ - "task_id": queued_task_id.to_string(), - "status": "queued", - "priority": format!("{:?}", priority), - "message": "Task has been queued for execution", - "endpoints": { - "status": format!("/v2/task/{}/status", queued_task_id), - "cancel": format!("/v2/task/{}/cancel", queued_task_id), - "output": format!("/task-output/{}", queued_task_id) - } - })), - ).into_response() - } - Err(crate::scheduler::SchedulerError::QueueFull) => { - error!("Queue is full, rejecting task for repo: {}", req.repo); - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "Queue is full", - "message": "The task queue is currently at capacity. Please try again later.", - "retry_after": 60 - })), - ).into_response() - } - Err(crate::scheduler::SchedulerError::TaskExists(existing_id)) => { - error!("Duplicate task ID {}", existing_id); - ( - StatusCode::CONFLICT, - Json(serde_json::json!({ - "error": "Task already exists", - "task_id": existing_id.to_string(), - "message": "A task with this ID already exists in the system" - })), - ).into_response() - } - Err(e) => { - error!("Failed to queue task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "error": "Failed to queue task", - "message": "An internal error occurred while queueing the task" - })), - ).into_response() - } - } -} - -/// Advanced priority determination logic -async fn determine_priority(req: &BuildRequest, app_state: &AppState) -> TaskPriority { - // Check MR-based priority - if let Some(mr) = &req.mr { - let mr_lower = mr.to_lowercase(); - - // Critical priority for hotfixes and urgent changes - if mr_lower.contains("hotfix") || mr_lower.contains("urgent") || mr_lower.contains("critical") { - return TaskPriority::Critical; - } - - // High priority for releases and important features - if mr_lower.contains("release") || mr_lower.contains("feature") { - return TaskPriority::High; - } - - // Low priority for documentation and minor changes - if mr_lower.contains("docs") || mr_lower.contains("refactor") || mr_lower.contains("minor") { - return TaskPriority::Low; - } - } - - // Check repository importance (you could maintain a config for this) - if req.repo.contains("core") || req.repo.contains("main") || req.repo.contains("production") { - return TaskPriority::High; - } - - // Check build complexity - simpler builds get higher priority for faster feedback - if let Some(args) = &req.args { - if args.len() <= 2 { - return TaskPriority::High; - } - if args.len() > 5 { - return TaskPriority::Low; - } - } - - // Check current system load - if let Some(handle) = get_scheduler_handle() { - if let Ok(stats) = handle.get_stats().await { - // If system is under load, prioritize smaller/faster tasks - if stats.pending_tasks > 20 && stats.idle_workers < 2 { - return TaskPriority::Low; - } - } - } - - TaskPriority::Normal -} - -/// Health check endpoint that includes scheduler status -#[utoipa::path( - get, - path = "/v2/health", - responses( - (status = 200, description = "System health status", body = serde_json::Value), - (status = 503, description = "System unhealthy", body = serde_json::Value) - ) -)] -pub async fn health_check_with_scheduler() -> impl IntoResponse { - let mut health_status = serde_json::json!({ - "status": "healthy", - "timestamp": chrono::Utc::now().to_rfc3339(), - "version": env!("CARGO_PKG_VERSION") - }); - - // Add scheduler health information - if let Some(handle) = get_scheduler_handle() { - match handle.get_stats().await { - Ok(stats) => { - health_status["scheduler"] = serde_json::json!({ - "status": "operational", - "stats": stats - }); - - // Determine if system is healthy based on scheduler stats - let is_healthy = stats.total_workers > 0 && - stats.pending_tasks < stats.total_workers * 10; - - if !is_healthy { - health_status["status"] = serde_json::Value::String("degraded".to_string()); - } - } - Err(e) => { - error!("Failed to get scheduler stats for health check: {}", e); - health_status["scheduler"] = serde_json::json!({ - "status": "error", - "error": format!("{}", e) - }); - health_status["status"] = serde_json::Value::String("unhealthy".to_string()); - } - } - } else { - health_status["scheduler"] = serde_json::json!({ - "status": "not_initialized" - }); - health_status["status"] = serde_json::Value::String("unhealthy".to_string()); - } - - let status_code = match health_status["status"].as_str() { - Some("healthy") => StatusCode::OK, - _ => StatusCode::SERVICE_UNAVAILABLE, - }; - - (status_code, Json(health_status)).into_response() -} - -/// Example of how to integrate scheduler into your main application -pub async fn setup_application_with_scheduler(app_state: AppState) -> Result> { - // Initialize the scheduler - let _scheduler_handle = initialize_scheduler(app_state.clone()).await?; - - // Create router with enhanced endpoints - let app = axum::Router::new() - .route("/v2/task", axum::routing::post(enhanced_task_handler)) - .route("/v2/health", axum::routing::get(health_check_with_scheduler)) - // Add scheduler management routes - .merge(crate::api_scheduler::scheduler_routes()) - // Include original routes for backward compatibility - .merge(crate::api::routers()) - .with_state(app_state); - - Ok(app) -} diff --git a/orion-server/src/scheduler.rs b/orion-server/src/scheduler.rs index e8af443ba..948fa9c9a 100644 --- a/orion-server/src/scheduler.rs +++ b/orion-server/src/scheduler.rs @@ -1,724 +1,522 @@ -use std::collections::{HashMap, VecDeque}; +use crate::model::builds; +use dashmap::DashMap; +use orion::ws::WSMessage; +use rand::Rng; +use sea_orm::{ + prelude::DateTimeUtc, + ActiveModelTrait, ActiveValue::Set, DatabaseConnection, +}; +use serde::{Deserialize, Serialize}; +use std::collections::VecDeque; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::sync::{mpsc, Mutex, RwLock}; -use tokio::time; +use tokio::sync::{mpsc::UnboundedSender, Mutex, Notify}; +use utoipa::ToSchema; use uuid::Uuid; -use serde::Serialize; - -use crate::api::{BuildRequest, WorkerStatus, AppState}; -use orion::ws::WSMessage; -/// State of queued tasks -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum TaskState { - /// Waiting for scheduling - Pending, - /// Assigned to a worker - Assigned(String), // worker_id - /// Currently running - Running(String), // worker_id - /// Completed successfully - Completed, - /// Failed with error - Failed, - /// Timed out - Timeout, - /// Cancelled - Cancelled, +/// Request payload for creating a new build task +#[derive(Debug, Clone, Deserialize, ToSchema)] +pub struct BuildRequest { + pub repo: String, + pub buck_hash: String, + pub buckconfig_hash: String, + pub args: Option>, + pub mr: Option, } -/// Task item in the queue +/// Pending task waiting for dispatch #[derive(Debug, Clone)] -pub struct QueuedTask { +pub struct PendingTask { pub task_id: Uuid, - pub build_request: BuildRequest, - pub target: String, // buck2 targets - pub state: TaskState, + pub request: BuildRequest, + pub target: String, pub created_at: Instant, - pub timeout_duration: Duration, - pub retry_count: u32, - pub max_retries: u32, - pub assigned_worker: Option, - pub priority: TaskPriority, } -/// Task priority levels -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -pub enum TaskPriority { - Low = 0, - Normal = 1, - High = 2, - Critical = 3, +/// Task queue configuration +#[derive(Debug, Clone)] +pub struct TaskQueueConfig { + /// Maximum queue length + pub max_queue_size: usize, + /// Maximum wait time for tasks in queue + pub max_wait_time: Duration, + /// Queue cleanup interval + pub cleanup_interval: Duration, } -impl Default for TaskPriority { +impl Default for TaskQueueConfig { fn default() -> Self { - TaskPriority::Normal + Self { + max_queue_size: 1000, + max_wait_time: Duration::from_secs(300), // 5 minutes + cleanup_interval: Duration::from_secs(30), // Cleanup every 30 seconds + } } } -impl QueuedTask { - /// Create a new QueuedTask from BuildRequest - pub fn new( - build_request: BuildRequest, - target: String, - config: &SchedulerConfig, - priority: Option, - ) -> Self { +/// Simple FIFO task queue +#[derive(Debug)] +pub struct TaskQueue { + /// Queue storage (FIFO) + queue: VecDeque, + /// Queue configuration + config: TaskQueueConfig, +} + +impl TaskQueue { + pub fn new(config: TaskQueueConfig) -> Self { Self { - task_id: Uuid::now_v7(), - build_request, - target, - state: TaskState::Pending, - created_at: Instant::now(), - timeout_duration: config.default_task_timeout, - retry_count: 0, - max_retries: config.default_max_retries, - assigned_worker: None, - priority: priority.unwrap_or_default(), + queue: VecDeque::new(), + config, } } - /// Check if task has timed out - pub fn is_timed_out(&self) -> bool { - Instant::now().duration_since(self.created_at) > self.timeout_duration - } + /// Add task to the end of queue + pub fn enqueue(&mut self, task: PendingTask) -> Result<(), String> { + // Check if queue is full + if self.queue.len() >= self.config.max_queue_size { + return Err("Queue is full".to_string()); + } - /// Check if task can be retried - pub fn can_retry(&self) -> bool { - self.retry_count < self.max_retries + self.queue.push_back(task); + Ok(()) } - /// Reset task for retry - pub fn reset_for_retry(&mut self) { - self.retry_count += 1; - self.state = TaskState::Pending; - self.assigned_worker = None; - self.created_at = Instant::now(); + /// Remove task from the front of queue + pub fn dequeue(&mut self) -> Option { + self.queue.pop_front() } -} -/// Scheduler configuration -#[derive(Debug, Clone)] -pub struct SchedulerConfig { - /// Maximum queue length - pub max_queue_length: usize, - /// Default task timeout duration - pub default_task_timeout: Duration, - /// Default maximum retry count - pub default_max_retries: u32, - /// Scheduler check interval - pub scheduler_interval: Duration, - /// Task cleanup interval - pub cleanup_interval: Duration, - /// Retention time for completed tasks - pub completed_task_retention: Duration, -} + /// Clean up expired tasks + 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()); + false + } else { + true + } + }); + + expired_tasks + } -impl Default for SchedulerConfig { - fn default() -> Self { - Self { - max_queue_length: 1000, - default_task_timeout: Duration::from_secs(3600), // 1 hour - default_max_retries: 3, - scheduler_interval: Duration::from_millis(100), - cleanup_interval: Duration::from_secs(60), - completed_task_retention: Duration::from_secs(3600), // 1 hour + /// Get queue statistics + pub fn get_stats(&self) -> TaskQueueStats { + TaskQueueStats { + total_queued: self.queue.len(), + oldest_task_age_seconds: self.queue.front() + .map(|task| Instant::now().duration_since(task.created_at).as_secs()), } } } -/// Scheduler statistics -#[derive(Debug, Clone, Serialize)] -pub struct SchedulerStats { - pub pending_tasks: usize, - pub running_tasks: usize, - pub completed_tasks: usize, - pub failed_tasks: usize, - pub timeout_tasks: usize, - pub total_workers: usize, - pub idle_workers: usize, - pub busy_workers: usize, +/// Queue statistics +#[derive(Debug, Clone, Serialize, ToSchema)] +pub struct TaskQueueStats { + pub total_queued: usize, + /// Age of oldest task in seconds + pub oldest_task_age_seconds: Option, } -/// Task scheduler -pub struct TaskScheduler { - /// Task queue sorted by priority - task_queue: Arc>>, - /// Registry of all task states - task_registry: Arc>>, - /// Scheduler configuration - config: SchedulerConfig, - /// Application state reference - app_state: AppState, - /// Scheduler control channel - control_tx: mpsc::UnboundedSender, - control_rx: Arc>>, +/// Information about an active build task +#[derive(Clone)] +pub struct BuildInfo { + 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 } -/// Scheduler commands +/// Information about a connected worker #[derive(Debug)] -pub enum SchedulerCommand { - /// Add new task - AddTask { - task: QueuedTask, - response_tx: tokio::sync::oneshot::Sender>, - }, - /// Cancel task - CancelTask { - task_id: Uuid, - response_tx: tokio::sync::oneshot::Sender>, - }, - /// Task completion notification - TaskCompleted { - task_id: Uuid, - success: bool, - worker_id: String, - }, - /// Worker status update - WorkerStatusUpdate { - worker_id: String, - status: TaskState, - }, - /// Get statistics - GetStats { - response_tx: tokio::sync::oneshot::Sender, - }, - /// Get task by ID - GetTask { - task_id: Uuid, - response_tx: tokio::sync::oneshot::Sender>, - }, - /// Stop scheduler - Stop, +pub struct WorkerInfo { + pub sender: UnboundedSender, + pub status: WorkerStatus, + pub last_heartbeat: DateTimeUtc, } -/// Scheduler error types -#[derive(Debug, thiserror::Error)] -pub enum SchedulerError { - #[error("Queue is full")] - QueueFull, - #[error("Task not found: {0}")] - TaskNotFound(Uuid), - #[error("Task already exists: {0}")] - TaskExists(Uuid), - #[error("No available workers")] - NoWorkers, - #[error("Worker not found: {0}")] - WorkerNotFound(String), - #[error("Invalid task state transition from {from:?} to {to:?}")] - InvalidStateTransition { from: TaskState, to: TaskState }, - #[error("Internal error: {0}")] - Internal(String), +/// Task scheduler - manages task queue and worker assignment +#[derive(Clone)] +pub struct TaskScheduler { + /// Pending task queue + pub pending_tasks: Arc>, + /// Event notifier for new tasks or available workers + pub task_notifier: Arc, + /// Worker information + pub workers: Arc>, + /// Active build tasks + pub active_builds: Arc>, + /// Database connection + pub conn: DatabaseConnection, } impl TaskScheduler { - /// Create new task scheduler - pub fn new(config: SchedulerConfig, app_state: AppState) -> Self { - let (control_tx, control_rx) = mpsc::unbounded_channel(); - + /// Create new task scheduler instance + pub fn new( + conn: DatabaseConnection, + workers: Arc>, + active_builds: Arc>, + queue_config: Option, + ) -> Self { + let config = queue_config.unwrap_or_default(); Self { - task_queue: Arc::new(Mutex::new(VecDeque::new())), - task_registry: Arc::new(RwLock::new(HashMap::new())), - config, - app_state, - control_tx, - control_rx: Arc::new(Mutex::new(control_rx)), + pending_tasks: Arc::new(Mutex::new(TaskQueue::new(config))), + task_notifier: Arc::new(Notify::new()), + workers, + active_builds, + conn, } } - /// Get scheduler control handle - pub fn get_handle(&self) -> SchedulerHandle { - SchedulerHandle { - control_tx: self.control_tx.clone(), + /// Add task to queue + pub async fn enqueue_task(&self, request: BuildRequest, target: String) -> Result { + let task_id = Uuid::now_v7(); + + let pending_task = PendingTask { + task_id, + request, + target, + created_at: Instant::now(), + }; + + { + 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) } - /// Start scheduler - pub async fn run(&self) { - let scheduler_task = self.run_scheduler_loop(); - let cleanup_task = self.run_cleanup_loop(); - - tokio::select! { - _ = scheduler_task => { - tracing::info!("Scheduler loop ended"); - } - _ = cleanup_task => { - tracing::info!("Cleanup loop ended"); - } - } + /// Get queue statistics + pub async fn get_queue_stats(&self) -> TaskQueueStats { + let queue = self.pending_tasks.lock().await; + queue.get_stats() } - /// Main scheduler loop - async fn run_scheduler_loop(&self) { - let mut interval = time::interval(self.config.scheduler_interval); - let mut control_rx = self.control_rx.lock().await; - - loop { - tokio::select! { - _ = interval.tick() => { - if let Err(e) = self.process_queue().await { - tracing::error!("Error processing queue: {}", e); - } - - if let Err(e) = self.check_timeouts().await { - tracing::error!("Error checking timeouts: {}", e); - } - } - - command = control_rx.recv() => { - match command { - Some(cmd) => { - if let SchedulerCommand::Stop = cmd { - tracing::info!("Scheduler stopping"); - break; - } - self.handle_command(cmd).await; - } - None => { - tracing::info!("Control channel closed, stopping scheduler"); - break; - } - } - } - } - } + /// Clean up expired tasks + pub async fn cleanup_expired_tasks(&self) -> Vec { + let mut queue = self.pending_tasks.lock().await; + queue.cleanup_expired() } - /// Process tasks in the queue - async fn process_queue(&self) -> Result<(), SchedulerError> { - let mut queue = self.task_queue.lock().await; - - // Get available workers - let idle_workers: Vec = self.app_state - .workers + /// Check if there are available workers + pub fn has_idle_workers(&self) -> bool { + self.workers + .iter() + .any(|entry| matches!(entry.value().status, WorkerStatus::Idle)) + } + + /// Get list of idle workers + pub fn get_idle_workers(&self) -> Vec { + self.workers .iter() .filter(|entry| matches!(entry.value().status, WorkerStatus::Idle)) .map(|entry| entry.key().clone()) - .collect(); - + .collect() + } + + /// Try to dispatch queued tasks (concurrent safe) + pub async fn process_pending_tasks(&self) { + // Get available workers + let idle_workers = self.get_idle_workers(); if idle_workers.is_empty() { - return Ok(()); + return; } - // Try to assign tasks to available workers - let mut assigned_count = 0; - let mut worker_iter = idle_workers.iter().cycle(); + // 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); - while assigned_count < idle_workers.len() && !queue.is_empty() { - if let Some(mut task) = queue.pop_front() { - if matches!(task.state, TaskState::Pending) { - if let Some(worker_id) = worker_iter.next() { - // Assign task to worker - match self.assign_task_to_worker(&mut task, worker_id.clone()).await { - Ok(()) => { - // Update task state - let mut registry = self.task_registry.write().await; - registry.insert(task.task_id, task); - assigned_count += 1; - } - Err(e) => { - tracing::error!("Failed to assign task {} to worker {}: {}", - task.task_id, worker_id, e); - // Put task back in queue - queue.push_front(task); - break; - } - } - } + // Batch dequeue tasks + { + let mut queue = self.pending_tasks.lock().await; + for _ in 0..max_tasks { + if let Some(task) = queue.dequeue() { + tasks_to_dispatch.push(task); } else { - // Task state is not Pending, put back in queue - queue.push_back(task); + break; } } } - - Ok(()) + + // 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); + } + }) + }).collect(); + + // Wait for all dispatch tasks to complete + for future in dispatch_futures { + let _ = future.await; + } + } } - /// Assign task to specified worker - async fn assign_task_to_worker(&self, task: &mut QueuedTask, worker_id: String) -> Result<(), SchedulerError> { + /// Dispatch single task + async fn dispatch_task(&self, pending_task: PendingTask) -> Result<(), String> { + let idle_workers = self.get_idle_workers(); + if idle_workers.is_empty() { + return Err("No idle workers available".to_string()); + } + + // 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(); + + // Create log file + 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); + return Err(format!("Failed to create log file: {e}")); + } + }; + + // Create build information + let build_info = BuildInfo { + repo: pending_task.request.repo.clone(), + target: pending_task.target.clone(), + args: pending_task.request.args.clone(), + start_at: chrono::Utc::now(), + mr: pending_task.request.mr.clone(), + _worker_id: chosen_id.clone(), + log_file, + }; + + // Save to database + let model = builds::ActiveModel { + build_id: Set(pending_task.task_id), + output_file: Set(format!("{}/{}", get_build_log_dir(), pending_task.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(&self.conn).await { + tracing::error!("Failed to insert queued task into DB: {}", e); + return Err(format!("Failed to create task in database: {e}")); + } + // Create WebSocket message let msg = WSMessage::Task { - id: task.task_id.to_string(), - repo: task.build_request.repo.clone(), - target: task.target.clone(), - args: task.build_request.args.clone(), - mr: task.build_request.mr.clone().unwrap_or_default(), + id: pending_task.task_id.to_string(), + repo: pending_task.request.repo, + target: pending_task.target, + args: pending_task.request.args, + mr: pending_task.request.mr.unwrap_or_default(), }; // Send task to worker - if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { + if let Some(mut worker) = self.workers.get_mut(&chosen_id) { if worker.sender.send(msg).is_ok() { - // Update worker status - worker.status = WorkerStatus::Busy(task.task_id.to_string()); - - // Update task state - task.state = TaskState::Assigned(worker_id.clone()); - task.assigned_worker = Some(worker_id.clone()); - - tracing::info!("Task {} assigned to worker {}", task.task_id, worker_id); + 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); Ok(()) } else { - Err(SchedulerError::Internal(format!("Failed to send task to worker {}", worker_id))) + Err(format!("Failed to send task to worker {chosen_id}")) } } else { - Err(SchedulerError::Internal(format!("Worker {} not found", worker_id))) + Err(format!("Worker {chosen_id} not found")) } } - /// Check for timed out tasks - async fn check_timeouts(&self) -> Result<(), SchedulerError> { - let now = Instant::now(); - let mut registry = self.task_registry.write().await; - let mut timed_out_tasks = Vec::new(); - - for (task_id, task) in registry.iter() { - if matches!(task.state, TaskState::Assigned(_) | TaskState::Running(_)) { - if now.duration_since(task.created_at) > task.timeout_duration { - timed_out_tasks.push(*task_id); - } - } - } - - for task_id in timed_out_tasks { - if let Some(task) = registry.get_mut(&task_id) { - tracing::warn!("Task {} timed out", task_id); - - // Release worker - if let Some(worker_id) = &task.assigned_worker { - if let Some(mut worker) = self.app_state.workers.get_mut(worker_id) { - worker.status = WorkerStatus::Idle; - } - } - - // Check if task can be retried - if task.retry_count < task.max_retries { - task.retry_count += 1; - task.state = TaskState::Pending; - task.assigned_worker = None; - task.created_at = now; // Reset creation time - - // Re-add to queue - let mut queue = self.task_queue.lock().await; - self.insert_task_by_priority(&mut queue, task.clone()); - - tracing::info!("Task {} requeued for retry ({}/{})", - task_id, task.retry_count, task.max_retries); - } else { - task.state = TaskState::Timeout; - tracing::error!("Task {} exceeded max retries and timed out", task_id); - } - } - } - - Ok(()) - } - - /// Insert task to queue by priority - fn insert_task_by_priority(&self, queue: &mut VecDeque, task: QueuedTask) { - let position = queue - .iter() - .position(|existing_task| existing_task.priority < task.priority) - .unwrap_or(queue.len()); - - queue.insert(position, task); - } - - /// Handle scheduler commands - async fn handle_command(&self, command: SchedulerCommand) { - match command { - SchedulerCommand::AddTask { task, response_tx } => { - let result = self.add_task_internal(task).await; - let _ = response_tx.send(result); - } - SchedulerCommand::CancelTask { task_id, response_tx } => { - let result = self.cancel_task_internal(task_id).await; - let _ = response_tx.send(result); - } - SchedulerCommand::TaskCompleted { task_id, success, worker_id } => { - self.handle_task_completion(task_id, success, worker_id).await; - } - SchedulerCommand::WorkerStatusUpdate { worker_id, status } => { - self.handle_worker_status_update(worker_id, status).await; - } - SchedulerCommand::GetStats { response_tx } => { - let stats = self.get_stats().await; - let _ = response_tx.send(stats); - } - SchedulerCommand::GetTask { task_id, response_tx } => { - let task = self.get_task_internal(task_id).await; - let _ = response_tx.send(task); - } - SchedulerCommand::Stop => { - // Already handled in main loop - } - } + /// Notify about new task or available worker + pub fn notify_task_available(&self) { + self.task_notifier.notify_one(); } - /// Internal add task logic - async fn add_task_internal(&self, task: QueuedTask) -> Result { - let mut queue = self.task_queue.lock().await; - - // Check if queue is full - if queue.len() >= self.config.max_queue_length { - return Err(SchedulerError::QueueFull); - } - - // Check if task already exists - let registry = self.task_registry.read().await; - if registry.contains_key(&task.task_id) { - return Err(SchedulerError::TaskExists(task.task_id)); - } - drop(registry); - - let task_id = task.task_id; - - // Add to registry - let mut registry = self.task_registry.write().await; - registry.insert(task_id, task.clone()); - drop(registry); - - // Insert by priority - self.insert_task_by_priority(&mut queue, task); - - tracing::info!("Task {} added to queue", task_id); - Ok(task_id) - } + /// Start queue management background task (event-driven + periodic cleanup) + pub async fn start_queue_manager(self) { + let cleanup_interval = { + let queue = self.pending_tasks.lock().await; + queue.config.cleanup_interval + }; - /// Internal cancel task logic - async fn cancel_task_internal(&self, task_id: Uuid) -> Result<(), SchedulerError> { - let mut registry = self.task_registry.write().await; - - if let Some(task) = registry.get_mut(&task_id) { - match &task.state { - TaskState::Pending => { - // Remove from queue - let mut queue = self.task_queue.lock().await; - queue.retain(|t| t.task_id != task_id); - task.state = TaskState::Cancelled; - } - TaskState::Assigned(worker_id) | TaskState::Running(worker_id) => { - // Notify worker to cancel task (if needed) - if let Some(mut worker) = self.app_state.workers.get_mut(worker_id) { - worker.status = WorkerStatus::Idle; + // Task dispatcher: wait for notifications or process periodically + let dispatch_scheduler = self.clone(); + let dispatch_task = tokio::spawn(async move { + loop { + // Wait for notification or timeout + tokio::select! { + // Wait for new task or worker available notification + _ = dispatch_scheduler.task_notifier.notified() => { + dispatch_scheduler.process_pending_tasks().await; + } + // Periodic check (prevent missing notifications) + _ = tokio::time::sleep(Duration::from_secs(5)) => { + dispatch_scheduler.process_pending_tasks().await; } - task.state = TaskState::Cancelled; - } - _ => { - return Err(SchedulerError::Internal("Cannot cancel completed task".to_string())); } } - - tracing::info!("Task {} cancelled", task_id); - Ok(()) - } else { - Err(SchedulerError::TaskNotFound(task_id)) - } - } + }); - /// Handle task completion - async fn handle_task_completion(&self, task_id: Uuid, success: bool, worker_id: String) { - let mut registry = self.task_registry.write().await; - - if let Some(task) = registry.get_mut(&task_id) { - task.state = if success { - TaskState::Completed - } else { - TaskState::Failed - }; + // Cleaner: periodically clean up expired tasks + let cleanup_scheduler = self.clone(); + let cleanup_task = tokio::spawn(async move { + let mut interval = tokio::time::interval(cleanup_interval); - // Release worker - if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { - worker.status = WorkerStatus::Idle; - } - - tracing::info!("Task {} completed with success: {}", task_id, success); - } - } - - /// Handle worker status update - async fn handle_worker_status_update(&self, worker_id: String, status: TaskState) { - if let Some(mut worker) = self.app_state.workers.get_mut(&worker_id) { - match status { - TaskState::Running(task_id) => { - worker.status = WorkerStatus::Busy(task_id.clone()); - // Update task state in registry - let mut registry = self.task_registry.write().await; - if let Some(task) = registry.get_mut(&task_id.parse::().unwrap_or_default()) { - task.state = TaskState::Running(worker_id); + 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()); + + // Log expired task information + for task in expired_tasks { + tracing::debug!("Expired task: {} ({})", task.task_id, task.request.repo); } } - _ => { - // Other status updates can be handled here - tracing::debug!("Worker {} status updated to {:?}", worker_id, status); - } } - } - } + }); - /// Get task by ID (internal method) - async fn get_task_internal(&self, task_id: Uuid) -> Option { - let registry = self.task_registry.read().await; - registry.get(&task_id).cloned() - } - - /// Get statistics - async fn get_stats(&self) -> SchedulerStats { - let registry = self.task_registry.read().await; - - let mut stats = SchedulerStats { - pending_tasks: 0, - running_tasks: 0, - completed_tasks: 0, - failed_tasks: 0, - timeout_tasks: 0, - total_workers: self.app_state.workers.len(), - idle_workers: 0, - busy_workers: 0, - }; - - // Count task states - for task in registry.values() { - match task.state { - TaskState::Pending => stats.pending_tasks += 1, - TaskState::Assigned(_) | TaskState::Running(_) => stats.running_tasks += 1, - TaskState::Completed => stats.completed_tasks += 1, - TaskState::Failed => stats.failed_tasks += 1, - TaskState::Timeout => stats.timeout_tasks += 1, - TaskState::Cancelled => {} // Not counted in statistics + // Wait for tasks to complete (actually runs forever) + tokio::select! { + _ = dispatch_task => { + tracing::error!("Task dispatcher unexpectedly stopped"); } - } - - // Count worker states - for worker in self.app_state.workers.iter() { - match worker.value().status { - WorkerStatus::Idle => stats.idle_workers += 1, - WorkerStatus::Busy(_) => stats.busy_workers += 1, + _ = cleanup_task => { + tracing::error!("Task cleanup unexpectedly stopped"); } } - - stats } +} - /// Cleanup loop - async fn run_cleanup_loop(&self) { - let mut interval = time::interval(self.config.cleanup_interval); - - loop { - interval.tick().await; - self.cleanup_completed_tasks().await; - } - } +/// Get build log directory +pub fn get_build_log_dir() -> &'static str { + use once_cell::sync::Lazy; + static BUILD_LOG_DIR: Lazy = + Lazy::new(|| std::env::var("BUILD_LOG_DIR").expect("BUILD_LOG_DIR must be set")); + &BUILD_LOG_DIR +} - /// Clean up completed tasks - async fn cleanup_completed_tasks(&self) { - let now = Instant::now(); - let retention = self.config.completed_task_retention; - - let mut registry = self.task_registry.write().await; - let mut to_remove = Vec::new(); - - for (task_id, task) in registry.iter() { - if matches!(task.state, TaskState::Completed | TaskState::Failed | TaskState::Timeout | TaskState::Cancelled) { - if now.duration_since(task.created_at) > retention { - to_remove.push(*task_id); - } - } - } - - for task_id in to_remove { - registry.remove(&task_id); - tracing::debug!("Cleaned up completed task {}", task_id); - } +/// Create log file +pub fn create_log_file(task_id: &str) -> Result { + let log_path = format!("{}/{}", get_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)?; } -} -/// Scheduler control handle -#[derive(Clone)] -pub struct SchedulerHandle { - control_tx: mpsc::UnboundedSender, + // Create or open the log file in append mode + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) } -impl SchedulerHandle { - /// Add task to queue - pub async fn add_task(&self, task: QueuedTask) -> Result { - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - - self.control_tx.send(SchedulerCommand::AddTask { - task, - response_tx, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; - - response_rx.await - .map_err(|_| SchedulerError::Internal("Response channel closed".to_string()))? - } +#[cfg(test)] +mod tests { + use super::*; - /// Cancel task - pub async fn cancel_task(&self, task_id: Uuid) -> Result<(), SchedulerError> { - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + /// Test task queue basic functionality + #[test] + fn test_task_queue_fifo() { + let config = TaskQueueConfig::default(); + let mut queue = TaskQueue::new(config); - self.control_tx.send(SchedulerCommand::CancelTask { - task_id, - response_tx, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + // Create test tasks + let task1 = PendingTask { + task_id: Uuid::now_v7(), + request: BuildRequest { + repo: "test1".to_string(), + buck_hash: "hash1".to_string(), + buckconfig_hash: "config1".to_string(), + args: None, + mr: None, + }, + target: "target1".to_string(), + created_at: Instant::now(), + }; - response_rx.await - .map_err(|_| SchedulerError::Internal("Response channel closed".to_string()))? - } - - /// Notify task completed - pub async fn notify_task_completed(&self, task_id: Uuid, success: bool, worker_id: String) -> Result<(), SchedulerError> { - self.control_tx.send(SchedulerCommand::TaskCompleted { - task_id, - success, - worker_id, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + let task2 = PendingTask { + task_id: Uuid::now_v7(), + request: BuildRequest { + repo: "test2".to_string(), + buck_hash: "hash2".to_string(), + buckconfig_hash: "config2".to_string(), + args: None, + mr: None, + }, + target: "target2".to_string(), + created_at: Instant::now(), + }; - Ok(()) - } - - /// Update worker status - pub async fn update_worker_status(&self, worker_id: String, status: TaskState) -> Result<(), SchedulerError> { - self.control_tx.send(SchedulerCommand::WorkerStatusUpdate { - worker_id, - status, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + // Test FIFO behavior + assert!(queue.enqueue(task1.clone()).is_ok()); + assert!(queue.enqueue(task2.clone()).is_ok()); - Ok(()) - } - - /// Get task by ID - pub async fn get_task(&self, task_id: Uuid) -> Result, SchedulerError> { - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - self.control_tx.send(SchedulerCommand::GetTask { - task_id, - response_tx, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + 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"); - response_rx.await - .map_err(|_| SchedulerError::Internal("Response channel closed".to_string())) } - /// Get statistics - pub async fn get_stats(&self) -> Result { - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + /// Test queue capacity limit + #[test] + fn test_queue_capacity() { + let config = TaskQueueConfig { + max_queue_size: 2, + max_wait_time: Duration::from_secs(60), + cleanup_interval: Duration::from_secs(30), + }; + let mut queue = TaskQueue::new(config); - self.control_tx.send(SchedulerCommand::GetStats { - response_tx, - }).map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + let task = PendingTask { + task_id: Uuid::now_v7(), + request: BuildRequest { + repo: "test".to_string(), + buck_hash: "hash".to_string(), + buckconfig_hash: "config".to_string(), + args: None, + mr: None, + }, + target: "target".to_string(), + created_at: Instant::now(), + }; - response_rx.await - .map_err(|_| SchedulerError::Internal("Response channel closed".to_string())) - } - - /// Stop scheduler - pub async fn stop(&self) -> Result<(), SchedulerError> { - self.control_tx.send(SchedulerCommand::Stop) - .map_err(|_| SchedulerError::Internal("Scheduler not running".to_string()))?; + // Fill queue to capacity + assert!(queue.enqueue(task.clone()).is_ok()); + assert!(queue.enqueue(task.clone()).is_ok()); - 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 f7264994f..144866dac 100644 --- a/orion-server/src/server.rs +++ b/orion-server/src/server.rs @@ -2,13 +2,11 @@ use crate::api::{self, AppState}; use crate::model::builds; use axum::Router; use axum::routing::get; -use dashmap::DashMap; 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; @@ -24,7 +22,7 @@ use utoipa_swagger_ui::SwaggerUi; api::task_query_by_mr, ), components( - schemas(api::BuildRequest, api::TaskStatus, api::TaskStatusEnum, api::BuildDTO) + schemas(crate::scheduler::BuildRequest, api::TaskStatus, api::TaskStatusEnum, api::BuildDTO) ), tags( (name = "Build", description = "Build related endpoints") @@ -41,14 +39,13 @@ pub async fn start_server(port: u16) { .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()), - }; + let state = AppState::new(conn, None); // Start background health check task tokio::spawn(start_health_check_task(state.clone())); + + // Start queue manager + tokio::spawn(api::start_queue_manager(state.clone())); let app = Router::new() .route("/", get(|| async { "Hello, World!" })) @@ -103,7 +100,7 @@ async fn start_health_check_task(state: AppState) { let now = chrono::Utc::now(); // Find workers that haven't sent heartbeat within timeout period - for entry in state.workers.iter() { + for entry in state.scheduler.workers.iter() { if now.signed_duration_since(entry.value().last_heartbeat) > chrono::Duration::from_std(worker_timeout).unwrap() { @@ -119,17 +116,17 @@ async fn start_health_check_task(state: AppState) { // Remove dead workers and handle their tasks for worker_id in dead_workers { - if let Some((_, worker_info)) = state.workers.remove(&worker_id) { + if let Some((_, worker_info)) = state.scheduler.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 { + if let crate::scheduler::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); + state.scheduler.active_builds.remove(&task_id); let update_res = builds::Entity::update_many() .set(builds::ActiveModel { From 82eb941572a75ffb053462d236d4318c1934d226 Mon Sep 17 00:00:00 2001 From: Xiaoyang Han Date: Mon, 11 Aug 2025 10:19:24 +0800 Subject: [PATCH 5/5] fix dep Signed-off-by: Xiaoyang Han --- orion-server/src/api.rs | 58 ----------------------------------------- 1 file changed, 58 deletions(-) diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index e4db1a353..543fdf022 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -38,64 +38,6 @@ use tokio::sync::mpsc::UnboundedSender; use utoipa::ToSchema; 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, - buck_hash: String, - buckconfig_hash: String, - args: Option>, - mr: Option, -} - -/// Information about an active build task -#[derive(Clone)] -pub struct BuildInfo { - 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)]