diff --git a/orion-server/Cargo.toml b/orion-server/Cargo.toml index 142495dd3..654915fec 100644 --- a/orion-server/Cargo.toml +++ b/orion-server/Cargo.toml @@ -32,4 +32,4 @@ utoipa.workspace = true utoipa-swagger-ui = { workspace = true, features = ["axum"] } chrono = { version = "0.4", features = ["serde"] } reqwest.workspace = true - +thiserror = "2.0" \ No newline at end of file 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..543fdf022 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,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)] @@ -119,9 +63,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 +94,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 +134,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 +209,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 +223,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 +259,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 +281,66 @@ 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 +382,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 +409,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 +509,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 +539,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 +548,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 +569,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 +606,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 +618,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/main.rs b/orion-server/src/main.rs index 2db33d7cb..ad1c37485 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -2,6 +2,8 @@ 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..948fa9c9a --- /dev/null +++ b/orion-server/src/scheduler.rs @@ -0,0 +1,522 @@ +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::UnboundedSender, Mutex, Notify}; +use utoipa::ToSchema; +use uuid::Uuid; + +/// 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, +} + +/// Pending task waiting for dispatch +#[derive(Debug, Clone)] +pub struct PendingTask { + pub task_id: Uuid, + pub request: BuildRequest, + pub target: String, + pub created_at: Instant, +} + +/// 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 TaskQueueConfig { + fn default() -> Self { + Self { + max_queue_size: 1000, + max_wait_time: Duration::from_secs(300), // 5 minutes + cleanup_interval: Duration::from_secs(30), // Cleanup every 30 seconds + } + } +} + +/// 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 { + queue: VecDeque::new(), + config, + } + } + + /// 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()); + } + + self.queue.push_back(task); + Ok(()) + } + + /// Remove task from the front of queue + pub fn dequeue(&mut self) -> Option { + self.queue.pop_front() + } + + /// 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 + } + + /// 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()), + } + } +} + +/// 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, +} + +/// 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, +} + +/// 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 instance + pub fn new( + conn: DatabaseConnection, + workers: Arc>, + active_builds: Arc>, + queue_config: Option, + ) -> Self { + let config = queue_config.unwrap_or_default(); + Self { + pending_tasks: Arc::new(Mutex::new(TaskQueue::new(config))), + task_notifier: Arc::new(Notify::new()), + workers, + active_builds, + conn, + } + } + + /// 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) + } + + /// Get queue statistics + pub async fn get_queue_stats(&self) -> TaskQueueStats { + let queue = self.pending_tasks.lock().await; + queue.get_stats() + } + + /// Clean up expired tasks + pub async fn cleanup_expired_tasks(&self) -> Vec { + let mut queue = self.pending_tasks.lock().await; + queue.cleanup_expired() + } + + /// 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() + } + + /// 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; + } + + // Process tasks in batches, up to the number of idle workers + let max_tasks = idle_workers.len(); + let mut tasks_to_dispatch = Vec::with_capacity(max_tasks); + + // Batch dequeue tasks + { + let mut queue = self.pending_tasks.lock().await; + for _ in 0..max_tasks { + if let Some(task) = queue.dequeue() { + tasks_to_dispatch.push(task); + } else { + break; + } + } + } + + // 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; + } + } + } + + /// 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: 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.workers.get_mut(&chosen_id) { + if worker.sender.send(msg).is_ok() { + worker.status = WorkerStatus::Busy(pending_task.task_id.to_string()); + self.active_builds.insert(pending_task.task_id.to_string(), build_info); + tracing::info!("Queued task {} dispatched to worker {}", pending_task.task_id, chosen_id); + Ok(()) + } else { + Err(format!("Failed to send task to worker {chosen_id}")) + } + } else { + Err(format!("Worker {chosen_id} not found")) + } + } + + /// Notify about new task or available worker + pub fn notify_task_available(&self) { + self.task_notifier.notify_one(); + } + + /// 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 + }; + + // 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; + } + } + } + }); + + // 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); + + 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); + } + } + } + }); + + // Wait for tasks to complete (actually runs forever) + tokio::select! { + _ = dispatch_task => { + tracing::error!("Task dispatcher unexpectedly stopped"); + } + _ = cleanup_task => { + tracing::error!("Task cleanup unexpectedly stopped"); + } + } + } +} + +/// 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 +} + +/// 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)?; + } + + // Create or open the log file in append mode + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Test task queue basic functionality + #[test] + fn test_task_queue_fifo() { + let config = TaskQueueConfig::default(); + let mut queue = TaskQueue::new(config); + + // Create test tasks + let task1 = PendingTask { + task_id: Uuid::now_v7(), + 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(), + }; + + 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(), + }; + + // Test FIFO behavior + assert!(queue.enqueue(task1.clone()).is_ok()); + assert!(queue.enqueue(task2.clone()).is_ok()); + + + let dequeued1 = queue.dequeue().unwrap(); + assert_eq!(dequeued1.task_id, task1.task_id); + assert_eq!(dequeued1.request.repo, "test1"); + + let dequeued2 = queue.dequeue().unwrap(); + assert_eq!(dequeued2.task_id, task2.task_id); + assert_eq!(dequeued2.request.repo, "test2"); + + } + + /// Test queue capacity limit + #[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); + + 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(), + }; + + // Fill queue to capacity + assert!(queue.enqueue(task.clone()).is_ok()); + assert!(queue.enqueue(task.clone()).is_ok()); + + // Should fail when full + assert!(queue.enqueue(task).is_err()); + } +} diff --git a/orion-server/src/server.rs b/orion-server/src/server.rs index 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 {