diff --git a/api-model/src/buck2/api.rs b/api-model/src/buck2/api.rs index 878f8c598..a8f0dce84 100644 --- a/api-model/src/buck2/api.rs +++ b/api-model/src/buck2/api.rs @@ -14,7 +14,6 @@ use utoipa::ToSchema; use crate::buck2::{status::Status, types::ProjectRelativePath}; /// Parameters required to build a task. -#[allow(dead_code)] #[derive(Debug, Deserialize, Serialize, ToSchema)] pub struct TaskBuildRequest { /// The repository base path @@ -47,7 +46,6 @@ pub struct RetryBuildRequest { } /// Result of a task build operation containing status and metadata. Used by Orion-Server -#[allow(dead_code)] #[derive(Debug, Deserialize, Serialize, ToSchema)] pub struct OrionBuildResult { /// Unique identifier for the build task diff --git a/api-model/src/buck2/types.rs b/api-model/src/buck2/types.rs index f483a3b97..0609eed29 100644 --- a/api-model/src/buck2/types.rs +++ b/api-model/src/buck2/types.rs @@ -5,7 +5,6 @@ use serde::{Deserialize, Serialize}; use utoipa::ToSchema; /// Task phase when in buck2 build -#[allow(dead_code)] #[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq, Eq)] pub enum TaskPhase { DownloadingSource, diff --git a/api-model/src/buck2/ws.rs b/api-model/src/buck2/ws.rs index b602af060..7325a13f8 100644 --- a/api-model/src/buck2/ws.rs +++ b/api-model/src/buck2/ws.rs @@ -8,7 +8,6 @@ use crate::buck2::{ }; /// Message protocol for WebSocket communication between worker and server. -#[allow(dead_code)] #[derive(Debug, Serialize, Deserialize, Clone)] #[serde(tag = "type")] pub enum WSMessage { diff --git a/ceres/src/model/change_list.rs b/ceres/src/model/change_list.rs index 672ef518b..729036d78 100644 --- a/ceres/src/model/change_list.rs +++ b/ceres/src/model/change_list.rs @@ -302,7 +302,7 @@ pub enum RequirementsState { } #[derive(Debug, Clone, serde::Deserialize, serde::Serialize, utoipa::ToSchema)] -#[allow(dead_code)] + pub struct VerifyClPayload { pub assignees: Vec, } @@ -386,12 +386,6 @@ impl ClDiffFile { } } } -#[derive(Serialize)] -pub struct BuckFile { - pub buck: ObjectHash, - pub buck_config: ObjectHash, - pub path: PathBuf, -} #[cfg(test)] mod tests { diff --git a/ceres/src/pack/monorepo.rs b/ceres/src/pack/monorepo.rs index ab7f2aa8c..d192ae22f 100644 --- a/ceres/src/pack/monorepo.rs +++ b/ceres/src/pack/monorepo.rs @@ -1,6 +1,6 @@ use std::{ collections::{HashMap, HashSet}, - path::{Component, Path, PathBuf}, + path::{Component, PathBuf}, str::FromStr, sync::{ Arc, @@ -39,9 +39,9 @@ use tokio::sync::{RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; use crate::{ - api_service::{ApiHandler, cache::GitObjectCache, mono_api_service::MonoApiService, tree_ops}, + api_service::{ApiHandler, cache::GitObjectCache, mono_api_service::MonoApiService}, code_edit::{on_push::OnpushCodeEdit, utils::get_changed_files}, - model::change_list::{BuckFile, ClDiffFile}, + model::change_list::ClDiffFile, pack::RepoHandler, protocol::import_refs::{RefCommand, Refs}, }; @@ -589,93 +589,6 @@ impl MonoRepo { Ok(cl_link) } - #[allow(dead_code)] - async fn search_buck_under_cl(&self, cl_path: &Path) -> Result, MegaError> { - let mut res = vec![]; - let mono_stg = self.storage.mono_storage(); - let mono_api_service: MonoApiService = self.into(); - - let mut path = Some(cl_path); - let mut path_q = Vec::new(); - while let Some(p) = path { - path_q.push(p); - path = p.parent(); - } - if path_q.len() > 2 { - path_q.pop(); - path_q.pop(); - - let p = path_q[path_q.len() - 1]; - if p.parent().is_some() - && let Some(tree) = tree_ops::search_tree_by_path(&mono_api_service, p, None) - .await - .ok() - .flatten() - && let Some(buck) = self.try_extract_buck(tree, cl_path) - { - return Ok(vec![buck]); - }; - return Ok(vec![]); - } - - let mut search_trees: Vec<(PathBuf, Tree)> = vec![]; - - let diff_trees = self.diff_trees_from_cl().await?; - for (path, new, old) in diff_trees { - match (new, old) { - (None, _) => { - continue; - } - (Some(sha1), _) => { - let tree = mono_stg.get_tree_by_hash(&sha1.to_string()).await?.unwrap(); - search_trees.push((path, Tree::from_mega_model(tree))); - } - } - } - - for (path, tree) in search_trees { - if let Some(buck) = self.try_extract_buck(tree, &cl_path.join(path)) { - res.push(buck); - } - } - - Ok(res) - } - - fn try_extract_buck(&self, tree: Tree, cl_path: &Path) -> Option { - let mut buck = None; - let mut buck_config = None; - for item in tree.tree_items { - if item.is_blob() && item.name == "BUCK" { - buck = Some(item.id) - } - if item.is_blob() && item.name == ".buckconfig" { - buck_config = Some(item.id) - } - } - match (buck, buck_config) { - (Some(buck), Some(buck_config)) => Some(BuckFile { - buck, - buck_config, - path: cl_path.to_path_buf(), - }), - _ => None, - } - } - - async fn diff_trees_from_cl( - &self, - ) -> Result, Option)>, MegaError> { - let mono_stg = self.storage.mono_storage(); - let from_c = mono_stg.get_commit_by_hash(&self.from_hash).await?.unwrap(); - let from_tree: Tree = - Tree::from_mega_model(mono_stg.get_tree_by_hash(&from_c.tree).await?.unwrap()); - let to_c = mono_stg.get_commit_by_hash(&self.to_hash).await?.unwrap(); - let to_tree: Tree = - Tree::from_mega_model(mono_stg.get_tree_by_hash(&to_c.tree).await?.unwrap()); - diff_trees(&to_tree, &from_tree) - } - pub fn username(&self) -> String { self.username.clone().unwrap_or(String::from("Anonymous")) } @@ -841,35 +754,3 @@ impl MonoRepo { Ok(()) } } - -#[allow(dead_code)] -type DiffResult = Vec<(PathBuf, Option, Option)>; - -#[allow(dead_code)] -fn diff_trees(theirs: &Tree, base: &Tree) -> Result { - let their_items: HashMap<_, _> = get_plain_items(theirs).into_iter().collect(); - let base_items: HashMap<_, _> = get_plain_items(base).into_iter().collect(); - let all_paths: HashSet<_> = their_items.keys().chain(base_items.keys()).collect(); - - let mut diffs = Vec::new(); - - for path in all_paths { - let their_hash = their_items.get(path).cloned(); - let base_hash = base_items.get(path).cloned(); - if their_hash != base_hash { - diffs.push((path.clone(), their_hash, base_hash)); - } - } - Ok(diffs) -} - -#[allow(dead_code)] -fn get_plain_items(tree: &Tree) -> Vec<(PathBuf, ObjectHash)> { - let mut items = Vec::new(); - for item in tree.tree_items.iter() { - if item.is_tree() { - items.push((PathBuf::from(item.name.clone()), item.id)); - } - } - items -} diff --git a/io-orbit/src/bin/migrate_local_to_s3.rs b/io-orbit/src/bin/migrate_local_to_s3.rs index e134ae0e4..fbcc214e9 100644 --- a/io-orbit/src/bin/migrate_local_to_s3.rs +++ b/io-orbit/src/bin/migrate_local_to_s3.rs @@ -290,11 +290,11 @@ async fn migrate_all( // Periodically await some tasks to keep the number of in-memory // JoinHandles bounded. Semaphore still enforces the true I/O // concurrency; this only caps bookkeeping overhead. - if tasks.len() >= concurrency.saturating_mul(4).max(64) { - if let Some(t) = tasks.pop() { - t.await - .map_err(|e| MegaError::Other(format!("migration task panicked: {e}")))??; - } + if tasks.len() >= concurrency.saturating_mul(4).max(64) + && let Some(t) = tasks.pop() + { + t.await + .map_err(|e| MegaError::Other(format!("migration task panicked: {e}")))??; } } diff --git a/mono/src/api/guard/cedar_guard.rs b/mono/src/api/guard/cedar_guard.rs index 91af0e368..cb9a9d5b1 100644 --- a/mono/src/api/guard/cedar_guard.rs +++ b/mono/src/api/guard/cedar_guard.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, path::Path, str::FromStr}; +use std::{collections::HashMap, str::FromStr}; use axum::{ extract::{FromRef, FromRequestParts, Request, State}, @@ -216,27 +216,6 @@ async fn authorize( Ok(()) } -#[allow(dead_code)] -async fn get_blob_string(state: &MonoApiServiceState, path: &Path) -> Result { - // Use main as default branch - let refs = None; - let data = state - .api_handler(path.as_ref()) - .await? - .get_blob_as_string(path.into(), refs) - .await?; - - match data { - Some(content) => Ok(content), - None => { - Err(MegaError::Other(format!( - "Blob not found at path: {}", - path.display() - ))) - }?, - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/mono/src/git_protocol/ssh.rs b/mono/src/git_protocol/ssh.rs index 0cbee0a0f..425fb54b6 100644 --- a/mono/src/git_protocol/ssh.rs +++ b/mono/src/git_protocol/ssh.rs @@ -21,7 +21,7 @@ use tokio::{io::AsyncReadExt, sync::Mutex}; use crate::git_protocol::http::search_subsequence; type ClientMap = HashMap<(usize, ChannelId), Channel>; -#[allow(dead_code)] + #[derive(Clone)] pub struct SshServer { pub clients: Arc>, diff --git a/orion-server/Cargo.toml b/orion-server/Cargo.toml index 6f60281f1..7fd34866e 100644 --- a/orion-server/Cargo.toml +++ b/orion-server/Cargo.toml @@ -13,13 +13,11 @@ path = "src/main.rs" common = { workspace = true } io-orbit = { workspace = true } api-model = { workspace = true } -jupiter = { workspace = true } callisto = { workspace = true } http = { workspace = true } axum = { workspace = true, features = ["macros", "ws"] } tokio = { workspace = true, features = ["rt-multi-thread", "fs", "process"] } -tokio-retry = "0.3" tokio-stream = { workspace = true, features = ["sync"] } tokio-util = { workspace = true } tracing = { workspace = true } @@ -42,7 +40,6 @@ dashmap = { workspace = true } utoipa.workspace = true utoipa-swagger-ui = { workspace = true, features = ["axum"] } chrono = { workspace = true, features = ["serde"] } -reqwest.workspace = true anyhow = { workspace = true } async-trait = { workspace = true } once_cell = { workspace = true } diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 79d22d2b2..b5690d0ce 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,153 +1,71 @@ -use std::{ - collections::HashMap, convert::Infallible, net::SocketAddr, ops::ControlFlow, sync::Arc, - time::Duration, -}; +use std::net::SocketAddr; use anyhow::Result; -use api_model::buck2::{ - api::{OrionBuildResult, OrionServerResponse, RetryBuildRequest, TaskBuildRequest}, - status::Status, - types::{ - LogErrorResponse, LogEvent, LogLinesResponse, LogReadMode, ProjectRelativePath, - TargetLogLinesResponse, TargetLogQuery, TargetStatusResponse, TaskHistoryQuery, TaskPhase, +use api_model::{ + buck2::{ + api::{RetryBuildRequest, TaskBuildRequest}, + types::{ + LogErrorResponse, LogLinesResponse, TargetLogLinesResponse, TargetLogQuery, + TargetStatusResponse, TaskHistoryQuery, + }, }, - ws::{WSMessage, WSTargetBuildStatusEvent}, + common::{CommonPage, PageParams}, }; use axum::{ Json, Router, - extract::{ - ConnectInfo, Path, Query, State, WebSocketUpgrade, - ws::{Message, Utf8Bytes, WebSocket}, - }, + extract::{ConnectInfo, Path, Query, State, WebSocketUpgrade}, http::StatusCode, - response::{IntoResponse, Sse, sse::Event}, + response::IntoResponse, routing::{any, get, post}, }; -use callisto::{sea_orm_active_enums::OrionTargetStatusEnum, target_build_status}; -use chrono::{FixedOffset, Utc}; -use dashmap::DashMap; -use futures::stream::select; -use futures_util::{SinkExt, Stream, StreamExt}; -use rand::RngExt; -use sea_orm::{ - ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter as _, QueryOrder, - QuerySelect, prelude::DateTimeUtc, -}; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use tokio::sync::{ - RwLock, - mpsc::{self, UnboundedSender}, - watch, -}; -use tokio_stream::wrappers::IntervalStream; -use utoipa::ToSchema; -use uuid::Uuid; use crate::{ - auto_retry::AutoRetryJudger, - log::log_service::LogService, - model::{ - build_events::{BuildEvent, BuildEventDTO}, - build_targets::BuildTarget, - builds, - orion_tasks::{OrionTask, OrionTaskDTO}, - targets::{self, TargetState, TargetWithBuilds}, - tasks, - }, - orion_common::model::{CommonPage, PageParams}, - scheduler::{ - BuildEventPayload, BuildInfo, TaskQueueStats, TaskScheduler, WorkerInfo, WorkerStatus, + app_state::AppState, + model::dto::{ + BuildEventDTO, BuildEventState, BuildTargetDTO, MessageResponse, OrionClientInfo, + OrionClientQuery, OrionClientStatus, OrionTaskDTO, TargetSummaryDTO, TaskInfoDTO, }, - service::target_build_status_service::TargetBuildStatusService, + scheduler::TaskQueueStats, + service::{api_v2_service, ws_service}, }; -const RETRY_COUNT_MAX: i32 = 3; - -/// Enumeration of possible task statuses -#[derive(Debug, Serialize, Default, ToSchema, Clone)] -pub enum TaskStatusEnum { - /// Task is queued and waiting to be assigned to a worker - Pending, - Building, - Interrupted, // Task was interrupted, exit code is None - Failed, - Completed, - #[default] - NotFound, -} - -/// Shared application state containing worker connections, database, and active builds -#[derive(Clone)] -pub struct AppState { - pub scheduler: TaskScheduler, - pub conn: DatabaseConnection, - pub log_service: LogService, - pub target_status_cache: TargetStatusCache, - - shutdown_tx: watch::Sender, +/// Creates and configures all API routes +pub fn routers() -> Router { + Router::new() + .merge(system_routes()) + .merge(task_routes()) + .merge(build_routes()) + .merge(worker_routes()) + .merge(target_status_routes()) } -impl AppState { - /// Create new AppState instance - pub fn new( - conn: DatabaseConnection, - queue_config: Option, - log_service: LogService, - ) -> 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); - let target_status_cache = TargetStatusCache::new(); - - let (shutdown_tx, _) = watch::channel(false); - - Self { - scheduler, - conn, - log_service, - target_status_cache, - shutdown_tx, - } - } - - pub fn start_background_tasks(&self) { - let conn = self.conn.clone(); - let cache = self.target_status_cache.clone(); - let shutdown_rx = self.shutdown_tx.subscribe(); - - tokio::spawn(async move { - cache.auto_flush_loop(conn, shutdown_rx).await; - }); - } +fn system_routes() -> Router { + Router::new() + .route("/ws", any(ws_handler)) + .route("/v2/health", get(health_check_handler)) + .route("/queue-stats", get(queue_stats_handler)) } -/// Creates and configures all API routes -pub fn routers() -> Router { +fn task_routes() -> Router { Router::new() - .route("/ws", any(ws_handler)) .route("/task", post(task_handler)) + .route("/v2/task-handler", get(task_handler_v2)) .route("/task-build-list/{id}", get(task_build_list_handler)) .route("/task-output/{id}", get(task_output_handler)) .route("/task-history-output", get(task_history_output_handler)) - .route("/targets/{target_id}/logs", get(target_logs_handler)) .route("/tasks/{cl}", get(tasks_handler)) .route("/tasks/{task_id}/targets", get(task_targets_handler)) .route( "/tasks/{task_id}/targets/summary", get(task_targets_summary_handler), ) - .route("/queue-stats", get(queue_stats_handler)) - .route("/orion-clients-info", post(get_orion_clients_info)) - .route( - "/orion-client-status/{id}", - get(get_orion_client_status_by_id), - ) - .route("/retry-build", post(build_retry_handler)) - .route("/v2/task-handler", get(task_handler_v2)) - .route("/v2/health", get(health_check_handler)) .route("/v2/task-retry/{id}", post(task_retry_handler)) .route("/v2/task/{cl}", get(task_get_handler)) +} + +fn build_routes() -> Router { + Router::new() + .route("/retry-build", post(build_retry_handler)) .route("/v2/build-events/{task_id}", get(build_event_get_handler)) .route("/v2/targets/{task_id}", get(targets_get_handler)) .route("/v2/build-state/{build_id}", get(build_state_handler)) @@ -156,6 +74,20 @@ pub fn routers() -> Router { "/v2/latest_build_result/{task_id}", get(latest_build_result_handler), ) +} + +fn worker_routes() -> Router { + Router::new() + .route("/orion-clients-info", post(get_orion_clients_info)) + .route( + "/orion-client-status/{id}", + get(get_orion_client_status_by_id), + ) +} + +fn target_status_routes() -> Router { + Router::new() + .route("/targets/{target_id}/logs", get(target_logs_handler)) .route( "/v2/all-target-status/{task_id}", get(targets_status_handler), @@ -166,23 +98,17 @@ pub fn routers() -> Router { ) } -/// 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", + tag = "System", 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)) + api_v2_service::queue_stats(&state).await } /// Health check endpoint for Orion Server @@ -190,23 +116,14 @@ pub async fn queue_stats_handler(State(state): State) -> impl IntoResp #[utoipa::path( get, path = "/v2/health", + tag = "System", responses( (status = 200, description = "Service is healthy", body = serde_json::Value), (status = 503, description = "Service is unhealthy", body = serde_json::Value) ) )] pub async fn health_check_handler(State(state): State) -> impl IntoResponse { - // Simple health check: verify database connectivity - match tasks::Entity::find().limit(1).all(&state.conn).await { - Ok(_) => (StatusCode::OK, Json(json!({"status": "healthy"}))), - Err(e) => { - tracing::error!("Health check failed: {}", e); - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(json!({"status": "unhealthy", "error": "database connectivity check failed"})), - ) - } - } + api_v2_service::health_check(&state).await } /// Streams build output logs in real-time using Server-Sent Events (SSE) @@ -214,6 +131,7 @@ pub async fn health_check_handler(State(state): State) -> impl IntoRes #[utoipa::path( get, path = "/task-output/{id}", + tag = "Task", params( ("id" = String, Path, description = "Build ID for which to stream output logs") ), @@ -225,52 +143,8 @@ pub async fn health_check_handler(State(state): State) -> impl IntoRes pub async fn task_output_handler( State(state): State, Path(id): Path, -) -> Result>>, StatusCode> { - if !state.scheduler.active_builds.contains_key(&id) { - return Err(StatusCode::NOT_FOUND); - } - - // Use watch channel as stop signal for all streams - let (stop_tx, stop_rx) = watch::channel(true); - - let log_stop_rx = stop_rx.clone(); - // Log stream with termination condition - let log_stream = state - .log_service - .subscribe_for_build(id.clone()) - .map(|log_event| { - Ok::(Event::default().event("log").data(log_event.line)) - }) - .take_while(move |_| { - let stop_rx = log_stop_rx.clone(); - async move { *stop_rx.borrow() } - }); - - let heart_stop_rx = stop_rx.clone(); - // Heartbeat stream every 15 seconds with termination condition - let heartbeat_stream = IntervalStream::new(tokio::time::interval(Duration::from_secs(15))) - .map(|_| Ok::(Event::default().comment("heartbeat"))) - .take_while(move |_| { - let stop_rx_clone = heart_stop_rx.clone(); - async move { *stop_rx_clone.borrow() } - }); - - // Spawn a task to watch active_builds and send stop signal when build ends - let stop_tx_clone = stop_tx.clone(); - tokio::spawn(async move { - loop { - tokio::time::sleep(Duration::from_secs(1)).await; - if !state.scheduler.active_builds.contains_key(&id) { - let _ = stop_tx_clone.send(false); - break; - } - } - }); - - // Merge log and heartbeat streams - let stream = select(log_stream, heartbeat_stream); - - Ok(Sse::new(stream)) +) -> Result { + api_v2_service::task_output(&state, &id).await } /// Provides the ability to read historical task logs @@ -278,6 +152,7 @@ pub async fn task_output_handler( #[utoipa::path( get, path = "/task-history-output", + tag = "Task", params( ("task_id" = String, Query, description = "Task ID whose log to read"), ("build_id" = String, Query, description = "Build ID whose log to read"), @@ -296,46 +171,13 @@ pub async fn task_history_output_handler( State(state): State, Query(params): Query, ) -> Result, (StatusCode, Json)> { - // Determine which read method to call - let log_result = if matches!((params.start, params.end), (None, None)) { - state - .log_service - .read_full_log(¶ms.task_id, ¶ms.repo, ¶ms.build_id) - .await - } else { - // Unwrap start/end, default to 0 if needed - let start = params.start.unwrap_or(0); - let end = params.end.unwrap_or(usize::MAX); - state - .log_service - .read_log_range(¶ms.task_id, ¶ms.repo, ¶ms.build_id, start, end) - .await - }; - - // Handle result - let log_content = match log_result { - Ok(content) => content, - Err(e) => { - tracing::error!("read log failed: {:?}", e); - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to read log file".to_string(), - }), - )); - } - }; - - // Split the content into lines and count them - let lines: Vec = log_content.lines().map(str::to_string).collect(); - let len = lines.len(); - - Ok(Json(LogLinesResponse { data: lines, len })) + api_v2_service::task_history_output(&state, ¶ms).await } #[utoipa::path( get, path = "/targets/{target_id}/logs", + tag = "TargetStatus", params( ("target_id" = String, Path, description = "Target ID whose logs to read"), ("type" = String, Query, description = "full | segment"), @@ -358,160 +200,13 @@ pub async fn target_logs_handler( Path(target_id): Path, Query(params): Query, ) -> Result, (StatusCode, Json)> { - let target_uuid = match target_id.parse::() { - Ok(uuid) => uuid, - Err(_) => { - return Err(( - StatusCode::BAD_REQUEST, - Json(LogErrorResponse { - message: "Invalid target id".to_string(), - }), - )); - } - }; - - let target_model = match targets::Entity::find_by_id(target_uuid) - .one(&state.conn) - .await - { - Ok(Some(target)) => target, - Ok(None) => { - return Err(( - StatusCode::NOT_FOUND, - Json(LogErrorResponse { - message: "Target not found".to_string(), - }), - )); - } - Err(err) => { - tracing::error!("Failed to load target {}: {}", target_id, err); - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to read target".to_string(), - }), - )); - } - }; - - let build_model = if let Some(build_id) = params.build_id.as_ref() { - let build_uuid = build_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(LogErrorResponse { - message: "Invalid build id".to_string(), - }), - ) - })?; - - match builds::Entity::find_by_id(build_uuid) - .filter(builds::Column::TargetId.eq(target_uuid)) - .one(&state.conn) - .await - { - Ok(Some(build)) => build, - Ok(None) => { - return Err(( - StatusCode::NOT_FOUND, - Json(LogErrorResponse { - message: "Build not found for target".to_string(), - }), - )); - } - Err(err) => { - tracing::error!("Failed to load build {}: {}", build_uuid, err); - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to load build".to_string(), - }), - )); - } - } - } else { - match builds::Entity::find() - .filter(builds::Column::TargetId.eq(target_uuid)) - .order_by_desc(builds::Column::EndAt) - .order_by_desc(builds::Column::CreatedAt) - .one(&state.conn) - .await - { - Ok(Some(build)) => build, - Ok(None) => { - return Err(( - StatusCode::NOT_FOUND, - Json(LogErrorResponse { - message: "No builds for target".to_string(), - }), - )); - } - Err(err) => { - tracing::error!("Failed to load build for target {}: {}", target_uuid, err); - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to load build".to_string(), - }), - )); - } - } - }; - - let repo_segment = LogService::last_segment(&build_model.repo); - let log_result = if matches!(params.r#type, LogReadMode::Segment) { - let offset = params.offset.unwrap_or(0); - let limit = params.limit.unwrap_or(200); - state - .log_service - .read_log_range( - &target_model.task_id.to_string(), - &repo_segment, - &build_model.id.to_string(), - offset, - offset + limit, - ) - .await - } else { - state - .log_service - .read_full_log( - &target_model.task_id.to_string(), - &repo_segment, - &build_model.id.to_string(), - ) - .await - }; - - match log_result { - Ok(content) => { - let lines: Vec = content.lines().map(str::to_string).collect(); - let len = lines.len(); - Ok(Json(TargetLogLinesResponse { - data: lines, - len, - build_id: build_model.id.to_string(), - })) - } - Err(e) => { - tracing::error!( - "Failed to read logs for target {} build {}: {}", - target_uuid, - build_model.id, - e - ); - Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to read log".to_string(), - }), - )) - } - } + api_v2_service::target_logs(&state, &target_id, ¶ms).await } #[utoipa::path( post, path = "/v2/task", + tag = "Task", request_body = TaskBuildRequest, responses( (status = 200, description = "Task created", body = serde_json::Value), @@ -523,89 +218,14 @@ pub async fn task_handler_v2( State(state): State, Json(req): Json, ) -> impl IntoResponse { - let task_id = Uuid::now_v7(); - let _task_name = format!("CL-{}-{}", req.cl_link, task_id); - - // Create a new task in the database - if let Err(err) = - OrionTask::insert_task(task_id, &req.cl_link, &req.repo, &req.changes, &state.conn).await - { - tracing::error!("Failed to insert task into DB: {}", err); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "message": format!("Failed to insert task into database: {}", err) - })), - ) - .into_response(); - } - - // Check for available worker - let result: OrionBuildResult; - if state.scheduler.has_idle_workers() { - result = handle_immediate_task_dispatch_v2( - state.clone(), - task_id, - &req.repo, - &req.cl_link, - req.changes.clone(), - None, - ) - .await; - - ( - StatusCode::OK, - Json(OrionServerResponse { - task_id: task_id.to_string(), - results: vec![result], - }), - ) - .into_response() - } else { - tracing::info!( - "No idle workers available, attempting to enqueue task {}", - task_id - ); - match state - .scheduler - .enqueue_task_v2(task_id, &req.cl_link, req.repo, req.changes, 0) - .await - { - Ok(build_id) => { - tracing::info!("Build {}/{} queued successfully", task_id, build_id); - result = OrionBuildResult { - build_id: build_id.to_string(), - status: "queued".to_string(), - message: "Task queued for processing when workers become available".to_string(), - }; - ( - StatusCode::OK, - Json(OrionServerResponse { - task_id: task_id.to_string(), - results: vec![result], - }), - ) - .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() - } - } - } + api_v2_service::task_handler_v2(&state, req).await } /// Creates build tasks and returns the task ID and status (immediate or queued) #[utoipa::path( post, path = "/task", + tag = "Task", request_body = TaskBuildRequest, responses( (status = 200, description = "Task created", body = serde_json::Value), @@ -616,383 +236,13 @@ pub async fn task_handler( State(state): State, Json(req): Json, ) -> impl IntoResponse { - // create task id - let task_id = Uuid::now_v7(); - let task_name = format!("CL-{}-{}", req.cl_link, task_id); - - let mut results = Vec::new(); - - // Insert task into the database using the model's insert method - // TODO: replace with the new Task model - if let Err(err) = tasks::Model::insert_task( - task_id, - // TODO: replace with new Task, use cl_link as cl identifier - req.cl_id, - Some(task_name), - None, - chrono::Utc::now().into(), - &state.conn, - ) - .await - { - tracing::error!("Failed to insert task into DB: {}", err); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ - "message": format!("Failed to insert task into database: {}", err) - })), - ) - .into_response(); - } - - // Check if there are idle workers available - if state.scheduler.has_idle_workers() { - // Have idle workers, directly dispatch task (keep original logic) - let result: OrionBuildResult = handle_immediate_task_dispatch( - state.clone(), - task_id, - &req.repo, - &req.cl_link, - req.changes.clone(), - None, - ) - .await; - results.push(result); - } else { - // No idle workers, add task to queue - match state - .scheduler - .enqueue_task( - task_id, - &req.cl_link, - req.repo.clone(), - req.changes.clone(), - None, - 0, - ) - .await - { - Ok(build_id) => { - tracing::info!("Build {}/{} queued for later processing", task_id, build_id); - let result: OrionBuildResult = OrionBuildResult { - build_id: build_id.to_string(), - status: "queued".to_string(), - message: "Task queued for processing when workers become available".to_string(), - }; - results.push(result); - } - Err(e) => { - tracing::warn!("Failed to queue task: {}", e); - let result: OrionBuildResult = OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: format!("Unable to queue task: {}", e), - }; - results.push(result); - } - } - } - - ( - StatusCode::OK, - Json(OrionServerResponse { - task_id: task_id.to_string(), - results, - }), - ) - .into_response() -} - -async fn handle_immediate_task_dispatch_v2( - state: AppState, - task_id: Uuid, - repo: &str, - cl_link: &str, - changes: Vec>, - // TODO: if reused for retry here, use targets - _targets: Option>, -) -> OrionBuildResult { - // Create new build event ID - let build_event_id = Uuid::now_v7(); - - // Select an idle worker - let chosen_id = match state - .scheduler - .search_and_claim_worker(&build_event_id.to_string()) - { - Some(chosen_id) => { - tracing::info!("Selected idle worker {} for task {}", chosen_id, task_id); - chosen_id - } - None => { - tracing::error!("No idle workers available for task {}", task_id); - return OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: "No available workers at the moment".to_string(), - }; - } - }; - - // create and insert target path - let target_id = Uuid::now_v7(); - let target_path = - match BuildTarget::insert_default_target(target_id, task_id, &state.conn).await { - Ok(default_path) => default_path, - Err(err) => { - tracing::error!("Failed to prepare target for task {}: {}", task_id, err); - state.scheduler.release_worker(&chosen_id).await; - return OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: format!("Failed to prepare target for task {}", task_id), - }; - } - }; - - // Create new build event for initial try - let event = BuildEventPayload::new( - build_event_id, - task_id, - cl_link.to_string(), - repo.to_string(), - 0, - ); - - let start_at = chrono::Utc::now(); - let _start_at_tz = start_at.with_timezone(&FixedOffset::east_opt(0).unwrap()); - let build_info = BuildInfo { - event_payload: event.clone(), - changes: changes.clone(), - target_id, - target_path, - worker_id: chosen_id.clone(), - auto_retry_judger: AutoRetryJudger::new(), - started_at: start_at, - }; - - // Store build event using BuildInfo structure - match callisto::build_events::Model::insert_build( - build_event_id, - task_id, - repo.to_string(), - &state.conn, - ) - .await - { - Ok(_) => { - tracing::info!( - "Created build event record in DB with ID {} for task {}", - build_info.event_payload.build_event_id, - task_id, - ); - } - Err(e) => { - tracing::error!( - "Failed to insert build event into DB for task {} exit with error: {}", - task_id, - e - ); - state.scheduler.release_worker(&chosen_id).await; - return OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: format!( - "Failed to insert build event into database for task {}: {}", - task_id, e - ), - }; - } - } - - return activate_worker(&build_info, &state.scheduler).await; -} - -async fn activate_worker(build_info: &BuildInfo, scheduler: &TaskScheduler) -> OrionBuildResult { - // Create WS Message - let msg = WSMessage::TaskBuild { - build_id: build_info.event_payload.build_event_id.to_string(), - repo: build_info.event_payload.repo.clone(), - changes: build_info.changes.clone(), - cl_link: build_info.event_payload.cl_link.clone(), - }; - - // Send task to the selected worker - if let Some(worker) = scheduler.workers.get_mut(&build_info.worker_id) - && worker.sender.send(msg).is_ok() - { - scheduler.active_builds.insert( - build_info.event_payload.build_event_id.to_string(), - build_info.clone(), - ); - tracing::info!( - "Build {}/{} dispatched immediately to worker {}", - build_info.event_payload.task_id, - build_info.event_payload.build_event_id, - build_info.worker_id - ); - return OrionBuildResult { - build_id: build_info.event_payload.build_event_id.to_string(), - status: "dispatched".to_string(), - message: format!("Build dispatched to worker {}", build_info.worker_id), - }; - } - - tracing::error!( - "Failed to dispatch task {} to worker {}", - build_info.event_payload.task_id, - build_info.worker_id - ); - scheduler.release_worker(&build_info.worker_id).await; - OrionBuildResult { - build_id: build_info.event_payload.build_event_id.to_string(), - status: "error".to_string(), - message: "Failed to dispatch task to worker".to_string(), - } -} - -async fn handle_immediate_task_dispatch( - state: AppState, - task_id: Uuid, - repo: &str, - cl_link: &str, - changes: Vec>, - // TODO: if reused for retry here, use targets - _targets: Option>, -) -> OrionBuildResult { - // Find all idle workers - 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 OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: "No available workers at the moment".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 new build event - let build_id = Uuid::now_v7(); - - // TODO: use empty string temporary until target db is implemented - let target_path = String::new(); - let target_model = match state.scheduler.ensure_target(task_id, &target_path).await { - Ok(target) => target, - Err(err) => { - tracing::error!("Failed to prepare target {}: {}", target_path, err); - return OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: format!("Failed to prepare target {}", target_path), - }; - } - }; - - let start_at = chrono::Utc::now(); - let start_at_tz = start_at.with_timezone(&FixedOffset::east_opt(0).unwrap()); - - // Mark target as building - if let Err(e) = targets::update_state( - &state.conn, - target_model.id, - TargetState::Building, - Some(start_at_tz), - None, - None, - ) - .await - { - tracing::error!("Failed to update target state to Building: {}", e); - } - - let event = BuildEventPayload::new(build_id, task_id, cl_link.to_string(), repo.to_string(), 0); - - // Create build information structure - let build_info = BuildInfo { - event_payload: event.clone(), - changes: changes.clone(), - target_id: target_model.id, - target_path: target_model.target_path.clone(), - worker_id: chosen_id.clone(), - auto_retry_judger: AutoRetryJudger::new(), - started_at: start_at, - }; - - // Use the model's insert_build method for direct insertion - if let Err(err) = crate::model::builds::Model::insert_build( - build_id, - task_id, - target_model.id, - repo.to_string(), - &state.conn, - ) - .await - { - tracing::error!("Failed to insert builds into DB: {}", err); - return OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: format!("Failed to insert builds into database: {}", err), - }; - } - tracing::info!( - "Created build record in DB with ID {} for task {}", - build_id, - task_id - ); - - // Create WebSocket message for the worker (use first build's args) - let msg = WSMessage::TaskBuild { - build_id: build_id.to_string(), - repo: repo.to_string(), - changes: changes.clone(), - cl_link: cl_link.to_string(), - }; - - // Send task to the selected worker - if let Some(mut worker) = state.scheduler.workers.get_mut(&chosen_id) - && worker.sender.send(msg).is_ok() - { - worker.status = WorkerStatus::Busy { - build_id: build_id.to_string(), - phase: None, - }; - state - .scheduler - .active_builds - .insert(build_id.to_string(), build_info); - tracing::info!( - "Build {}/{} dispatched immediately to worker {}", - task_id, - build_id, - chosen_id - ); - return OrionBuildResult { - build_id: build_id.to_string(), - status: "dispatched".to_string(), - message: format!("Build dispatched to worker {}", chosen_id), - }; - } - - // If we reach here, sending failed - OrionBuildResult { - build_id: "".to_string(), - status: "error".to_string(), - message: "Failed to dispatch task to worker".to_string(), - } + api_v2_service::task_handler_v1(&state, req).await } #[utoipa::path( get, path = "/task-build-list/{id}", + tag = "Task", params( ("id" = String, Path, description = "Task ID to get build IDs for") ), @@ -1007,31 +257,7 @@ pub async fn task_build_list_handler( State(state): State, Path(id): Path, ) -> impl IntoResponse { - let db = &state.conn; - - let task_id = match id.parse::() { - Ok(uuid) => uuid, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "Invalid task ID format"})), - ) - .into_response(); - } - }; - - match tasks::Model::get_builds_by_task_id(task_id, db).await { - Some(build_ids) => { - let build_ids_str: Vec = - build_ids.into_iter().map(|uuid| uuid.to_string()).collect(); - (StatusCode::OK, Json(build_ids_str)).into_response() - } - None => ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Task not found"})), - ) - .into_response(), - } + api_v2_service::task_build_list(&state, &id).await } /// Handles WebSocket upgrade requests from workers @@ -1041,680 +267,13 @@ async fn ws_handler( ConnectInfo(addr): ConnectInfo, State(state): State, ) -> impl IntoResponse { - tracing::info!("{addr} connected. Waiting for registration..."); - ws.on_upgrade(move |socket| handle_socket(socket, addr, state)) -} - -/// Manages WebSocket connection lifecycle for worker communication -/// Handles message sending/receiving and connection cleanup -async fn handle_socket(socket: WebSocket, who: SocketAddr, state: AppState) { - let (tx, mut rx) = mpsc::unbounded_channel::(); - let mut worker_id: Option = None; - - let (mut sender, mut receiver) = socket.split(); - - // Task for sending messages to the worker - let send_task = tokio::spawn(async move { - while let Some(msg) = rx.recv().await { - let msg_str = serde_json::to_string(&msg).unwrap(); - if sender - .send(Message::Text(Utf8Bytes::from(msg_str))) - .await - .is_err() - { - tracing::warn!("Failed to send message to {who}, client disconnected."); - break; - } - } - }); - - let state_clone = state.clone(); - let tx_clone = tx.clone(); - - // Task for receiving messages from the worker - let recv_task = tokio::spawn(async move { - let mut worker_id_inner: Option = None; - while let Some(Ok(msg)) = receiver.next().await { - if process_message(msg, who, &state_clone, &mut worker_id_inner, &tx_clone) - .await - .is_break() - { - break; - } - } - worker_id_inner - }); - - tokio::select! { - _ = send_task => { }, - result = recv_task => { - if let Ok(final_worker_id) = result { - worker_id = final_worker_id; - } - }, - } - - // Cleanup worker connection when socket closes - if let Some(id) = &worker_id { - tracing::info!("Cleaning up for worker: {id} from {who}."); - state.scheduler.workers.remove(id); - } else { - tracing::info!("Cleaning up unregistered connection from {who}."); - } - - tracing::info!("Websocket context {who} destroyed"); -} - -/// Processes individual WebSocket messages from workers -/// Handles registration, heartbeats, build output, and completion messages -async fn process_message( - msg: Message, - who: SocketAddr, - state: &AppState, - worker_id: &mut Option, - tx: &UnboundedSender, -) -> ControlFlow<(), ()> { - match msg { - Message::Text(t) => { - let ws_msg: Result = serde_json::from_str(&t); - if let Err(e) = ws_msg { - tracing::warn!("Failed to parse message from {who}: {e}"); - return ControlFlow::Continue(()); - } - let ws_msg = ws_msg.unwrap(); - - // Handle worker registration (must be first message) - if worker_id.is_none() { - if let WSMessage::Register { - id, - hostname, - orion_version, - } = ws_msg - { - tracing::info!("Worker from {who} registered as: {id}"); - state.scheduler.workers.insert( - id.clone(), - WorkerInfo { - sender: tx.clone(), - status: WorkerStatus::Idle, - last_heartbeat: chrono::Utc::now(), - start_time: chrono::Utc::now(), - hostname, - orion_version, - }, - ); - *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." - ); - return ControlFlow::Break(()); - } - return ControlFlow::Continue(()); - } - - // Process messages from registered workers - let current_worker_id = worker_id.as_ref().unwrap(); - match ws_msg { - WSMessage::Register { .. } => { - tracing::warn!( - "Worker {current_worker_id} sent Register message again. Ignoring." - ); - } - WSMessage::Heartbeat => { - if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { - worker.last_heartbeat = chrono::Utc::now(); - tracing::debug!("Received heartbeat from {current_worker_id}"); - - // If the worker was previously in Error state, a successful heartbeat now restores it to Idle. - if let WorkerStatus::Error(_) = worker.status { - worker.status = WorkerStatus::Idle; - tracing::info!( - "Worker {current_worker_id} recovered from Error to Idle via heartbeat." - ); - } - } - } - WSMessage::TaskBuildOutput { build_id, output } => { - // Write build output to the associated log file - if let Some(build_info) = state.scheduler.active_builds.get(&build_id) { - let log_event = LogEvent { - task_id: build_info.event_payload.task_id.to_string(), - repo_name: LogService::last_segment( - &build_info.event_payload.repo.clone(), - ) - .to_string(), - build_id: build_id.clone(), - line: output.clone(), - is_end: false, - }; - // Publish the log event to the log stream - state.log_service.publish(log_event.clone()); - - // Debug output showing the published log - tracing::debug!( - "Published log for build_id {} (task: {}, repo: {}): {}", - build_id, - build_info.event_payload.task_id, - build_info.event_payload.repo, - output - ); - } else { - tracing::warn!("Received output for unknown task: {}", build_id); - } - - // Judge auto retry by output - if let Some(mut build_info) = state.scheduler.active_builds.get_mut(&build_id) { - build_info.auto_retry_judger.judge_by_output(&output); - } - } - WSMessage::TaskBuildCompleteV2 { - build_id, - success, - exit_code, - message, - } => { - // Handle build completion - tracing::info!( - "Build {build_id} completed by worker {current_worker_id} with exit code: {exit_code:?}" - ); - - // Get build information - let ( - mut auto_retry_judger, - mut retry_count, - repo, - changes, - cl_link, - task_id, - _target_id, - _target_path, - ) = if let Some(build_info) = state.scheduler.active_builds.get(&build_id) { - ( - build_info.auto_retry_judger.clone(), - build_info.event_payload.retry_count, - build_info.event_payload.repo.clone(), - build_info.changes.clone(), - build_info.event_payload.cl_link.clone(), - build_info.event_payload.task_id, - build_info.target_id, - build_info.target_path.clone(), - ) - } else { - tracing::error!("Not found build {build_id}"); - return ControlFlow::Continue(()); - }; - - // Judge auto retry by exit code - auto_retry_judger.judge_by_exit_code(exit_code.unwrap_or(0)); - - let can_auto_retry = auto_retry_judger.get_can_auto_retry(); - - if can_auto_retry && retry_count < RETRY_COUNT_MAX { - tracing::info!( - "Build {build_id} will retry, current retry count: {retry_count}" - ); - - // Add retry count - retry_count += 1; - - // Update build information - if let Some(mut build_info) = - state.scheduler.active_builds.get_mut(&build_id) - { - build_info.event_payload.retry_count = retry_count; - // New AutoRetryJudger - build_info.auto_retry_judger = AutoRetryJudger::new(); - } - - // Update database - let _ = builds::Entity::update_many() - .set(builds::ActiveModel { - retry_count: Set(retry_count), - ..Default::default() - }) - .filter(builds::Column::Id.eq(build_id.parse::().unwrap())) - .exec(&state.conn) - .await; - - // Send task to this worker - let msg = WSMessage::TaskBuild { - build_id: build_id.clone(), - repo: repo.clone(), - cl_link, - changes, - }; - if let Some(worker) = state.scheduler.workers.get_mut(current_worker_id) - && worker.sender.send(msg).is_ok() - { - tracing::info!( - "Retry build: {}, worker: {}", - build_id, - current_worker_id - ); - return ControlFlow::Continue(()); - } - } - - // Send final log event - let log_event = LogEvent { - task_id: task_id.to_string(), - repo_name: LogService::last_segment(&repo).to_string(), - build_id: build_id.to_string(), - line: String::from(""), - is_end: true, - }; - state.log_service.publish(log_event); - - // Remove from active - state.scheduler.active_builds.remove(&build_id); - - // TODO: update to jupiter's build_event's model - if BuildEvent::update_build_complete_result( - &build_id, - exit_code, - success, - &message, - &state.conn, - ) - .await - .is_err() - { - tracing::error!( - "Failed to update build complete result for build: {}", - build_id - ); - } - - // Update target state - let target_state = match (success, exit_code) { - (true, Some(0)) => TargetState::Completed, - (_, None) => TargetState::Interrupted, - _ => TargetState::Failed, - }; - let mut _error_summary = None; - if matches!(target_state, TargetState::Failed) { - let repo_segment = LogService::last_segment(&repo); - if let Ok(log_content) = state - .log_service - .read_full_log( - &task_id.to_string(), - &repo_segment, - &build_id.to_string(), - ) - .await - { - _error_summary = - find_caused_by_next_line_in_content(&log_content).await; - } - } - - if let Err(e) = - BuildTarget::update_build_targets(target_state, &build_id, &state.conn) - .await - { - tracing::error!( - "unable to update build targets for build {}: {}", - build_id, - e, - ) - } - - // Mark the worker as idle or error depending on whether the task succeeds. - if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { - worker.status = if success { - WorkerStatus::Idle - } else { - WorkerStatus::Error(message) - }; - } - - // Notify scheduler to process queued tasks - state.scheduler.notify_task_available(); - } - WSMessage::TaskBuildComplete { - build_id, - success, - exit_code, - message, - } => { - // Handle build completion - tracing::info!( - "Build {build_id} completed by worker {current_worker_id} with exit code: {exit_code:?}" - ); - - // Get build information - let ( - mut auto_retry_judger, - mut retry_count, - repo, - changes, - cl_link, - task_id, - target_id, - _target_path, - ) = if let Some(build_info) = state.scheduler.active_builds.get(&build_id) { - ( - build_info.auto_retry_judger.clone(), - build_info.event_payload.retry_count, - build_info.event_payload.repo.clone(), - build_info.changes.clone(), - build_info.event_payload.cl_link.clone(), - build_info.event_payload.task_id, - build_info.target_id, - build_info.target_path.clone(), - ) - } else { - tracing::error!("Not found build {build_id}"); - return ControlFlow::Continue(()); - }; - - // Judge auto retry by exit code - auto_retry_judger.judge_by_exit_code(exit_code.unwrap_or(0)); - - let can_auto_retry = auto_retry_judger.get_can_auto_retry(); - - if can_auto_retry && retry_count < RETRY_COUNT_MAX { - tracing::info!( - "Build {build_id} will retry, current retry count: {retry_count}" - ); - - // Add retry count - retry_count += 1; - - // Update build information - if let Some(mut build_info) = - state.scheduler.active_builds.get_mut(&build_id) - { - build_info.event_payload.retry_count = retry_count; - // New AutoRetryJudger - build_info.auto_retry_judger = AutoRetryJudger::new(); - } - - // Update database - let _ = builds::Entity::update_many() - .set(builds::ActiveModel { - retry_count: Set(retry_count), - ..Default::default() - }) - .filter(builds::Column::Id.eq(build_id.parse::().unwrap())) - .exec(&state.conn) - .await; - - // Send task to this worker - let msg = WSMessage::TaskBuild { - build_id: build_id.clone(), - repo: repo.clone(), - cl_link, - changes, - }; - if let Some(worker) = state.scheduler.workers.get_mut(current_worker_id) - && worker.sender.send(msg).is_ok() - { - tracing::info!( - "Retry build: {}, worker: {}", - build_id, - current_worker_id - ); - return ControlFlow::Continue(()); - } - } - - // Send final log event - let log_event = LogEvent { - task_id: task_id.to_string(), - repo_name: LogService::last_segment(&repo).to_string(), - build_id: build_id.to_string(), - line: String::from(""), - is_end: true, - }; - state.log_service.publish(log_event); - - // Remove from active - state.scheduler.active_builds.remove(&build_id); - - // Update database with final state - let end_at = Utc::now().with_timezone(&FixedOffset::east_opt(0).unwrap()); - // TODO: update to jupiter's build_event's model - let _ = builds::Entity::update_many() - .set(builds::ActiveModel { - exit_code: Set(exit_code), - end_at: Set(Some(end_at)), - retry_count: Set(retry_count), - ..Default::default() - }) - .filter(builds::Column::Id.eq(build_id.parse::().unwrap())) - .exec(&state.conn) - .await; - - // Update target state - let target_uuid = target_id.to_string().parse::().ok(); - let target_state = match (success, exit_code) { - (true, Some(0)) => TargetState::Completed, - (_, None) => TargetState::Interrupted, - _ => TargetState::Failed, - }; - let mut error_summary = None; - if matches!(target_state, TargetState::Failed) { - let repo_segment = LogService::last_segment(&repo); - if let Ok(log_content) = state - .log_service - .read_full_log( - &task_id.to_string(), - &repo_segment, - &build_id.to_string(), - ) - .await - { - error_summary = find_caused_by_next_line_in_content(&log_content).await; - } - } - if let Some(target_uuid) = target_uuid { - if let Err(e) = targets::update_state( - &state.conn, - target_uuid, - target_state, - None, - Some(end_at), - error_summary, - ) - .await - { - tracing::error!( - "Failed to update target state for build {}: {}", - build_id, - e - ); - } - } else { - tracing::warn!( - "Unable to parse target id {} for build {}", - target_id, - build_id - ); - } - - // Mark the worker as idle or error depending on whether the task succeeds. - if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { - worker.status = if success { - WorkerStatus::Idle - } else { - WorkerStatus::Error(message) - }; - } - - // Notify scheduler to process queued tasks - state.scheduler.notify_task_available(); - } - WSMessage::TaskPhaseUpdate { build_id, phase } => { - tracing::info!( - "Task phase updated by orion worker {current_worker_id} with: {phase:?}" - ); - if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) { - if let WorkerStatus::Busy { build_id: id, .. } = &worker.status { - if &build_id == id { - worker.status = WorkerStatus::Busy { - build_id, - phase: Some(phase), - }; - } else { - tracing::warn!( - "Ignoring TaskPhaseUpdate for worker {current_worker_id}: \ - task_id mismatch (expected {build_id}, got {id})" - ); - } - } else { - tracing::warn!( - "Ignoring TaskPhaseUpdate for worker {current_worker_id}: \ - worker not in Busy state (current status: {:?})", - worker.status - ); - } - } else { - tracing::warn!( - "Ignoring TaskPhaseUpdate: unknown worker {current_worker_id}" - ); - } - } - WSMessage::TargetBuildStatusBatch { events } => { - tracing::info!( - "Target build status updated by orion worker {current_worker_id}, {} items", - events.len() - ); - - for update in events { - state.target_status_cache.insert_event(update).await; - } - } - _ => {} - } - } - Message::Close(_) => { - tracing::info!("Client {who} sent close message."); - if let Some(id) = worker_id.take() - && let Some(mut worker) = state.scheduler.workers.get_mut(&id) - { - worker.status = WorkerStatus::Lost; - tracing::info!("Worker {id} marked as Lost due to connection close"); - } - - return ControlFlow::Break(()); - } - _ => {} - } - ControlFlow::Continue(()) -} - -/// Data transfer object for build information in API responses -#[derive(Debug, Serialize, ToSchema, Clone)] -pub struct BuildDTO { - pub id: String, - pub task_id: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub target_id: Option, - pub exit_code: Option, - pub start_at: String, - pub end_at: Option, - pub repo: String, - pub target: String, - pub args: Option, - pub output_file: String, - pub created_at: String, - pub retry_count: i32, - pub status: TaskStatusEnum, - pub cause_by: Option, -} - -impl BuildDTO { - /// Converts a database model to a DTO for API responses. - /// `target` is optional; empty string means target path missing (for compat). - pub fn from_model( - model: builds::Model, - target: Option<&targets::Model>, - status: TaskStatusEnum, - ) -> Self { - let target_path = target.map(|t| t.target_path.clone()).unwrap_or_default(); - Self { - id: model.id.to_string(), - task_id: model.task_id.to_string(), - target_id: target.map(|t| t.id.to_string()), - exit_code: model.exit_code, - start_at: model.start_at.with_timezone(&Utc).to_rfc3339(), - end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), - repo: model.repo, - target: target_path, - args: model.args.map(|v| json!(v)), - output_file: model.output_file, - created_at: model.created_at.with_timezone(&Utc).to_rfc3339(), - retry_count: model.retry_count, - status, - cause_by: None, - } - } - - /// Determine build status based on database fields and active builds - pub fn determine_status(model: &builds::Model, is_active: bool) -> TaskStatusEnum { - if is_active { - TaskStatusEnum::Building - } else if model.end_at.is_none() { - // Not in active_builds and end_at is None => still queued (pending) - TaskStatusEnum::Pending - } else if model.exit_code.is_none() { - TaskStatusEnum::Interrupted - } else if model.exit_code == Some(0) { - TaskStatusEnum::Completed - } else { - TaskStatusEnum::Failed - } - } -} - -pub type TargetDTO = targets::TargetWithBuilds; - -/// Task information including current status -#[derive(Debug, Serialize, ToSchema)] -pub struct TaskInfoDTO { - pub task_id: String, - pub cl_id: i64, - pub task_name: Option, - pub template: Option, - pub created_at: String, - pub build_list: Vec, - pub targets: Vec, -} - -/// Target summary counts for a task. -#[derive(Debug, Serialize, ToSchema)] -pub struct TargetSummaryDTO { - pub task_id: String, - pub pending: u64, - pub building: u64, - pub completed: u64, - pub failed: u64, - pub interrupted: u64, - pub uninitialized: u64, -} - -impl TaskInfoDTO { - fn from_model(model: tasks::Model, build_list: Vec, targets: Vec) -> Self { - Self { - task_id: model.id.to_string(), - cl_id: model.cl_id, - task_name: model.task_name, - template: model.template, - created_at: model.created_at.with_timezone(&Utc).to_rfc3339(), - build_list, - targets, - } - } + ws_service::ws_handler(ws, ConnectInfo(addr), State(state)).await } #[utoipa::path( get, path = "/tasks/{cl}", + tag = "Task", params( ("cl" = i64, Path, description = "CL number to filter tasks by") ), @@ -1728,108 +287,13 @@ pub async fn tasks_handler( State(state): State, Path(cl): Path, ) -> Result>, (StatusCode, Json)> { - let db = &state.conn; - let active_builds = state.scheduler.active_builds.clone(); - - match tasks::Entity::find() - .filter(tasks::Column::ClId.eq(cl)) - .all(db) - .await - { - Ok(task_models) => { - let mut tasks: Vec = Vec::new(); - - for m in task_models { - tasks.push(assemble_task_info(m, &state, &active_builds).await); - } - - Ok(Json(tasks)) - } - Err(e) => { - tracing::error!("Failed to fetch tasks: {e}"); - Err(( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Failed to fetch tasks"})), - )) - } - } -} - -async fn assemble_task_info( - task: tasks::Model, - state: &AppState, - active_builds: &Arc>, -) -> TaskInfoDTO { - let target_models = targets::Entity::find() - .filter(targets::Column::TaskId.eq(task.id)) - .all(&state.conn) - .await - .unwrap_or_else(|_| vec![]); - - let build_models = builds::Entity::find() - .filter(builds::Column::TaskId.eq(task.id)) - .all(&state.conn) - .await - .unwrap_or_else(|_| vec![]); - - let target_map: HashMap = - target_models.iter().cloned().map(|t| (t.id, t)).collect(); - - let mut build_list: Vec = Vec::new(); - let mut target_build_map: HashMap> = HashMap::new(); - for build_model in build_models { - let build_id_str = build_model.id.to_string(); - let is_active = active_builds.contains_key(&build_id_str); - let status = BuildDTO::determine_status(&build_model, is_active); - let mut dto = BuildDTO::from_model( - build_model.clone(), - target_map.get(&build_model.target_id), - status.clone(), - ); - - // Prefer persisted error summary; avoid reading full logs in the task summary path. - if matches!(status, TaskStatusEnum::Failed) - && let Some(t) = target_map.get(&build_model.target_id) - && let Some(summary) = &t.error_summary - { - dto.cause_by = Some(summary.clone()); - } - - target_build_map - .entry(build_model.target_id) - .or_default() - .push(dto.clone()); - build_list.push(dto); - } - - let mut target_list: Vec = Vec::new(); - for target in target_models { - let builds = target_build_map.remove(&target.id).unwrap_or_default(); - target_list.push(TargetWithBuilds::from_model(target, builds)); - } - - TaskInfoDTO::from_model(task, build_list, target_list) -} - -async fn find_caused_by_next_line_in_content(content: &str) -> Option { - let mut last_was_caused = false; - - for line in content.lines() { - if last_was_caused { - return Some(line.to_string()); - } - - if line.trim() == "Caused by:" { - last_was_caused = true; - } - } - - None + api_v2_service::tasks_by_cl(&state, cl).await } #[utoipa::path( get, path = "/tasks/{task_id}/targets", + tag = "Task", params( ("task_id" = String, Path, description = "Task ID to query targets for") ), @@ -1844,37 +308,13 @@ pub async fn task_targets_handler( State(state): State, Path(task_id): Path, ) -> Result, (StatusCode, Json)> { - let task_uuid = task_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "Invalid task ID"})), - ) - })?; - - let task_model = tasks::Entity::find_by_id(task_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Failed to fetch task"})), - ) - })? - .ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Task not found"})), - ) - })?; - - let info = assemble_task_info(task_model, &state, &state.scheduler.active_builds).await; - Ok(Json(info)) + api_v2_service::task_targets(&state, &task_id).await } #[utoipa::path( get, path = "/tasks/{task_id}/targets/summary", + tag = "Task", params( ("task_id" = String, Path, description = "Task ID to query target summary for") ), @@ -1888,95 +328,14 @@ pub async fn task_targets_summary_handler( State(state): State, Path(task_id): Path, ) -> Result, (StatusCode, Json)> { - let task_uuid = task_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ "message": "Invalid task ID" })), - ) - })?; - - let targets = targets::Entity::find() - .filter(targets::Column::TaskId.eq(task_uuid)) - .all(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch target summary: {e}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({ "message": "Failed to fetch target summary" })), - ) - })?; - - let mut summary = TargetSummaryDTO { - task_id, - pending: 0, - building: 0, - completed: 0, - failed: 0, - interrupted: 0, - uninitialized: 0, - }; - - for target in targets { - match target.state { - TargetState::Pending => summary.pending += 1, - TargetState::Building => summary.building += 1, - TargetState::Completed => summary.completed += 1, - TargetState::Failed => summary.failed += 1, - TargetState::Interrupted => summary.interrupted += 1, - TargetState::Uninitialized => summary.uninitialized += 1, - } - } - - Ok(Json(summary)) -} - -// Orion client information -#[derive(Debug, Serialize, ToSchema)] -pub struct OrionClientInfo { - pub client_id: String, - pub hostname: String, - pub orion_version: String, - #[schema(value_type = String, format = "date-time")] - pub start_time: DateTimeUtc, - #[schema(value_type = String, format = "date-time")] - pub last_heartbeat: DateTimeUtc, -} - -impl OrionClientInfo { - fn from_worker(client_id: impl Into, worker: &WorkerInfo) -> Self { - Self { - client_id: client_id.into(), - hostname: worker.hostname.clone(), - orion_version: worker.orion_version.clone(), - start_time: worker.start_time, - last_heartbeat: worker.last_heartbeat, - } - } -} - -// Orion client status -#[derive(Debug, Clone, Deserialize, Serialize, ToSchema, PartialEq, Eq)] -pub enum CoreWorkerStatus { - Idle, - Busy, - Error, - Lost, -} - -/// Additional query parameters for querying Orion clients. -/// When no extra conditions are required, this struct can be left empty. -#[derive(Debug, Deserialize, ToSchema, Clone)] -pub struct OrionClientQuery { - pub hostname: Option, - pub status: Option, - pub phase: Option, // Only in Busy status + api_v2_service::task_targets_summary(&state, &task_id).await } /// Endpoint to retrieve paginated Orion client information. #[utoipa::path( post, path = "/orion-clients-info", + tag = "Worker", request_body = PageParams, responses( (status = 200, description = "Paged Orion client information", body = CommonPage) @@ -1986,93 +345,14 @@ pub async fn get_orion_clients_info( State(state): State, Json(params): Json>, ) -> Result>, (StatusCode, Json)> { - let pagination = params.pagination; - let query = params.additional.clone(); - - let page = pagination.page.max(1); - // per_page must be in 1..=100 - let per_page = pagination.per_page.clamp(1u64, 100); - let offset = (page - 1) * per_page; - - let mut total: u64 = 0; - let mut items: Vec = Vec::with_capacity(per_page as usize); - - for entry in state.scheduler.workers.iter() { - let matches = query - .hostname - .as_ref() - .is_none_or(|h| entry.value().hostname.contains(h)) - && query - .status - .as_ref() - .is_none_or(|s| entry.value().status.status_type() == *s) - && query.phase.as_ref().is_none_or(|p| { - matches!( - entry.value().status, - WorkerStatus::Busy { phase: Some(ref x), .. } if *x == *p - ) - }); - - if matches { - total += 1; - - if total > offset && items.len() < per_page as usize { - items.push(OrionClientInfo::from_worker( - entry.key().clone(), - entry.value(), - )); - } - } - } - - Ok(Json(CommonPage { total, items })) -} - -// Orion client status -#[derive(Debug, Serialize, ToSchema)] -pub struct OrionClientStatus { - /// Core (Idle / Busy / Error / Lost) - pub core_status: CoreWorkerStatus, - /// Only when building - pub phase: Option, - /// Only when error - pub error_message: Option, -} - -impl OrionClientStatus { - pub fn from_worker_status(worker: &WorkerInfo) -> Self { - match &worker.status { - WorkerStatus::Idle => Self { - core_status: CoreWorkerStatus::Idle, - phase: None, - error_message: None, - }, - - WorkerStatus::Busy { phase, .. } => Self { - core_status: CoreWorkerStatus::Busy, - phase: phase.clone(), - error_message: None, - }, - - WorkerStatus::Error(msg) => Self { - core_status: CoreWorkerStatus::Error, - phase: None, - error_message: Some(msg.clone()), - }, - - WorkerStatus::Lost => Self { - core_status: CoreWorkerStatus::Lost, - phase: None, - error_message: None, - }, - } - } + api_v2_service::get_orion_clients_info(&state, params).await } /// Retrieve the current status of a specific Orion client by its ID. #[utoipa::path( get, path = "/orion-client-status/{id}", + tag = "Worker", params( ("id" = String, description = "Orion client Id") ), @@ -2085,24 +365,14 @@ pub async fn get_orion_client_status_by_id( State(state): State, Path(id): Path, ) -> Result, (StatusCode, Json)> { - let worker = state.scheduler.workers.get(&id).ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "message": "Orion client not found" - })), - ) - })?; - - let status = OrionClientStatus::from_worker_status(&worker); - - Ok(Json(status)) + api_v2_service::get_orion_client_status_by_id(&state, &id).await } /// Retry the build #[utoipa::path( post, path = "/retry-build", + tag = "Build", request_body = RetryBuildRequest, responses( (status = 200, description = "Retry created", body = serde_json::Value), @@ -2117,246 +387,13 @@ pub async fn build_retry_handler( State(state): State, Json(req): Json, ) -> impl IntoResponse { - let db = &state.conn; - - let build_id = match req.build_id.parse::() { - Ok(uuid) => uuid, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "Invalid build ID format"})), - ) - .into_response(); - } - }; - - if state.scheduler.active_builds.contains_key(&req.build_id) { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "The build already exists"})), - ) - .into_response(); - } - - let build = match builds::Entity::find_by_id(build_id).one(db).await { - Ok(o) => match o { - Some(build) => build, - None => { - return ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Build not found"})), - ) - .into_response(); - } - }, - Err(_) => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Database find failed"})), - ) - .into_response(); - } - }; - - let retry_count = build.retry_count + 1; - let target_model = match targets::Entity::find_by_id(build.target_id).one(db).await { - Ok(Some(target)) => target, - Ok(None) => { - tracing::error!("Target not found for build {}", build.id); - return ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Target not found for build"})), - ) - .into_response(); - } - Err(err) => { - tracing::error!("Failed to load target for build {}: {}", build.id, err); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Database find failed"})), - ) - .into_response(); - } - }; - - let idle_workers = state.scheduler.get_idle_workers(); - if idle_workers.is_empty() { - // Generate a new build id for queued retry to avoid PK conflict with existing build. - let new_build_id = Uuid::now_v7(); - // No idle workers, add task to queue - match state - .scheduler - .enqueue_task_with_build_id( - new_build_id, - build.task_id, - &req.cl_link, - build.repo.clone(), - req.changes.clone(), - target_model.target_path.clone(), - retry_count, - ) - .await - { - Ok(()) => { - tracing::info!("Build {} queued for later processing", build.id); - ( - StatusCode::OK, - Json(serde_json::json!({"message": "Build queued for later processing"})), - ) - .into_response() - } - Err(e) => { - tracing::warn!("Failed to queue retry build: {}", e); - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({"message": "No available workers at the moment"})), - ) - .into_response() - } - } - } else if immediate_work( - &state, - build_id, - &idle_workers, - &build, - &target_model, - &req, - retry_count, - ) - .await - { - ( - StatusCode::OK, - Json(serde_json::json!({ - "message": "Build retry dispatched immediately to worker" - })), - ) - .into_response() - } else { - tracing::warn!( - "Failed to dispatch build {} retry to worker; worker missing or send failed", - build.id, - ); - ( - StatusCode::BAD_GATEWAY, - Json(serde_json::json!({ - "message": "Failed to dispatch build retry to worker" - })), - ) - .into_response() - } -} - -// TODO: replace with new build -async fn immediate_work( - state: &AppState, - build_id: Uuid, - idle_workers: &[String], - build: &builds::Model, - target: &targets::Model, - req: &RetryBuildRequest, - retry_count: i32, -) -> bool { - // Randomly select an idle worker - let chosen_index = { - let mut rng = rand::rng(); - rng.random_range(0..idle_workers.len()) - }; - let chosen_id = idle_workers[chosen_index].clone(); - - let start_at = chrono::Utc::now(); - // Create build information - - let event = BuildEventPayload::new( - build.id, - build.task_id, - req.cl_link.clone(), - build.repo.clone(), - retry_count, - ); - let build_info = BuildInfo { - event_payload: event, - target_id: target.id, - target_path: target.target_path.clone(), - changes: req.changes.clone(), - worker_id: chosen_id.clone(), - auto_retry_judger: AutoRetryJudger::new(), - started_at: start_at, - }; - - // Send build to worker - let msg = WSMessage::TaskBuild { - build_id: build.id.to_string(), - repo: build.repo.to_string(), - changes: req.changes.clone(), - cl_link: req.cl_link.to_string(), - }; - - if let Some(mut worker) = state.scheduler.workers.get_mut(&chosen_id) - && worker.sender.send(msg).is_ok() - { - worker.status = WorkerStatus::Busy { - build_id: build_id.to_string(), - phase: None, - }; - if let Err(e) = targets::update_state( - &state.conn, - target.id, - TargetState::Building, - Some(start_at.with_timezone(&FixedOffset::east_opt(0).unwrap())), - None, - None, - ) - .await - { - tracing::error!("Failed to update target state to Building: {}", e); - } - // Insert active build - state - .scheduler - .active_builds - .insert(build.id.to_string(), build_info); - tracing::info!( - "Build {} retry dispatched immediately to worker {}", - build.id, - chosen_id - ); - true - } else { - tracing::warn!( - "Failed to dispatch build {} retry to worker {}; worker missing or send failed", - build.id, - chosen_id - ); - false - } -} - -#[derive(ToSchema, Serialize)] -pub struct MessageResponse { - pub message: String, -} - -#[derive(ToSchema, Serialize)] -pub struct BuildTargetDTO { - pub id: String, - pub task_id: String, - pub path: String, - pub latest_state: String, -} - -#[derive(ToSchema, Serialize)] -pub enum BuildEventState { - #[allow(dead_code)] - Pending, - Running, - Success, - Failure, + api_v2_service::build_retry(&state, req).await } #[utoipa::path( post, path = "/v2/task-retry/{id}", + tag = "Task", params(("id" = String, description = "Task ID to retry task")), responses( (status = 200, description = "Task queued for retry", body = MessageResponse), @@ -2368,49 +405,13 @@ pub async fn task_retry_handler( State(state): State, Path(id): Path, ) -> Result, (StatusCode, Json)> { - let task_uuid = id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(MessageResponse { - message: "Invalid task ID format".to_string(), - }), - ) - })?; - - // Verify task exists in orion_tasks table - let task = callisto::orion_tasks::Entity::find_by_id(task_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch task {}: {}", id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })? - .ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(MessageResponse { - message: "Task not found".to_string(), - }), - ) - })?; - - // For now, just return success message - // TODO: Implement actual retry logic - create new build events, update scheduler, etc. - tracing::info!("Task retry requested for task {} (CL: {})", id, task.cl); - - Ok(Json(MessageResponse { - message: format!("Task {} queued for retry", id), - })) + api_v2_service::task_retry(&state, &id).await } #[utoipa::path( get, path = "/v2/task/{cl}", + tag = "Task", params(("cl" = String, Path, description = "Change List")), responses( (status = 200, description = "Get task successfully", body = OrionTaskDTO), @@ -2423,34 +424,13 @@ pub async fn task_get_handler( State(state): State, Path(cl): Path, ) -> Result, (StatusCode, Json)> { - let tasks: Vec = callisto::orion_tasks::Entity::find() - .filter(callisto::orion_tasks::Column::Cl.eq(&cl)) - .all(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch tasks by CL {}: {}", &cl, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Database error"})), - ) - })?; - - match tasks.len() { - 0 => Err(( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Not found task"})), - )), - 1 => Ok(Json(OrionTaskDTO::from(&tasks[0]))), - _ => Err(( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "Multiple tasks"})), - )), - } + api_v2_service::task_get(&state, &cl).await } #[utoipa::path( get, path = "/v2/build-events/{task_id}", + tag = "Build", params(("task_id" = String, Path, description = "Task ID")), responses( (status = 200, description = "Get build events successfully", body = Vec), @@ -2463,56 +443,13 @@ pub async fn build_event_get_handler( State(state): State, Path(task_id): Path, ) -> Result>, (StatusCode, Json)> { - let task_uuid = task_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"message": "Invalid task ID"})), - ) - })?; - - // First, verify the task exists - let task_exists = callisto::orion_tasks::Entity::find_by_id(task_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to verify task existence {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Database error"})), - ) - })? - .is_some(); - - if !task_exists { - return Err(( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"message": "Task not found"})), - )); - } - - let build_events = callisto::build_events::Entity::find() - .filter(callisto::build_events::Column::TaskId.eq(task_uuid)) - .all(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch build events for task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"message": "Database error"})), - ) - })?; - - let dtos: Vec = build_events - .into_iter() - .map(|m| BuildEventDTO::from(&m)) - .collect(); - - Ok(Json(dtos)) + api_v2_service::build_event_get(&state, &task_id).await } #[utoipa::path( get, path = "/v2/targets/{task_id}", + tag = "Build", params(("task_id" = String, Path, description = "Task ID")), responses( (status = 200, description = "Get targets successfully", body = Vec), @@ -2524,70 +461,14 @@ pub async fn targets_get_handler( State(state): State, Path(task_id): Path, ) -> Result>, (StatusCode, Json)> { - let task_uuid = task_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(MessageResponse { - message: "Invalid task ID".to_string(), - }), - ) - })?; - - // First, verify the task exists - let task_exists = callisto::orion_tasks::Entity::find_by_id(task_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to verify task existence {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })? - .is_some(); - - if !task_exists { - return Err(( - StatusCode::NOT_FOUND, - Json(MessageResponse { - message: "Task not found".to_string(), - }), - )); - } - - let build_targets = callisto::build_targets::Entity::find() - .filter(callisto::build_targets::Column::TaskId.eq(task_uuid)) - .all(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch build targets for task {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })?; - - let dtos: Vec = build_targets - .into_iter() - .map(|build_target| BuildTargetDTO { - id: build_target.id.to_string(), - task_id: build_target.task_id.to_string(), - path: build_target.path, - latest_state: build_target.latest_state, - }) - .collect(); - - Ok(Json(dtos)) + api_v2_service::targets_get(&state, &task_id).await } /// Get complete log for a specific build event #[utoipa::path( get, path = "/v2/builds/{build_id}/logs", + tag = "Build", params(("build_id" = String, Path, description = "Build event ID")), responses( (status = 200, description = "Complete log content", body = api_model::buck2::types::LogLinesResponse), @@ -2600,87 +481,14 @@ pub async fn build_logs_handler( State(state): State, Path(build_id): Path, ) -> Result, (StatusCode, Json)> { - let build_uuid = build_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(LogErrorResponse { - message: "Invalid build ID".to_string(), - }), - ) - })?; - - // Get build event to extract log storage path - let build_event = callisto::build_events::Entity::find_by_id(build_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch build event {}: {}", build_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Database error".to_string(), - }), - ) - })? - .ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(LogErrorResponse { - message: "Build event not found".to_string(), - }), - ) - })?; - - // Get the associated Orion task to extract repository name - let orion_task = callisto::orion_tasks::Entity::find_by_id(build_event.task_id) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch Orion task {}: {}", build_event.task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Database error".to_string(), - }), - ) - })? - .ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(LogErrorResponse { - message: "Task not found".to_string(), - }), - ) - })?; - - let task_id = build_event.task_id.to_string(); - let repo_name = &orion_task.repo_name; - - // Read complete log using existing LogService - let log_content = state - .log_service - .read_full_log(&task_id, repo_name, &build_id) - .await - .map_err(|e| { - tracing::error!("Failed to read log for build {}: {}", build_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(LogErrorResponse { - message: "Failed to read log".to_string(), - }), - ) - })?; - - // Format response - let lines: Vec = log_content.lines().map(str::to_string).collect(); - let len = lines.len(); - Ok(Json(LogLinesResponse { data: lines, len })) + api_v2_service::build_logs(&state, &build_id).await } /// Get build state by build ID #[utoipa::path( get, path = "/v2/build-state/{build_id}", + tag = "Build", params(("build_id" = String, Path, description = "Build ID")), responses( (status = 200, description = "Get build state successfully", body = BuildEventState), @@ -2691,57 +499,14 @@ pub async fn build_state_handler( State(state): State, Path(build_id): Path, ) -> Result, (StatusCode, Json)> { - let build_uuid = build_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(MessageResponse { - message: "Invalid build ID".to_string(), - }), - ) - })?; - - // First check if the build is currently active (running or pending) - if state.scheduler.active_builds.contains_key(&build_id) { - return Ok(Json(BuildEventState::Running)); - } - - // If not active, check the build_events table - let build_event = callisto::build_events::Entity::find_by_id(build_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to fetch build event {}: {}", build_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })? - .ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json(MessageResponse { - message: "Build not found".to_string(), - }), - ) - })?; - - // Determine state based on build event fields - let state_enum = match (build_event.end_at, build_event.exit_code) { - (None, _) => BuildEventState::Running, // Still running but not in active_builds (shouldn't happen) - (Some(_), Some(0)) => BuildEventState::Success, - (Some(_), Some(_)) => BuildEventState::Failure, - (Some(_), None) => BuildEventState::Failure, // Interrupted case - }; - - Ok(Json(state_enum)) + api_v2_service::build_state(&state, &build_id).await } /// Get latest build result by task ID #[utoipa::path( get, path = "/v2/latest_build_result/{task_id}", + tag = "Build", params(("task_id" = String, Path, description = "Task ID")), responses( (status = 200, description = "Get latest build result successfully", body = BuildEventState), @@ -2752,240 +517,14 @@ pub async fn latest_build_result_handler( State(state): State, Path(task_id): Path, ) -> Result, (StatusCode, Json)> { - let task_uuid = task_id.parse::().map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json(MessageResponse { - message: "Invalid task ID".to_string(), - }), - ) - })?; - - // Verify task exists - let task_exists = callisto::orion_tasks::Entity::find_by_id(task_uuid) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!("Failed to verify task existence {}: {}", task_id, e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })? - .is_some(); - - if !task_exists { - return Err(( - StatusCode::NOT_FOUND, - Json(MessageResponse { - message: "Task not found".to_string(), - }), - )); - } - - // Find latest build event for this task - let latest_build_event = callisto::build_events::Entity::find() - .filter(callisto::build_events::Column::TaskId.eq(task_uuid)) - .order_by_desc(callisto::build_events::Column::StartAt) - .one(&state.conn) - .await - .map_err(|e| { - tracing::error!( - "Failed to fetch latest build event for task {}: {}", - task_id, - e - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(MessageResponse { - message: "Database error".to_string(), - }), - ) - })?; - - let build_event = match latest_build_event { - Some(event) => event, - None => { - return Err(( - StatusCode::NOT_FOUND, - Json(MessageResponse { - message: "No build events found for this task".to_string(), - }), - )); - } - }; - - // Determine state based on build event fields - let state_enum = match (build_event.end_at, build_event.exit_code) { - (None, _) => BuildEventState::Running, - (Some(_), Some(0)) => BuildEventState::Success, - (Some(_), Some(_)) => BuildEventState::Failure, - (Some(_), None) => BuildEventState::Failure, // Interrupted case - }; - - Ok(Json(state_enum)) -} - -#[derive(Hash, Eq, PartialEq, Clone)] -struct ActionKey { - package: String, - name: String, - configuration: String, - category: String, - identifier: String, - action: String, -} - -#[derive(Clone)] -pub struct TargetStatusCache { - /// task_id -> (ActionKey -> ActiveModel) - inner: Arc>>>, -} - -impl TargetStatusCache { - pub fn new() -> Self { - Self { - inner: Arc::new(RwLock::new(HashMap::new())), - } - } - - /// Insert a single WebSocket event into the cache - pub async fn insert_event(&self, event: WSTargetBuildStatusEvent) { - let task_id = match Uuid::parse_str(&event.context.task_id) { - Ok(id) => id, - Err(_) => { - tracing::error!("Invalid task_id: {}", event.context.task_id); - return; - } - }; - - let status = OrionTargetStatusEnum::from_ws_status(&event.target.new_status); - - let key = ActionKey { - package: event.target.configured_target_package.clone(), - name: event.target.configured_target_name.clone(), - configuration: event.target.configured_target_configuration.clone(), - category: event.target.category.clone(), - identifier: event.target.identifier.clone(), - action: event.target.action.clone(), - }; - - let active_model = TargetBuildStatusService::new_active_model( - Uuid::new_v4(), // id is not auto-incremented - task_id, - event.target.configured_target_package, - event.target.configured_target_name, - event.target.configured_target_configuration, - event.target.category, - event.target.identifier, - event.target.action, - status, - ); - - let mut guard = self.inner.write().await; - let task_map = guard.entry(task_id).or_default(); - - task_map.insert(key, active_model); - } - - /// Flush all cached entries - pub async fn flush_all(&self) -> Vec { - let mut guard = self.inner.write().await; - let mut result = Vec::new(); - - for (_, action_map) in guard.drain() { - result.extend(action_map.into_values()); - } - - result - } - - /// Flush cached entries for a specific task - pub async fn _flush_task(&self, task_id: Uuid) -> Vec { - let mut guard = self.inner.write().await; - - guard - .remove(&task_id) - .map(|m| m.into_values().collect()) - .unwrap_or_default() - } - - pub async fn auto_flush_loop( - self, - conn: DatabaseConnection, - mut shutdown: tokio::sync::watch::Receiver, - ) { - let mut ticker = tokio::time::interval(std::time::Duration::from_millis(500)); - - loop { - tokio::select! { - _ = ticker.tick() => { - let models = self.flush_all().await; - - if models.is_empty() { - continue; - } - - if let Err(e) = TargetBuildStatusService::upsert_batch(&conn, models).await - { - tracing::error!("Auto flush failed: {:?}", e); - } - } - - _ = shutdown.changed() => { - tracing::info!("TargetStatusCache auto flush shutting down"); - - let models = self.flush_all().await; - - if !models.is_empty() { - let _ = TargetBuildStatusService::upsert_batch(&conn, models).await; - } - - break; - } - } - } - } -} - -impl Default for TargetStatusCache { - fn default() -> Self { - Self::new() - } -} - -pub trait FromWsStatus { - fn from_ws_status(status: &str) -> Self; - fn as_str(&self) -> &str; -} - -impl FromWsStatus for OrionTargetStatusEnum { - fn from_ws_status(status: &str) -> Self { - match status.trim().to_ascii_lowercase().as_str() { - "pending" => Self::Pending, - "running" => Self::Running, - "success" | "succeeded" => Self::Success, - "failed" => Self::Failed, - _ => Self::Pending, - } - } - - fn as_str(&self) -> &str { - match self { - Self::Pending => "PENDING", - Self::Running => "RUNNING", - Self::Success => "SUCCESS", - Self::Failed => "FAILED", - } - } + api_v2_service::latest_build_result(&state, &task_id).await } /// Get target status with task_id #[utoipa::path( get, path = "/v2/all-target-status/{task_id}", + tag = "TargetStatus", params( ("task_id" = String, Path, description = "Task ID whose target belong"), ), @@ -2998,51 +537,16 @@ pub async fn targets_status_handler( State(state): State, Path(task_id): Path, ) -> Result>, (StatusCode, String)> { - use sea_orm::prelude::Uuid; - - let task_uuid = match Uuid::parse_str(&task_id) { - Ok(uuid) => uuid, - Err(_) => return Err((StatusCode::BAD_REQUEST, "Invalid task_id".to_string())), - }; - - let targets = match TargetBuildStatusService::fetch_by_task_id(&state.conn, task_uuid).await { - Ok(list) => list, - Err(e) => { - tracing::error!( - "Failed to fetch target status for task_id {}: {:?}", - task_id, - e - ); - return Err((StatusCode::INTERNAL_SERVER_ERROR, "Database error".into())); - } - }; - - if targets.is_empty() { - return Err((StatusCode::NOT_FOUND, "No target status found".to_string())); - } - - let response: Vec = targets - .into_iter() - .map(|t| TargetStatusResponse { - id: t.id.to_string(), - task_id: t.task_id.to_string(), - package: t.target_package, - name: t.target_name, - configuration: t.target_configuration, - category: t.category, - identifier: t.identifier, - action: t.action, - status: t.status.as_str().to_owned(), - }) - .collect(); - - Ok(Json(response)) + crate::service::target_status_cache_service::targets_status_by_task_id(&state.conn, &task_id) + .await + .map(Json) } /// Get target status with target id #[utoipa::path( get, path = "/v2/target-status/{target_id}", + tag = "TargetStatus", params( ("target_id" = String, Path, description = "target_id ID"), ), @@ -3055,43 +559,9 @@ pub async fn single_target_status_handle( State(state): State, Path(target_id): Path, ) -> Result, (StatusCode, String)> { - // 解析 target_id 为 UUID - let target_uuid = match Uuid::parse_str(&target_id) { - Ok(id) => id, - Err(e) => { - return Err(( - StatusCode::BAD_REQUEST, - format!("Invalid target_id '{}': {}", target_id, e), - )); - } - }; - - // 查询数据库 - let target = match TargetBuildStatusService::find_by_id(&state.conn, target_uuid).await { - Ok(Some(t)) => t, - Ok(None) => return Err((StatusCode::NOT_FOUND, "Target not found".to_string())), - Err(e) => { - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Database query failed: {}", e), - )); - } - }; - - // 构建响应 - let response = TargetStatusResponse { - id: target.id.to_string(), - task_id: target.task_id.to_string(), - package: target.target_package, - name: target.target_name, - configuration: target.target_configuration, - category: target.category, - identifier: target.identifier, - action: target.action, - status: target.status.as_str().to_string(), - }; - - Ok(Json(response)) + crate::service::target_status_cache_service::target_status_by_id(&state.conn, &target_id) + .await + .map(Json) } #[cfg(test)] diff --git a/orion-server/src/api_doc.rs b/orion-server/src/api_doc.rs new file mode 100644 index 000000000..d41e8e5bd --- /dev/null +++ b/orion-server/src/api_doc.rs @@ -0,0 +1,73 @@ +use api_model::buck2::types::TaskPhase; +use utoipa::OpenApi; + +use crate::api; + +/// OpenAPI documentation configuration. +#[derive(OpenApi)] +#[openapi( + paths( + // Task domain + api::task_handler, + api::task_build_list_handler, + api::task_output_handler, + api::task_history_output_handler, + api::tasks_handler, + api::task_targets_handler, + api::task_targets_summary_handler, + api::task_retry_handler, + api::task_get_handler, + // Build domain + api::build_retry_handler, + api::build_event_get_handler, + api::targets_get_handler, + api::build_state_handler, + api::build_logs_handler, + api::latest_build_result_handler, + // Worker domain + api::get_orion_clients_info, + api::get_orion_client_status_by_id, + // Target status domain + api::target_logs_handler, + api::targets_status_handler, + api::single_target_status_handle, + // System domain + api::health_check_handler + ), + components( + schemas( + crate::scheduler::BuildRequest, + api_model::buck2::types::LogLinesResponse, + api_model::buck2::types::TargetLogLinesResponse, + api_model::buck2::types::LogErrorResponse, + crate::model::task_status::TaskStatusEnum, + crate::model::dto::BuildDTO, + crate::model::dto::TargetDTO, + crate::model::dto::TargetSummaryDTO, + api_model::buck2::types::TargetLogQuery, + api_model::buck2::types::LogReadMode, + api_model::buck2::types::TaskHistoryQuery, + crate::model::dto::TaskInfoDTO, + crate::model::dto::OrionClientInfo, + crate::model::dto::OrionClientStatus, + crate::model::dto::CoreWorkerStatus, + crate::model::dto::OrionClientQuery, + crate::entity::targets::TargetState, + TaskPhase, + crate::model::dto::MessageResponse, + crate::model::dto::BuildEventDTO, + crate::model::dto::OrionTaskDTO, + crate::model::dto::BuildTargetDTO, + crate::model::dto::BuildEventState, + api_model::buck2::types::TargetStatusResponse, + ) + ), + tags( + (name = "Task", description = "Task lifecycle and task query endpoints"), + (name = "Build", description = "Build dispatch/retry/state/log endpoints"), + (name = "Worker", description = "Orion worker status and listing endpoints"), + (name = "TargetStatus", description = "Target log and target status endpoints"), + (name = "System", description = "Health and system-level endpoints") + ) +)] +pub struct ApiDoc; diff --git a/orion-server/src/app_state.rs b/orion-server/src/app_state.rs new file mode 100644 index 000000000..65ac21884 --- /dev/null +++ b/orion-server/src/app_state.rs @@ -0,0 +1,55 @@ +use std::sync::Arc; + +use dashmap::DashMap; +use sea_orm::DatabaseConnection; +use tokio::sync::watch; + +use crate::{ + log::log_service::LogService, scheduler::TaskScheduler, + service::target_status_cache_service::TargetStatusCache, +}; + +/// Shared application state containing worker connections, database, and active builds. +#[derive(Clone)] +pub struct AppState { + pub scheduler: TaskScheduler, + pub conn: DatabaseConnection, + pub log_service: LogService, + pub target_status_cache: TargetStatusCache, + shutdown_tx: watch::Sender, +} + +impl AppState { + pub fn new( + conn: DatabaseConnection, + queue_config: Option, + log_service: LogService, + ) -> 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); + let target_status_cache = TargetStatusCache::new(); + let (shutdown_tx, _) = watch::channel(false); + + Self { + scheduler, + conn, + log_service, + target_status_cache, + shutdown_tx, + } + } + + pub fn start_background_tasks(&self) { + let conn = self.conn.clone(); + let cache = self.target_status_cache.clone(); + let shutdown_rx = self.shutdown_tx.subscribe(); + tokio::spawn(async move { + cache.auto_flush_loop(conn, shutdown_rx).await; + }); + } + + pub async fn start_queue_manager(self) { + self.scheduler.start_queue_manager().await; + } +} diff --git a/orion-server/src/buck2.rs b/orion-server/src/buck2.rs index 2d61d29b3..e0f1af230 100644 --- a/orion-server/src/buck2.rs +++ b/orion-server/src/buck2.rs @@ -1,8 +1,4 @@ -use std::{fs, path::Path, process::Command, time::Duration}; - use once_cell::sync::OnceCell; -use tokio_retry::{Retry, strategy::ExponentialBackoff}; -use uuid::Uuid; /// Mono base URL for file/blob API, set at startup from config (or default). static MONO_BASE_URL: OnceCell = OnceCell::new(); @@ -11,174 +7,3 @@ static MONO_BASE_URL: OnceCell = OnceCell::new(); pub fn set_mono_base_url(url: String) { let _ = MONO_BASE_URL.set(url); } - -/// 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 (from config, or default). -fn base_url() -> String { - MONO_BASE_URL - .get() - .cloned() - .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 -#[allow(dead_code)] -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 -#[allow(dead_code)] -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 - //TODO: Cleanup should happen in a finally block or using RAII to ensure temporary files are removed even if buck2 command fails. Consider using a guard or defer pattern. TempDirGuard , for example. - if Path::new(&tmp_dir_path).exists() { - fs::remove_dir_all(&tmp_dir_path)?; - } - - Ok(last_line) -} - -#[cfg(test)] -mod tests { - use std::{fs, path::Path}; - - use super::*; - - #[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 '{test_dir}' does not exist. Skipping test."); - } - } - - /// 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(()) - } -} diff --git a/orion-server/src/model/builds.rs b/orion-server/src/entity/builds.rs similarity index 52% rename from orion-server/src/model/builds.rs rename to orion-server/src/entity/builds.rs index 826f253da..390d17dae 100644 --- a/orion-server/src/model/builds.rs +++ b/orion-server/src/entity/builds.rs @@ -1,7 +1,6 @@ //! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.15 -use chrono::Utc; -use sea_orm::{ActiveValue::Set, ConnectionTrait, entity::prelude::*}; +use sea_orm::entity::prelude::*; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -56,47 +55,3 @@ impl Related for Entity { } impl ActiveModelBehavior for ActiveModel {} - -impl Model { - /// Create a new build ActiveModel for database insertion - pub fn create_build( - build_id: Uuid, - task_id: Uuid, - target_id: Uuid, - repo: String, - args: Option, - ) -> ActiveModel { - let now = Utc::now().into(); - let repo_leaf = repo - .trim_end_matches('/') - .rsplit('/') - .next() - .unwrap_or(&repo) - .to_string(); - ActiveModel { - id: Set(build_id), - task_id: Set(task_id), - target_id: Set(target_id), - exit_code: Set(None), - start_at: Set(now), - end_at: Set(None), - repo: Set(repo), - args: Set(args), - output_file: Set(format!("{}/{}/{}.log", task_id, repo_leaf, build_id)), - created_at: Set(now), - retry_count: Set(0), - } - } - - /// Insert a single build directly into the database - pub async fn insert_build( - build_id: Uuid, - task_id: Uuid, - target_id: Uuid, - repo: String, - db: &impl ConnectionTrait, - ) -> Result { - let build_model = Self::create_build(build_id, task_id, target_id, repo, None); - build_model.insert(db).await - } -} diff --git a/orion-server/src/entity/mod.rs b/orion-server/src/entity/mod.rs new file mode 100644 index 000000000..340c8e6c0 --- /dev/null +++ b/orion-server/src/entity/mod.rs @@ -0,0 +1,3 @@ +pub mod builds; +pub mod targets; +pub mod tasks; diff --git a/orion-server/src/entity/targets.rs b/orion-server/src/entity/targets.rs new file mode 100644 index 000000000..37aac8787 --- /dev/null +++ b/orion-server/src/entity/targets.rs @@ -0,0 +1,122 @@ +use std::fmt::Display; + +use chrono::Utc; +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; +use uuid::Uuid; + +#[derive( + Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Serialize, Deserialize, Copy, ToSchema, +)] +#[sea_orm(rs_type = "String", db_type = "Text")] +pub enum TargetState { + #[sea_orm(string_value = "Pending")] + Pending, + #[sea_orm(string_value = "Building")] + Building, + #[sea_orm(string_value = "Completed")] + Completed, + #[sea_orm(string_value = "Failed")] + Failed, + #[sea_orm(string_value = "Interrupted")] + Interrupted, + #[sea_orm(string_value = "Uninitialized")] + Uninitialized, +} + +impl From for TargetState { + fn from(s: String) -> Self { + match s.as_str() { + "Pending" => TargetState::Pending, + "Building" => TargetState::Building, + "Completed" => TargetState::Completed, + "Failed" => TargetState::Failed, + "Interrupted" => TargetState::Interrupted, + "Uninitialized" => TargetState::Uninitialized, + _ => TargetState::Pending, // Default to Pending for unknown states + } + } +} + +impl Display for TargetState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + TargetState::Pending => "Pending", + TargetState::Building => "Building", + TargetState::Completed => "Completed", + TargetState::Failed => "Failed", + TargetState::Interrupted => "Interrupted", + TargetState::Uninitialized => "Uninitialized", + }; + write!(f, "{}", s) + } +} + +/// Target DTO with a generic builds payload. +#[derive(Debug, Serialize, ToSchema, Clone)] +pub struct TargetWithBuilds { + pub id: String, + pub target_path: String, + pub state: TargetState, + pub start_at: Option, + pub end_at: Option, + pub error_summary: Option, + pub builds: Vec, +} + +impl TargetWithBuilds { + pub fn from_model(model: Model, builds: Vec) -> Self { + Self { + id: model.id.to_string(), + target_path: model.target_path, + state: model.state, + start_at: model.start_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), + end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), + error_summary: model.error_summary, + builds, + } + } +} + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "targets")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: Uuid, + pub task_id: Uuid, + pub target_path: String, + pub state: TargetState, + pub start_at: Option, + pub end_at: Option, + pub error_summary: Option, + pub created_at: DateTimeWithTimeZone, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::tasks::Entity", + from = "Column::TaskId", + to = "super::tasks::Column::Id", + on_update = "NoAction", + on_delete = "Cascade" + )] + Tasks, + #[sea_orm(has_many = "super::builds::Entity")] + Builds, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Tasks.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Builds.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/orion-server/src/entity/tasks.rs b/orion-server/src/entity/tasks.rs new file mode 100644 index 000000000..7d4feb4f1 --- /dev/null +++ b/orion-server/src/entity/tasks.rs @@ -0,0 +1,39 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.15 + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "tasks")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: Uuid, + pub cl_id: i64, + pub task_name: Option, + #[sea_orm(column_type = "JsonBinary", nullable)] + pub template: Option, + pub created_at: DateTimeWithTimeZone, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm(has_many = "super::builds::Entity")] + Builds, + #[sea_orm(has_many = "super::targets::Entity")] + Targets, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Builds.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Targets.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/orion-server/src/lib.rs b/orion-server/src/lib.rs index 8dd15b4d8..5b298098c 100644 --- a/orion-server/src/lib.rs +++ b/orion-server/src/lib.rs @@ -1,9 +1,12 @@ pub mod api; +pub mod api_doc; +pub mod app_state; pub mod auto_retry; pub mod buck2; +pub mod entity; pub mod log; pub mod model; -pub mod orion_common; +pub mod repository; pub mod scheduler; pub mod server; pub mod service; diff --git a/orion-server/src/main.rs b/orion-server/src/main.rs index 658885cd7..38a337893 100644 --- a/orion-server/src/main.rs +++ b/orion-server/src/main.rs @@ -1,13 +1,3 @@ -mod api; -mod auto_retry; -mod buck2; -mod log; -mod model; -mod orion_common; -mod scheduler; -mod server; -mod service; - /// Orion Build Server /// A distributed build system that manages build tasks and worker nodes #[tokio::main] @@ -20,5 +10,5 @@ async fn main() { // Load environment variables from .env file (optional) dotenvy::dotenv().ok(); - server::start_server().await; + orion_server::server::start_server().await; } diff --git a/orion-server/src/model/build_events.rs b/orion-server/src/model/build_events.rs deleted file mode 100644 index 9208b5fc0..000000000 --- a/orion-server/src/model/build_events.rs +++ /dev/null @@ -1,54 +0,0 @@ -use chrono::Utc; -use sea_orm::{ActiveValue::Set, ColumnTrait, ConnectionTrait, DbErr, EntityTrait, QueryFilter}; -use serde::Serialize; -use utoipa::ToSchema; -use uuid::Uuid; - -#[derive(ToSchema, Serialize)] -pub struct BuildEventDTO { - pub id: String, - pub task_id: String, - pub retry_count: i32, - pub exit_code: Option, - pub log: Option, - pub log_output_file: String, - pub start_at: String, - pub end_at: Option, -} - -impl From<&callisto::build_events::Model> for BuildEventDTO { - fn from(model: &callisto::build_events::Model) -> Self { - Self { - id: model.id.to_string(), - task_id: model.task_id.to_string(), - retry_count: model.retry_count, - exit_code: model.exit_code, - log: model.log.clone(), - log_output_file: model.log_output_file.clone(), - start_at: model.start_at.to_string(), - end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_string()), - } - } -} - -pub struct BuildEvent; - -impl BuildEvent { - pub async fn update_build_complete_result( - build_id: &str, - exit_code: Option, - _success: bool, - _message: &str, - db_connection: &impl ConnectionTrait, - ) -> Result<(), DbErr> { - callisto::build_events::Entity::update_many() - .filter(callisto::build_events::Column::Id.eq(build_id.parse::().unwrap())) - .set(callisto::build_events::ActiveModel { - exit_code: Set(exit_code), - ..Default::default() - }) - .exec(db_connection) - .await?; - Ok(()) - } -} diff --git a/orion-server/src/model/dto.rs b/orion-server/src/model/dto.rs new file mode 100644 index 000000000..905b5fcb9 --- /dev/null +++ b/orion-server/src/model/dto.rs @@ -0,0 +1,226 @@ +use api_model::buck2::types::TaskPhase; +use chrono::Utc; +use sea_orm::prelude::DateTimeUtc; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use utoipa::ToSchema; + +use crate::{ + entity::{builds, targets}, + model::task_status::TaskStatusEnum, + scheduler::{WorkerInfo, WorkerStatus}, +}; + +/// Data transfer object for build information in API responses. +#[derive(Debug, Serialize, ToSchema, Clone)] +pub struct BuildDTO { + pub id: String, + pub task_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub target_id: Option, + pub exit_code: Option, + pub start_at: String, + pub end_at: Option, + pub repo: String, + pub target: String, + pub args: Option, + pub output_file: String, + pub created_at: String, + pub retry_count: i32, + pub status: TaskStatusEnum, + pub cause_by: Option, +} + +impl BuildDTO { + pub fn from_model( + model: builds::Model, + target: Option<&targets::Model>, + status: TaskStatusEnum, + ) -> Self { + let target_path = target.map(|t| t.target_path.clone()).unwrap_or_default(); + Self { + id: model.id.to_string(), + task_id: model.task_id.to_string(), + target_id: target.map(|t| t.id.to_string()), + exit_code: model.exit_code, + start_at: model.start_at.with_timezone(&Utc).to_rfc3339(), + end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), + repo: model.repo, + target: target_path, + args: model.args.map(|v| json!(v)), + output_file: model.output_file, + created_at: model.created_at.with_timezone(&Utc).to_rfc3339(), + retry_count: model.retry_count, + status, + cause_by: None, + } + } + + pub fn determine_status(model: &builds::Model, is_active: bool) -> TaskStatusEnum { + if is_active { + TaskStatusEnum::Building + } else if model.end_at.is_none() { + TaskStatusEnum::Pending + } else if model.exit_code.is_none() { + TaskStatusEnum::Interrupted + } else if model.exit_code == Some(0) { + TaskStatusEnum::Completed + } else { + TaskStatusEnum::Failed + } + } +} + +pub type TargetDTO = targets::TargetWithBuilds; + +#[derive(Debug, Serialize, ToSchema)] +pub struct TaskInfoDTO { + pub task_id: String, + pub cl_id: i64, + pub task_name: Option, + pub template: Option, + pub created_at: String, + pub build_list: Vec, + pub targets: Vec, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct TargetSummaryDTO { + pub task_id: String, + pub pending: u64, + pub building: u64, + pub completed: u64, + pub failed: u64, + pub interrupted: u64, + pub uninitialized: u64, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct OrionClientInfo { + pub client_id: String, + pub hostname: String, + pub orion_version: String, + #[schema(value_type = String, format = "date-time")] + pub start_time: DateTimeUtc, + #[schema(value_type = String, format = "date-time")] + pub last_heartbeat: DateTimeUtc, +} + +#[derive(Debug, Clone, Deserialize, Serialize, ToSchema, PartialEq, Eq)] +pub enum CoreWorkerStatus { + Idle, + Busy, + Error, + Lost, +} + +#[derive(Debug, Deserialize, ToSchema, Clone)] +pub struct OrionClientQuery { + pub hostname: Option, + pub status: Option, + pub phase: Option, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct OrionClientStatus { + pub core_status: CoreWorkerStatus, + pub phase: Option, + pub error_message: Option, +} + +impl OrionClientStatus { + pub fn from_worker_status(worker: &WorkerInfo) -> Self { + match &worker.status { + WorkerStatus::Idle => Self { + core_status: CoreWorkerStatus::Idle, + phase: None, + error_message: None, + }, + WorkerStatus::Busy { phase, .. } => Self { + core_status: CoreWorkerStatus::Busy, + phase: phase.clone(), + error_message: None, + }, + WorkerStatus::Error(msg) => Self { + core_status: CoreWorkerStatus::Error, + phase: None, + error_message: Some(msg.clone()), + }, + WorkerStatus::Lost => Self { + core_status: CoreWorkerStatus::Lost, + phase: None, + error_message: None, + }, + } + } +} + +#[derive(ToSchema, Serialize)] +pub struct MessageResponse { + pub message: String, +} + +#[derive(ToSchema, Serialize)] +pub struct BuildTargetDTO { + pub id: String, + pub task_id: String, + pub path: String, + pub latest_state: String, +} + +#[derive(ToSchema, Serialize)] +pub struct BuildEventDTO { + pub id: String, + pub task_id: String, + pub retry_count: i32, + pub exit_code: Option, + pub log: Option, + pub log_output_file: String, + pub start_at: String, + pub end_at: Option, +} + +impl From<&callisto::build_events::Model> for BuildEventDTO { + fn from(model: &callisto::build_events::Model) -> Self { + Self { + id: model.id.to_string(), + task_id: model.task_id.to_string(), + retry_count: model.retry_count, + exit_code: model.exit_code, + log: model.log.clone(), + log_output_file: model.log_output_file.clone(), + start_at: model.start_at.to_string(), + end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_string()), + } + } +} + +#[derive(ToSchema, Serialize, Deserialize)] +pub struct OrionTaskDTO { + pub id: String, + pub changes: Value, + pub repo_name: String, + pub cl: String, + pub created_at: String, +} + +impl From<&callisto::orion_tasks::Model> for OrionTaskDTO { + fn from(model: &callisto::orion_tasks::Model) -> Self { + Self { + id: model.id.to_string(), + changes: model.changes.clone(), + repo_name: model.repo_name.clone(), + cl: model.cl.clone(), + created_at: model.created_at.with_timezone(&Utc).to_string(), + } + } +} + +#[derive(ToSchema, Serialize)] +pub enum BuildEventState { + #[allow(dead_code)] + Pending, + Running, + Success, + Failure, +} diff --git a/orion-server/src/model/mod.rs b/orion-server/src/model/mod.rs index cffc6fd33..423bdc032 100644 --- a/orion-server/src/model/mod.rs +++ b/orion-server/src/model/mod.rs @@ -1,6 +1,2 @@ -pub mod build_events; -pub mod build_targets; -pub mod builds; -pub mod orion_tasks; -pub mod targets; -pub mod tasks; +pub mod dto; +pub mod task_status; diff --git a/orion-server/src/model/targets.rs b/orion-server/src/model/targets.rs deleted file mode 100644 index 72e37fd23..000000000 --- a/orion-server/src/model/targets.rs +++ /dev/null @@ -1,231 +0,0 @@ -use std::fmt::Display; - -use chrono::Utc; -use sea_orm::{ - ActiveModelTrait, ActiveValue::Set, ConnectionTrait, DbErr, RuntimeErr, entity::prelude::*, - sqlx, -}; -use serde::{Deserialize, Serialize}; -use utoipa::ToSchema; -use uuid::Uuid; - -#[derive( - Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Serialize, Deserialize, Copy, ToSchema, -)] -#[sea_orm(rs_type = "String", db_type = "Text")] -pub enum TargetState { - #[sea_orm(string_value = "Pending")] - Pending, - #[sea_orm(string_value = "Building")] - Building, - #[sea_orm(string_value = "Completed")] - Completed, - #[sea_orm(string_value = "Failed")] - Failed, - #[sea_orm(string_value = "Interrupted")] - Interrupted, - #[sea_orm(string_value = "Uninitialized")] - Uninitialized, -} - -impl From for TargetState { - fn from(s: String) -> Self { - match s.as_str() { - "Pending" => TargetState::Pending, - "Building" => TargetState::Building, - "Completed" => TargetState::Completed, - "Failed" => TargetState::Failed, - "Interrupted" => TargetState::Interrupted, - "Uninitialized" => TargetState::Uninitialized, - _ => TargetState::Pending, // Default to Pending for unknown states - } - } -} - -impl Display for TargetState { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let s = match self { - TargetState::Pending => "Pending", - TargetState::Building => "Building", - TargetState::Completed => "Completed", - TargetState::Failed => "Failed", - TargetState::Interrupted => "Interrupted", - TargetState::Uninitialized => "Uninitialized", - }; - write!(f, "{}", s) - } -} - -/// Target DTO with a generic builds payload. -#[derive(Debug, Serialize, ToSchema, Clone)] -pub struct TargetWithBuilds { - pub id: String, - pub target_path: String, - pub state: TargetState, - pub start_at: Option, - pub end_at: Option, - pub error_summary: Option, - pub builds: Vec, -} - -impl TargetWithBuilds { - pub fn from_model(model: Model, builds: Vec) -> Self { - Self { - id: model.id.to_string(), - target_path: model.target_path, - state: model.state, - start_at: model.start_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), - end_at: model.end_at.map(|dt| dt.with_timezone(&Utc).to_rfc3339()), - error_summary: model.error_summary, - builds, - } - } -} - -#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] -#[sea_orm(table_name = "targets")] -pub struct Model { - #[sea_orm(primary_key, auto_increment = false)] - pub id: Uuid, - pub task_id: Uuid, - pub target_path: String, - pub state: TargetState, - pub start_at: Option, - pub end_at: Option, - pub error_summary: Option, - pub created_at: DateTimeWithTimeZone, -} - -#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] -pub enum Relation { - #[sea_orm( - belongs_to = "super::tasks::Entity", - from = "Column::TaskId", - to = "super::tasks::Column::Id", - on_update = "NoAction", - on_delete = "Cascade" - )] - Tasks, - #[sea_orm(has_many = "super::builds::Entity")] - Builds, -} - -impl Related for Entity { - fn to() -> RelationDef { - Relation::Tasks.def() - } -} - -impl Related for Entity { - fn to() -> RelationDef { - Relation::Builds.def() - } -} - -impl ActiveModelBehavior for ActiveModel {} - -impl Model { - /// Create a new target ActiveModel ready for insertion - pub fn create_target(task_id: Uuid, target_path: impl Into) -> ActiveModel { - let now = Utc::now().into(); - ActiveModel { - id: Set(Uuid::now_v7()), - task_id: Set(task_id), - target_path: Set(target_path.into()), - state: Set(TargetState::Pending), - start_at: Set(None), - end_at: Set(None), - error_summary: Set(None), - created_at: Set(now), - } - } - - /// Insert a target directly into the database - pub async fn insert_target( - task_id: Uuid, - target_path: impl Into, - db: &impl ConnectionTrait, - ) -> Result { - let target_model = Self::create_target(task_id, target_path); - target_model.insert(db).await - } -} - -impl Entity { - /// Find an existing target for the same task and path, or create a new one. - pub async fn find_or_create( - db: &impl ConnectionTrait, - task_id: Uuid, - target_path: impl Into + Clone, - ) -> Result { - let target_path_owned: String = target_path.clone().into(); - if let Some(target) = Entity::find() - .filter(Column::TaskId.eq(task_id)) - .filter(Column::TargetPath.eq(target_path_owned.clone())) - .one(db) - .await? - { - return Ok(target); - } - - // Handle potential race: unique(task_id, target_path) enforced in migration - match Model::insert_target(task_id, target_path_owned.clone(), db).await { - Ok(model) => Ok(model), - Err(DbErr::Exec(RuntimeErr::SqlxError(sqlx::Error::Database(db_err)))) => { - if let Some(code) = db_err.code() - && code == "23505" - { - return Entity::find() - .filter(Column::TaskId.eq(task_id)) - .filter(Column::TargetPath.eq(target_path_owned)) - .one(db) - .await? - .ok_or_else(|| DbErr::RecordNotFound("target".into())); - } - Err(DbErr::Exec(RuntimeErr::SqlxError(sqlx::Error::Database( - db_err, - )))) - } - Err(e) => Err(e), - } - } -} - -/// Update target state and optional timing/error metadata. -pub async fn update_state( - db: &impl ConnectionTrait, - target_id: Uuid, - state: TargetState, - start_at: Option, - end_at: Option, - error_summary: Option, -) -> Result<(), DbErr> { - let mut active = ActiveModel { - id: Set(target_id), - state: Set(state), - ..Default::default() - }; - - // When transitioning back to Building/Pending, clear stale end_at/error_summary - match state { - TargetState::Building | TargetState::Pending => { - active.end_at = Set(None); - active.error_summary = Set(None); - } - _ => {} - } - - if let Some(start_at) = start_at { - active.start_at = Set(Some(start_at)); - } - - if let Some(end_at) = end_at { - active.end_at = Set(Some(end_at)); - } - - if let Some(summary) = error_summary { - active.error_summary = Set(Some(summary)); - } - - active.update(db).await.map(|_| ()) -} diff --git a/orion-server/src/model/task_status.rs b/orion-server/src/model/task_status.rs new file mode 100644 index 000000000..b01a29fe2 --- /dev/null +++ b/orion-server/src/model/task_status.rs @@ -0,0 +1,15 @@ +use serde::Serialize; +use utoipa::ToSchema; + +/// Enumeration of possible task statuses +#[derive(Debug, Serialize, Default, ToSchema, Clone)] +pub enum TaskStatusEnum { + /// Task is queued and waiting to be assigned to a worker + Pending, + Building, + Interrupted, // Task was interrupted, exit code is None + Failed, + Completed, + #[default] + NotFound, +} diff --git a/orion-server/src/model/tasks.rs b/orion-server/src/model/tasks.rs deleted file mode 100644 index f595de9f1..000000000 --- a/orion-server/src/model/tasks.rs +++ /dev/null @@ -1,94 +0,0 @@ -//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.15 - -use sea_orm::{ActiveValue::Set, ConnectionTrait, DbErr, entity::prelude::*}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] -#[sea_orm(table_name = "tasks")] -pub struct Model { - #[sea_orm(primary_key, auto_increment = false)] - pub id: Uuid, - pub cl_id: i64, - pub task_name: Option, - #[sea_orm(column_type = "JsonBinary", nullable)] - pub template: Option, - pub created_at: DateTimeWithTimeZone, -} - -#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] -pub enum Relation { - #[sea_orm(has_many = "super::builds::Entity")] - Builds, - #[sea_orm(has_many = "super::targets::Entity")] - Targets, -} - -impl Related for Entity { - fn to() -> RelationDef { - Relation::Builds.def() - } -} - -impl Related for Entity { - fn to() -> RelationDef { - Relation::Targets.def() - } -} - -impl ActiveModelBehavior for ActiveModel {} - -impl Model { - /// Retrieves build IDs associated with a task ID - pub async fn get_builds_by_task_id( - task_id: Uuid, - db: &DatabaseConnection, - ) -> Option> { - // Query the builds table for all builds with the given task_id - match super::builds::Entity::find() - .filter(super::builds::Column::TaskId.eq(task_id)) - .all(db) - .await - { - Ok(builds) => { - // Extract the IDs from the builds - let build_ids: Vec = builds.into_iter().map(|build| build.id).collect(); - Some(build_ids) - } - Err(e) => { - tracing::error!("Failed to fetch builds for task_id {}: {}", task_id, e); - None - } - } - } - - /// Create a new task ActiveModel for database insertion - pub fn create_task( - task_id: Uuid, - cl_id: i64, - task_name: Option, - template: Option, - created_at: DateTimeWithTimeZone, - ) -> ActiveModel { - ActiveModel { - id: Set(task_id), - cl_id: Set(cl_id), - task_name: Set(task_name), - template: Set(template), - created_at: Set(created_at), - } - } - - /// Insert a task directly into the database - pub async fn insert_task( - task_id: Uuid, - cl_id: i64, - task_name: Option, - template: Option, - created_at: DateTimeWithTimeZone, - db: &impl ConnectionTrait, - ) -> Result { - let task_model = Self::create_task(task_id, cl_id, task_name, template, created_at); - task_model.insert(db).await - } -} diff --git a/orion-server/src/orion_common/mod.rs b/orion-server/src/orion_common/mod.rs deleted file mode 100644 index 65880be0e..000000000 --- a/orion-server/src/orion_common/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod model; diff --git a/orion-server/src/orion_common/model.rs b/orion-server/src/orion_common/model.rs deleted file mode 100644 index 63c591202..000000000 --- a/orion-server/src/orion_common/model.rs +++ /dev/null @@ -1,29 +0,0 @@ -use serde::{Deserialize, Serialize}; -use utoipa::ToSchema; - -#[derive(Deserialize, ToSchema, Clone)] -pub struct Pagination { - pub page: u64, - pub per_page: u64, -} - -impl Default for Pagination { - fn default() -> Self { - Pagination { - page: 1, - per_page: 20, - } - } -} - -#[derive(Deserialize, ToSchema)] -pub struct PageParams { - pub pagination: Pagination, - pub additional: T, -} - -#[derive(PartialEq, Eq, Debug, Clone, Default, Serialize, Deserialize, ToSchema)] -pub struct CommonPage { - pub total: u64, - pub items: Vec, -} diff --git a/orion-server/src/repository/build_events.rs b/orion-server/src/repository/build_events.rs new file mode 100644 index 000000000..3f121b8be --- /dev/null +++ b/orion-server/src/repository/build_events.rs @@ -0,0 +1,22 @@ +use sea_orm::{ActiveValue::Set, ColumnTrait, ConnectionTrait, DbErr, EntityTrait, QueryFilter}; +use uuid::Uuid; + +pub struct BuildEvent; + +impl BuildEvent { + pub async fn update_build_complete_result( + build_id: &str, + exit_code: Option, + db_connection: &impl ConnectionTrait, + ) -> Result<(), DbErr> { + callisto::build_events::Entity::update_many() + .filter(callisto::build_events::Column::Id.eq(build_id.parse::().unwrap())) + .set(callisto::build_events::ActiveModel { + exit_code: Set(exit_code), + ..Default::default() + }) + .exec(db_connection) + .await?; + Ok(()) + } +} diff --git a/orion-server/src/model/build_targets.rs b/orion-server/src/repository/build_targets.rs similarity index 98% rename from orion-server/src/model/build_targets.rs rename to orion-server/src/repository/build_targets.rs index 018b22c30..a2ad9fb9c 100644 --- a/orion-server/src/model/build_targets.rs +++ b/orion-server/src/repository/build_targets.rs @@ -4,7 +4,8 @@ use sea_orm::{ }; use uuid::Uuid; -use crate::model::targets::TargetState; +use crate::entity::targets::TargetState; + /// A collection of utility methods for the `build_targets` database table. pub struct BuildTarget; @@ -91,6 +92,7 @@ impl BuildTarget { Ok(result) } + #[allow(dead_code)] pub async fn update_build_targets( target_state: TargetState, build_id: &str, diff --git a/orion-server/src/repository/builds.rs b/orion-server/src/repository/builds.rs new file mode 100644 index 000000000..f5d4f8558 --- /dev/null +++ b/orion-server/src/repository/builds.rs @@ -0,0 +1,50 @@ +use chrono::Utc; +use sea_orm::{ActiveModelTrait, ActiveValue::Set, ConnectionTrait, DbErr}; +use serde_json::Value; +use uuid::Uuid; + +pub struct BuildRepository; + +impl BuildRepository { + /// Create a new build ActiveModel for database insertion + pub fn create_build( + build_id: Uuid, + task_id: Uuid, + target_id: Uuid, + repo: String, + args: Option, + ) -> crate::entity::builds::ActiveModel { + let now = Utc::now().into(); + let repo_leaf = repo + .trim_end_matches('/') + .rsplit('/') + .next() + .unwrap_or(&repo) + .to_string(); + crate::entity::builds::ActiveModel { + id: Set(build_id), + task_id: Set(task_id), + target_id: Set(target_id), + exit_code: Set(None), + start_at: Set(now), + end_at: Set(None), + repo: Set(repo), + args: Set(args), + output_file: Set(format!("{}/{}/{}.log", task_id, repo_leaf, build_id)), + created_at: Set(now), + retry_count: Set(0), + } + } + + /// Insert a single build directly into the database + pub async fn insert_build( + build_id: Uuid, + task_id: Uuid, + target_id: Uuid, + repo: String, + db: &impl ConnectionTrait, + ) -> Result { + let build_model = Self::create_build(build_id, task_id, target_id, repo, None); + build_model.insert(db).await + } +} diff --git a/orion-server/src/repository/mod.rs b/orion-server/src/repository/mod.rs new file mode 100644 index 000000000..cffc6fd33 --- /dev/null +++ b/orion-server/src/repository/mod.rs @@ -0,0 +1,6 @@ +pub mod build_events; +pub mod build_targets; +pub mod builds; +pub mod orion_tasks; +pub mod targets; +pub mod tasks; diff --git a/orion-server/src/model/orion_tasks.rs b/orion-server/src/repository/orion_tasks.rs similarity index 58% rename from orion-server/src/model/orion_tasks.rs rename to orion-server/src/repository/orion_tasks.rs index b9e50cd14..54145169b 100644 --- a/orion-server/src/model/orion_tasks.rs +++ b/orion-server/src/repository/orion_tasks.rs @@ -1,16 +1,11 @@ use api_model::buck2::{status::Status, types::ProjectRelativePath}; -use chrono::Utc; use sea_orm::{ActiveModelTrait, ConnectionTrait, DbErr, IntoActiveModel}; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, to_value}; -use utoipa::ToSchema; +use serde_json::to_value; use uuid::Uuid; -#[allow(dead_code)] pub struct OrionTask; impl OrionTask { - #[allow(dead_code)] fn create_task( task_id: Uuid, cl_link: &str, @@ -26,7 +21,6 @@ impl OrionTask { }) } - #[allow(dead_code)] pub async fn insert_task( task_id: Uuid, cl_link: &str, @@ -39,24 +33,3 @@ impl OrionTask { task_model.into_active_model().insert(db).await } } - -#[derive(ToSchema, Serialize, Deserialize)] -pub struct OrionTaskDTO { - pub id: String, - pub changes: Value, - pub repo_name: String, - pub cl: String, - pub created_at: String, -} - -impl From<&callisto::orion_tasks::Model> for OrionTaskDTO { - fn from(model: &callisto::orion_tasks::Model) -> Self { - Self { - id: model.id.to_string(), - changes: model.changes.clone(), - repo_name: model.repo_name.clone(), - cl: model.cl.clone(), - created_at: model.created_at.with_timezone(&Utc).to_string(), - } - } -} diff --git a/orion-server/src/repository/targets.rs b/orion-server/src/repository/targets.rs new file mode 100644 index 000000000..da87be439 --- /dev/null +++ b/orion-server/src/repository/targets.rs @@ -0,0 +1,110 @@ +use chrono::Utc; +use sea_orm::{ + ActiveModelTrait, ActiveValue::Set, ColumnTrait, ConnectionTrait, DbErr, EntityTrait, + QueryFilter, RuntimeErr, sqlx, +}; +use uuid::Uuid; + +use crate::entity::targets::{self, TargetState}; + +pub struct TargetRepository; + +impl TargetRepository { + /// Create a new target ActiveModel ready for insertion + pub fn create_target(task_id: Uuid, target_path: impl Into) -> targets::ActiveModel { + let now = Utc::now().into(); + targets::ActiveModel { + id: Set(Uuid::now_v7()), + task_id: Set(task_id), + target_path: Set(target_path.into()), + state: Set(TargetState::Pending), + start_at: Set(None), + end_at: Set(None), + error_summary: Set(None), + created_at: Set(now), + } + } + + /// Insert a target directly into the database + pub async fn insert_target( + task_id: Uuid, + target_path: impl Into, + db: &impl ConnectionTrait, + ) -> Result { + let target_model = Self::create_target(task_id, target_path); + target_model.insert(db).await + } + + /// Find an existing target for the same task and path, or create a new one. + pub async fn find_or_create( + db: &impl ConnectionTrait, + task_id: Uuid, + target_path: impl Into + Clone, + ) -> Result { + let target_path_owned: String = target_path.clone().into(); + if let Some(target) = targets::Entity::find() + .filter(targets::Column::TaskId.eq(task_id)) + .filter(targets::Column::TargetPath.eq(target_path_owned.clone())) + .one(db) + .await? + { + return Ok(target); + } + + match Self::insert_target(task_id, target_path_owned.clone(), db).await { + Ok(model) => Ok(model), + Err(DbErr::Exec(RuntimeErr::SqlxError(sqlx::Error::Database(db_err)))) => { + if let Some(code) = db_err.code() + && code == "23505" + { + return targets::Entity::find() + .filter(targets::Column::TaskId.eq(task_id)) + .filter(targets::Column::TargetPath.eq(target_path_owned)) + .one(db) + .await? + .ok_or_else(|| DbErr::RecordNotFound("target".into())); + } + Err(DbErr::Exec(RuntimeErr::SqlxError(sqlx::Error::Database( + db_err, + )))) + } + Err(e) => Err(e), + } + } + + /// Update target state and optional timing/error metadata. + pub async fn update_state( + db: &impl ConnectionTrait, + target_id: Uuid, + state: TargetState, + start_at: Option, + end_at: Option, + error_summary: Option, + ) -> Result<(), DbErr> { + let mut active = targets::ActiveModel { + id: Set(target_id), + state: Set(state), + ..Default::default() + }; + + match state { + TargetState::Building | TargetState::Pending => { + active.end_at = Set(None); + active.error_summary = Set(None); + } + _ => {} + } + + if let Some(start_at) = start_at { + active.start_at = Set(Some(start_at)); + } + if let Some(end_at) = end_at { + active.end_at = Set(Some(end_at)); + } + if let Some(summary) = error_summary { + active.error_summary = Set(Some(summary)); + } + + active.update(db).await.map(|_| ()) + } +} diff --git a/orion-server/src/repository/tasks.rs b/orion-server/src/repository/tasks.rs new file mode 100644 index 000000000..3399f5ae9 --- /dev/null +++ b/orion-server/src/repository/tasks.rs @@ -0,0 +1,57 @@ +use sea_orm::{ + ActiveModelTrait, ActiveValue::Set, ColumnTrait, ConnectionTrait, DbErr, EntityTrait, + QueryFilter, +}; +use uuid::Uuid; + +pub struct TaskRepository; + +impl TaskRepository { + /// Retrieves build IDs associated with a task ID + pub async fn get_builds_by_task_id( + task_id: Uuid, + db: &sea_orm::DatabaseConnection, + ) -> Option> { + match crate::entity::builds::Entity::find() + .filter(crate::entity::builds::Column::TaskId.eq(task_id)) + .all(db) + .await + { + Ok(builds) => Some(builds.into_iter().map(|build| build.id).collect()), + Err(e) => { + tracing::error!("Failed to fetch builds for task_id {}: {}", task_id, e); + None + } + } + } + + /// Create a new task ActiveModel for database insertion + pub fn create_task( + task_id: Uuid, + cl_id: i64, + task_name: Option, + template: Option, + created_at: sea_orm::prelude::DateTimeWithTimeZone, + ) -> crate::entity::tasks::ActiveModel { + crate::entity::tasks::ActiveModel { + id: Set(task_id), + cl_id: Set(cl_id), + task_name: Set(task_name), + template: Set(template), + created_at: Set(created_at), + } + } + + /// Insert a task directly into the database + pub async fn insert_task( + task_id: Uuid, + cl_id: i64, + task_name: Option, + template: Option, + created_at: sea_orm::prelude::DateTimeWithTimeZone, + db: &impl ConnectionTrait, + ) -> Result { + let task_model = Self::create_task(task_id, cl_id, task_name, template, created_at); + task_model.insert(db).await + } +} diff --git a/orion-server/src/scheduler.rs b/orion-server/src/scheduler.rs index 1381df284..6ae8fdd82 100644 --- a/orion-server/src/scheduler.rs +++ b/orion-server/src/scheduler.rs @@ -19,20 +19,23 @@ use utoipa::ToSchema; use uuid::Uuid; use crate::{ - api::CoreWorkerStatus, auto_retry::AutoRetryJudger, - log::log_service::LogService, - model::{ - build_targets::{BuildTarget, BuildTargetDTO}, + entity::{ builds, targets::{self, TargetState}, }, + log::log_service::LogService, + model::dto::CoreWorkerStatus, + repository::{ + build_targets::{BuildTarget, BuildTargetDTO}, + targets::TargetRepository, + }, }; /// Request payload for creating a new build task -#[allow(dead_code)] #[derive(Debug, Clone, Deserialize, ToSchema)] pub struct BuildRequest { + #[allow(dead_code)] pub changes: Vec>, /// Buck2 target path (e.g. //app:server). Optional for backward compatibility. #[serde(default, alias = "target_path")] @@ -188,12 +191,12 @@ pub struct PendingBuildEventV2 { pub struct BuildInfo { pub event_payload: BuildEventPayload, pub target_id: Uuid, + #[allow(dead_code)] pub target_path: String, pub changes: Vec>, #[allow(dead_code)] pub started_at: DateTimeUtc, pub auto_retry_judger: AutoRetryJudger, - #[allow(dead_code)] pub worker_id: String, } @@ -266,21 +269,6 @@ pub struct TaskScheduler { pub conn: DatabaseConnection, } -/// Errors when reading a log segment -#[allow(dead_code)] -#[derive(Debug)] -pub enum LogReadError { - NotFound, - OffsetOutOfRange { size: u64 }, - Io(std::io::Error), -} - -impl From for LogReadError { - fn from(e: std::io::Error) -> Self { - Self::Io(e) - } -} - impl TaskScheduler { /// Create new task scheduler instance pub fn new( @@ -308,7 +296,7 @@ impl TaskScheduler { target_path: &str, ) -> Result { // Find-or-create target for (task_id, target_path) - targets::Entity::find_or_create(&self.conn, task_id, target_path.to_string()).await + TargetRepository::find_or_create(&self.conn, task_id, target_path.to_string()).await } /// Bound corresponding task build ID to the given task and enqueue @@ -472,7 +460,6 @@ impl TaskScheduler { } /// Search available worker and claim the worker for current build - #[allow(dead_code)] pub fn search_and_claim_worker(&self, build_id: &str) -> Option { let idle_workers: Vec = self .workers @@ -496,7 +483,6 @@ impl TaskScheduler { } } - #[allow(dead_code)] pub async fn release_worker(&self, worker_id: &str) { tracing::info!("Releasing worker {} back to idle", worker_id); if let Some(mut worker) = self.workers.get_mut(worker_id) { @@ -625,7 +611,7 @@ impl TaskScheduler { if let Some(mut worker) = self.workers.get_mut(&chosen_id) { if worker.sender.send(msg).is_ok() { // Only mark Building after send succeeds - if let Err(e) = targets::update_state( + if let Err(e) = TargetRepository::update_state( &self.conn, //TODO: update target_id here pending_build_event.target_id.unwrap_or(Uuid::nil()), @@ -656,7 +642,7 @@ impl TaskScheduler { Ok(()) } else { // Send failed: best-effort mark target back to Pending - let _ = targets::update_state( + let _ = TargetRepository::update_state( &self.conn, // TODO: update target_id here pending_build_event.target_id.unwrap_or(Uuid::nil()), diff --git a/orion-server/src/server.rs b/orion-server/src/server.rs index a162560d8..df8010973 100644 --- a/orion-server/src/server.rs +++ b/orion-server/src/server.rs @@ -1,6 +1,5 @@ use std::{env, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration}; -use api_model::buck2::types::TaskPhase; use axum::{Router, routing::get}; use chrono::{FixedOffset, Utc}; use common::{ @@ -19,75 +18,18 @@ use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; use crate::{ - api::{self, AppState}, + api, + api_doc::ApiDoc, + app_state::AppState, buck2::set_mono_base_url, + entity::builds, log::{ log_service::LogService, store::{LogStore, io_orbit_store::IoOrbitLogStore, local_log_store, noop_log_store}, }, - model::builds, + repository::targets::TargetRepository, }; -/// OpenAPI documentation configuration -#[derive(OpenApi)] -#[openapi( - paths( - api::task_handler, - api::task_build_list_handler, - api::task_output_handler, - api::task_history_output_handler, - api::target_logs_handler, - api::tasks_handler, - api::task_targets_handler, - api::task_targets_summary_handler, - api::get_orion_clients_info, - api::get_orion_client_status_by_id, - api::build_retry_handler, - api::health_check_handler, - api::task_retry_handler, - api::task_get_handler, - api::build_event_get_handler, - api::targets_get_handler, - api::build_state_handler, - api::build_logs_handler, - api::latest_build_result_handler, - api::targets_status_handler, - api::single_target_status_handle, - ), - components( - schemas( - crate::scheduler::BuildRequest, - api_model::buck2::types::LogLinesResponse, - api_model::buck2::types::TargetLogLinesResponse, - api_model::buck2::types::LogErrorResponse, - api::TaskStatusEnum, - api::BuildDTO, - api::TargetDTO, - api::TargetSummaryDTO, - api_model::buck2::types::TargetLogQuery, - api_model::buck2::types::LogReadMode, - api_model::buck2::types::TaskHistoryQuery, - api::TaskInfoDTO, - api::OrionClientInfo, - api::OrionClientStatus, - api::CoreWorkerStatus, - api::OrionClientQuery, - crate::model::targets::TargetState, - TaskPhase, - api::MessageResponse, - crate::model::build_events::BuildEventDTO, - crate::model::orion_tasks::OrionTaskDTO, - api::BuildTargetDTO, - api::BuildEventState, - api_model::buck2::types::TargetStatusResponse, - ) - ), - tags( - (name = "Build", description = "Build related endpoints") - ) -)] -pub struct ApiDoc; - pub async fn init_log_service(config: Config) -> Result { // Read buffer size from environment, defaulting to 4096 if unset or invalid @@ -214,7 +156,7 @@ pub async fn start_server() { tokio::spawn(start_health_check_task(state.clone())); // Start queue manager - tokio::spawn(api::start_queue_manager(state.clone())); + tokio::spawn(state.clone().start_queue_manager()); // Start background dp operation state.start_background_tasks(); @@ -363,10 +305,10 @@ async fn start_health_check_task(state: AppState) { ); } - if let Err(e) = crate::model::targets::update_state( + if let Err(e) = TargetRepository::update_state( &state.conn, build_model.target_id, - crate::model::targets::TargetState::Interrupted, + crate::entity::targets::TargetState::Interrupted, None, Some(Utc::now().with_timezone(&utc_offset)), None, diff --git a/orion-server/src/service/api_v2_service.rs b/orion-server/src/service/api_v2_service.rs new file mode 100644 index 000000000..0fdf899f0 --- /dev/null +++ b/orion-server/src/service/api_v2_service.rs @@ -0,0 +1,1337 @@ +use std::{collections::HashMap, convert::Infallible, pin::Pin, sync::Arc, time::Duration}; + +use api_model::{ + buck2::{ + api::{OrionBuildResult, OrionServerResponse, TaskBuildRequest}, + status::Status, + types::{ + LogErrorResponse, LogLinesResponse, LogReadMode, ProjectRelativePath, + TargetLogLinesResponse, TargetLogQuery, TaskHistoryQuery, + }, + ws::WSMessage, + }, + common::{CommonPage, PageParams}, +}; +use axum::{ + Json, + http::StatusCode, + response::{IntoResponse, Response, Sse, sse::Event}, +}; +use dashmap::DashMap; +use futures::stream::select; +use futures_util::{Stream, StreamExt}; +use rand::RngExt; +use sea_orm::{ColumnTrait, EntityTrait, QueryFilter as _, QueryOrder, QuerySelect}; +use serde_json::{Value, json}; +use tokio::sync::watch; +use tokio_stream::wrappers::IntervalStream; +use uuid::Uuid; + +use crate::{ + app_state::AppState, + auto_retry::AutoRetryJudger, + entity::{builds, targets, targets::TargetState, tasks}, + log::log_service::LogService, + model::{ + dto::{ + BuildDTO, BuildEventDTO, BuildEventState, BuildTargetDTO, MessageResponse, + OrionClientInfo, OrionClientQuery, OrionClientStatus, OrionTaskDTO, TargetDTO, + TargetSummaryDTO, TaskInfoDTO, + }, + task_status::TaskStatusEnum, + }, + repository::{ + build_targets::BuildTarget, builds::BuildRepository, orion_tasks::OrionTask, + targets::TargetRepository, tasks::TaskRepository, + }, + scheduler::{BuildEventPayload, BuildInfo, TaskQueueStats, WorkerStatus}, +}; + +type MessageErrorResponse = (StatusCode, Json); +type JsonValueErrorResponse = (StatusCode, Json); +type LogSseStream = Pin> + Send>>; + +pub async fn task_output(state: &AppState, id: &str) -> Result, StatusCode> { + if !state.scheduler.active_builds.contains_key(id) { + return Err(StatusCode::NOT_FOUND); + } + + let (stop_tx, stop_rx) = watch::channel(true); + + let log_stop_rx = stop_rx.clone(); + let build_id = id.to_string(); + let log_stream = state + .log_service + .subscribe_for_build(build_id.clone()) + .map(|log_event| { + Ok::(Event::default().event("log").data(log_event.line)) + }) + .take_while(move |_| { + let stop_rx = log_stop_rx.clone(); + async move { *stop_rx.borrow() } + }); + + let heart_stop_rx = stop_rx.clone(); + let heartbeat_stream = IntervalStream::new(tokio::time::interval(Duration::from_secs(15))) + .map(|_| Ok::(Event::default().comment("heartbeat"))) + .take_while(move |_| { + let stop_rx_clone = heart_stop_rx.clone(); + async move { *stop_rx_clone.borrow() } + }); + + let stop_tx_clone = stop_tx.clone(); + let state_clone = state.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + if !state_clone.scheduler.active_builds.contains_key(&build_id) { + let _ = stop_tx_clone.send(false); + break; + } + } + }); + + Ok(Sse::new(Box::pin(select(log_stream, heartbeat_stream)))) +} + +fn message_response(message: impl Into) -> Json { + Json(MessageResponse { + message: message.into(), + }) +} + +fn message_error(status: StatusCode, message: impl Into) -> MessageErrorResponse { + (status, message_response(message)) +} + +fn value_error(status: StatusCode, message: impl Into) -> JsonValueErrorResponse { + (status, Json(json!({ "message": message.into() }))) +} + +fn parse_uuid_or_message_error( + raw_id: &str, + invalid_message: &str, +) -> Result { + raw_id + .parse::() + .map_err(|_| message_error(StatusCode::BAD_REQUEST, invalid_message)) +} + +fn parse_uuid_or_value_error( + raw_id: &str, + invalid_message: &str, +) -> Result { + raw_id + .parse::() + .map_err(|_| value_error(StatusCode::BAD_REQUEST, invalid_message)) +} + +async fn task_exists_by_id( + conn: &sea_orm::DatabaseConnection, + task_id: Uuid, +) -> Result { + callisto::orion_tasks::Entity::find_by_id(task_id) + .one(conn) + .await + .map(|task| task.is_some()) +} + +pub async fn task_retry( + state: &AppState, + id: &str, +) -> Result, MessageErrorResponse> { + let task_uuid = parse_uuid_or_message_error(id, "Invalid task ID format")?; + let task = callisto::orion_tasks::Entity::find_by_id(task_uuid) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch task {}: {}", id, e); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })? + .ok_or_else(|| message_error(StatusCode::NOT_FOUND, "Task not found"))?; + + tracing::info!("Task retry requested for task {} (CL: {})", id, task.cl); + Ok(message_response(format!("Task {} queued for retry", id))) +} + +pub async fn task_get( + state: &AppState, + cl: &str, +) -> Result, JsonValueErrorResponse> { + let tasks: Vec = callisto::orion_tasks::Entity::find() + .filter(callisto::orion_tasks::Column::Cl.eq(cl)) + .all(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch tasks by CL {}: {}", cl, e); + value_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + + match tasks.len() { + 0 => Err(value_error(StatusCode::NOT_FOUND, "Not found task")), + 1 => Ok(Json(OrionTaskDTO::from(&tasks[0]))), + _ => Err(value_error(StatusCode::BAD_REQUEST, "Multiple tasks")), + } +} + +pub async fn build_event_get( + state: &AppState, + task_id: &str, +) -> Result>, JsonValueErrorResponse> { + let task_uuid = parse_uuid_or_value_error(task_id, "Invalid task ID")?; + let task_exists = task_exists_by_id(&state.conn, task_uuid) + .await + .map_err(|e| { + tracing::error!("Failed to verify task existence {}: {}", task_id, e); + value_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + if !task_exists { + return Err(value_error(StatusCode::NOT_FOUND, "Task not found")); + } + + let build_events = callisto::build_events::Entity::find() + .filter(callisto::build_events::Column::TaskId.eq(task_uuid)) + .all(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch build events for task {}: {}", task_id, e); + value_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + + Ok(Json( + build_events + .into_iter() + .map(|m| BuildEventDTO::from(&m)) + .collect(), + )) +} + +pub async fn targets_get( + state: &AppState, + task_id: &str, +) -> Result>, MessageErrorResponse> { + let task_uuid = parse_uuid_or_message_error(task_id, "Invalid task ID")?; + let task_exists = task_exists_by_id(&state.conn, task_uuid) + .await + .map_err(|e| { + tracing::error!("Failed to verify task existence {}: {}", task_id, e); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + if !task_exists { + return Err(message_error(StatusCode::NOT_FOUND, "Task not found")); + } + + let build_targets = callisto::build_targets::Entity::find() + .filter(callisto::build_targets::Column::TaskId.eq(task_uuid)) + .all(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch build targets for task {}: {}", task_id, e); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + + Ok(Json( + build_targets + .into_iter() + .map(|build_target| BuildTargetDTO { + id: build_target.id.to_string(), + task_id: build_target.task_id.to_string(), + path: build_target.path, + latest_state: build_target.latest_state, + }) + .collect(), + )) +} + +pub async fn build_logs( + state: &AppState, + build_id: &str, +) -> Result, (StatusCode, Json)> { + let build_uuid = build_id.parse::().map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(LogErrorResponse { + message: "Invalid build ID".to_string(), + }), + ) + })?; + + let build_event = callisto::build_events::Entity::find_by_id(build_uuid) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch build event {}: {}", build_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Database error".to_string(), + }), + ) + })? + .ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json(LogErrorResponse { + message: "Build event not found".to_string(), + }), + ) + })?; + + let orion_task = callisto::orion_tasks::Entity::find_by_id(build_event.task_id) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch Orion task {}: {}", build_event.task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Database error".to_string(), + }), + ) + })? + .ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json(LogErrorResponse { + message: "Task not found".to_string(), + }), + ) + })?; + + let task_id = build_event.task_id.to_string(); + let repo_name = &orion_task.repo_name; + let log_content = state + .log_service + .read_full_log(&task_id, repo_name, build_id) + .await + .map_err(|e| { + tracing::error!("Failed to read log for build {}: {}", build_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to read log".to_string(), + }), + ) + })?; + + let lines: Vec = log_content.lines().map(str::to_string).collect(); + Ok(Json(LogLinesResponse { + len: lines.len(), + data: lines, + })) +} + +pub async fn build_state( + state: &AppState, + build_id: &str, +) -> Result, MessageErrorResponse> { + let build_uuid = parse_uuid_or_message_error(build_id, "Invalid build ID")?; + if state.scheduler.active_builds.contains_key(build_id) { + return Ok(Json(BuildEventState::Running)); + } + + let build_event = callisto::build_events::Entity::find_by_id(build_uuid) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch build event {}: {}", build_id, e); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })? + .ok_or_else(|| message_error(StatusCode::NOT_FOUND, "Build not found"))?; + + let state_enum = match (build_event.end_at, build_event.exit_code) { + (None, _) => BuildEventState::Running, + (Some(_), Some(0)) => BuildEventState::Success, + (Some(_), Some(_)) => BuildEventState::Failure, + (Some(_), None) => BuildEventState::Failure, + }; + Ok(Json(state_enum)) +} + +pub async fn latest_build_result( + state: &AppState, + task_id: &str, +) -> Result, MessageErrorResponse> { + let task_uuid = parse_uuid_or_message_error(task_id, "Invalid task ID")?; + let task_exists = task_exists_by_id(&state.conn, task_uuid) + .await + .map_err(|e| { + tracing::error!("Failed to verify task existence {}: {}", task_id, e); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + if !task_exists { + return Err(message_error(StatusCode::NOT_FOUND, "Task not found")); + } + + let latest_build_event = callisto::build_events::Entity::find() + .filter(callisto::build_events::Column::TaskId.eq(task_uuid)) + .order_by_desc(callisto::build_events::Column::StartAt) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!( + "Failed to fetch latest build event for task {}: {}", + task_id, + e + ); + message_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error") + })?; + + let build_event = latest_build_event.ok_or_else(|| { + message_error(StatusCode::NOT_FOUND, "No build events found for this task") + })?; + + let state_enum = match (build_event.end_at, build_event.exit_code) { + (None, _) => BuildEventState::Running, + (Some(_), Some(0)) => BuildEventState::Success, + (Some(_), Some(_)) => BuildEventState::Failure, + (Some(_), None) => BuildEventState::Failure, + }; + Ok(Json(state_enum)) +} + +pub async fn queue_stats(state: &AppState) -> (StatusCode, Json) { + let stats = state.scheduler.get_queue_stats().await; + (StatusCode::OK, Json(stats)) +} + +pub async fn health_check(state: &AppState) -> (StatusCode, Json) { + match tasks::Entity::find().limit(1).all(&state.conn).await { + Ok(_) => (StatusCode::OK, Json(json!({"status": "healthy"}))), + Err(e) => { + tracing::error!("Health check failed: {}", e); + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"status": "unhealthy", "error": "database connectivity check failed"})), + ) + } + } +} + +pub async fn task_history_output( + state: &AppState, + params: &TaskHistoryQuery, +) -> Result, (StatusCode, Json)> { + let log_result = if matches!((params.start, params.end), (None, None)) { + state + .log_service + .read_full_log(¶ms.task_id, ¶ms.repo, ¶ms.build_id) + .await + } else { + let start = params.start.unwrap_or(0); + let end = params.end.unwrap_or(usize::MAX); + state + .log_service + .read_log_range(¶ms.task_id, ¶ms.repo, ¶ms.build_id, start, end) + .await + }; + + let log_content = match log_result { + Ok(content) => content, + Err(e) => { + tracing::error!("read log failed: {:?}", e); + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to read log file".to_string(), + }), + )); + } + }; + + let lines: Vec = log_content.lines().map(str::to_string).collect(); + Ok(Json(LogLinesResponse { + len: lines.len(), + data: lines, + })) +} + +pub async fn target_logs( + state: &AppState, + target_id: &str, + params: &TargetLogQuery, +) -> Result, (StatusCode, Json)> { + let target_uuid = target_id.parse::().map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(LogErrorResponse { + message: "Invalid target id".to_string(), + }), + ) + })?; + + let target_model = match targets::Entity::find_by_id(target_uuid) + .one(&state.conn) + .await + { + Ok(Some(target)) => target, + Ok(None) => { + return Err(( + StatusCode::NOT_FOUND, + Json(LogErrorResponse { + message: "Target not found".to_string(), + }), + )); + } + Err(err) => { + tracing::error!("Failed to load target {}: {}", target_id, err); + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to read target".to_string(), + }), + )); + } + }; + + let build_model = if let Some(build_id) = params.build_id.as_ref() { + let build_uuid = build_id.parse::().map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(LogErrorResponse { + message: "Invalid build id".to_string(), + }), + ) + })?; + + match builds::Entity::find_by_id(build_uuid) + .filter(builds::Column::TargetId.eq(target_uuid)) + .one(&state.conn) + .await + { + Ok(Some(build)) => build, + Ok(None) => { + return Err(( + StatusCode::NOT_FOUND, + Json(LogErrorResponse { + message: "Build not found for target".to_string(), + }), + )); + } + Err(err) => { + tracing::error!("Failed to load build {}: {}", build_uuid, err); + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to load build".to_string(), + }), + )); + } + } + } else { + match builds::Entity::find() + .filter(builds::Column::TargetId.eq(target_uuid)) + .order_by_desc(builds::Column::EndAt) + .order_by_desc(builds::Column::CreatedAt) + .one(&state.conn) + .await + { + Ok(Some(build)) => build, + Ok(None) => { + return Err(( + StatusCode::NOT_FOUND, + Json(LogErrorResponse { + message: "No builds for target".to_string(), + }), + )); + } + Err(err) => { + tracing::error!("Failed to load build for target {}: {}", target_uuid, err); + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to load build".to_string(), + }), + )); + } + } + }; + + let repo_segment = LogService::last_segment(&build_model.repo); + let log_result = if matches!(params.r#type, LogReadMode::Segment) { + let offset = params.offset.unwrap_or(0); + let limit = params.limit.unwrap_or(200); + state + .log_service + .read_log_range( + &target_model.task_id.to_string(), + &repo_segment, + &build_model.id.to_string(), + offset, + offset + limit, + ) + .await + } else { + state + .log_service + .read_full_log( + &target_model.task_id.to_string(), + &repo_segment, + &build_model.id.to_string(), + ) + .await + }; + + match log_result { + Ok(content) => { + let lines: Vec = content.lines().map(str::to_string).collect(); + Ok(Json(TargetLogLinesResponse { + len: lines.len(), + data: lines, + build_id: build_model.id.to_string(), + })) + } + Err(e) => { + tracing::error!( + "Failed to read logs for target {} build {}: {}", + target_uuid, + build_model.id, + e + ); + Err(( + StatusCode::INTERNAL_SERVER_ERROR, + Json(LogErrorResponse { + message: "Failed to read log".to_string(), + }), + )) + } + } +} + +async fn assemble_task_info( + task: tasks::Model, + state: &AppState, + active_builds: &Arc>, +) -> TaskInfoDTO { + let target_models = targets::Entity::find() + .filter(targets::Column::TaskId.eq(task.id)) + .all(&state.conn) + .await + .unwrap_or_else(|_| vec![]); + let build_models = builds::Entity::find() + .filter(builds::Column::TaskId.eq(task.id)) + .all(&state.conn) + .await + .unwrap_or_else(|_| vec![]); + + let target_map: HashMap = + target_models.iter().cloned().map(|t| (t.id, t)).collect(); + let mut build_list: Vec = Vec::new(); + let mut target_build_map: HashMap> = HashMap::new(); + + for build_model in build_models { + let build_id_str = build_model.id.to_string(); + let is_active = active_builds.contains_key(&build_id_str); + let status = BuildDTO::determine_status(&build_model, is_active); + let mut dto = BuildDTO::from_model( + build_model.clone(), + target_map.get(&build_model.target_id), + status.clone(), + ); + if matches!(status, TaskStatusEnum::Failed) + && let Some(t) = target_map.get(&build_model.target_id) + && let Some(summary) = &t.error_summary + { + dto.cause_by = Some(summary.clone()); + } + target_build_map + .entry(build_model.target_id) + .or_default() + .push(dto.clone()); + build_list.push(dto); + } + + let mut target_list: Vec = Vec::new(); + for target in target_models { + let target_builds = target_build_map.remove(&target.id).unwrap_or_default(); + target_list.push(crate::entity::targets::TargetWithBuilds::from_model( + target, + target_builds, + )); + } + + TaskInfoDTO { + task_id: task.id.to_string(), + cl_id: task.cl_id, + task_name: task.task_name, + template: task.template, + created_at: task.created_at.with_timezone(&chrono::Utc).to_rfc3339(), + build_list, + targets: target_list, + } +} + +pub async fn tasks_by_cl( + state: &AppState, + cl: i64, +) -> Result>, (StatusCode, Json)> { + let active_builds = state.scheduler.active_builds.clone(); + let task_models = tasks::Entity::find() + .filter(tasks::Column::ClId.eq(cl)) + .all(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch tasks: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": "Failed to fetch tasks"})), + ) + })?; + let mut result = Vec::with_capacity(task_models.len()); + for model in task_models { + result.push(assemble_task_info(model, state, &active_builds).await); + } + Ok(Json(result)) +} + +pub async fn task_targets( + state: &AppState, + task_id: &str, +) -> Result, (StatusCode, Json)> { + let task_uuid = task_id.parse::().map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(json!({"message": "Invalid task ID"})), + ) + })?; + let task_model = tasks::Entity::find_by_id(task_uuid) + .one(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch task {}: {}", task_id, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": "Failed to fetch task"})), + ) + })? + .ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json(json!({"message": "Task not found"})), + ) + })?; + let info = assemble_task_info(task_model, state, &state.scheduler.active_builds).await; + Ok(Json(info)) +} + +pub async fn task_targets_summary( + state: &AppState, + task_id: &str, +) -> Result, (StatusCode, Json)> { + let task_uuid = task_id.parse::().map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(json!({ "message": "Invalid task ID" })), + ) + })?; + let target_models = targets::Entity::find() + .filter(targets::Column::TaskId.eq(task_uuid)) + .all(&state.conn) + .await + .map_err(|e| { + tracing::error!("Failed to fetch target summary: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "message": "Failed to fetch target summary" })), + ) + })?; + let mut summary = TargetSummaryDTO { + task_id: task_id.to_string(), + pending: 0, + building: 0, + completed: 0, + failed: 0, + interrupted: 0, + uninitialized: 0, + }; + for target in target_models { + match target.state { + TargetState::Pending => summary.pending += 1, + TargetState::Building => summary.building += 1, + TargetState::Completed => summary.completed += 1, + TargetState::Failed => summary.failed += 1, + TargetState::Interrupted => summary.interrupted += 1, + TargetState::Uninitialized => summary.uninitialized += 1, + } + } + Ok(Json(summary)) +} + +async fn immediate_retry_work( + state: &AppState, + build_id: Uuid, + idle_workers: &[String], + build: &builds::Model, + target: &targets::Model, + req: &api_model::buck2::api::RetryBuildRequest, + retry_count: i32, +) -> bool { + let chosen_index = { + let mut rng = rand::rng(); + rng.random_range(0..idle_workers.len()) + }; + let chosen_id = idle_workers[chosen_index].clone(); + let start_at = chrono::Utc::now(); + let event = BuildEventPayload::new( + build.id, + build.task_id, + req.cl_link.clone(), + build.repo.clone(), + retry_count, + ); + let build_info = BuildInfo { + event_payload: event, + target_id: target.id, + target_path: target.target_path.clone(), + changes: req.changes.clone(), + worker_id: chosen_id.clone(), + auto_retry_judger: AutoRetryJudger::new(), + started_at: start_at, + }; + let msg = api_model::buck2::ws::WSMessage::TaskBuild { + build_id: build.id.to_string(), + repo: build.repo.to_string(), + changes: req.changes.clone(), + cl_link: req.cl_link.to_string(), + }; + if let Some(mut worker) = state.scheduler.workers.get_mut(&chosen_id) + && worker.sender.send(msg).is_ok() + { + worker.status = WorkerStatus::Busy { + build_id: build_id.to_string(), + phase: None, + }; + if let Err(e) = TargetRepository::update_state( + &state.conn, + target.id, + TargetState::Building, + Some(start_at.with_timezone( + &chrono::FixedOffset::east_opt(0).unwrap_or_else(|| unreachable!()), + )), + None, + None, + ) + .await + { + tracing::error!("Failed to update target state to Building: {}", e); + } + state + .scheduler + .active_builds + .insert(build.id.to_string(), build_info); + true + } else { + false + } +} + +pub async fn build_retry( + state: &AppState, + req: api_model::buck2::api::RetryBuildRequest, +) -> Response { + let build_id = match req.build_id.parse::() { + Ok(uuid) => uuid, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(json!({"message": "Invalid build ID format"})), + ) + .into_response(); + } + }; + if state.scheduler.active_builds.contains_key(&req.build_id) { + return ( + StatusCode::BAD_REQUEST, + Json(json!({"message": "The build already exists"})), + ) + .into_response(); + } + let build = match builds::Entity::find_by_id(build_id).one(&state.conn).await { + Ok(Some(build)) => build, + Ok(None) => { + return ( + StatusCode::NOT_FOUND, + Json(json!({"message": "Build not found"})), + ) + .into_response(); + } + Err(_) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": "Database find failed"})), + ) + .into_response(); + } + }; + let retry_count = build.retry_count + 1; + let target_model = match targets::Entity::find_by_id(build.target_id) + .one(&state.conn) + .await + { + Ok(Some(target)) => target, + Ok(None) => { + return ( + StatusCode::NOT_FOUND, + Json(json!({"message": "Target not found for build"})), + ) + .into_response(); + } + Err(_) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": "Database find failed"})), + ) + .into_response(); + } + }; + + let idle_workers = state.scheduler.get_idle_workers(); + if idle_workers.is_empty() { + let new_build_id = Uuid::now_v7(); + match state + .scheduler + .enqueue_task_with_build_id( + new_build_id, + build.task_id, + &req.cl_link, + build.repo.clone(), + req.changes.clone(), + target_model.target_path.clone(), + retry_count, + ) + .await + { + Ok(()) => ( + StatusCode::OK, + Json(json!({"message": "Build queued for later processing"})), + ) + .into_response(), + Err(_) => ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"message": "No available workers at the moment"})), + ) + .into_response(), + } + } else if immediate_retry_work( + state, + build_id, + &idle_workers, + &build, + &target_model, + &req, + retry_count, + ) + .await + { + ( + StatusCode::OK, + Json(json!({"message": "Build retry dispatched immediately to worker"})), + ) + .into_response() + } else { + ( + StatusCode::BAD_GATEWAY, + Json(json!({"message": "Failed to dispatch build retry to worker"})), + ) + .into_response() + } +} + +async fn activate_worker( + build_info: &BuildInfo, + scheduler: &crate::scheduler::TaskScheduler, +) -> OrionBuildResult { + let msg = WSMessage::TaskBuild { + build_id: build_info.event_payload.build_event_id.to_string(), + repo: build_info.event_payload.repo.clone(), + changes: build_info.changes.clone(), + cl_link: build_info.event_payload.cl_link.clone(), + }; + if let Some(worker) = scheduler.workers.get_mut(&build_info.worker_id) + && worker.sender.send(msg).is_ok() + { + scheduler.active_builds.insert( + build_info.event_payload.build_event_id.to_string(), + build_info.clone(), + ); + return OrionBuildResult { + build_id: build_info.event_payload.build_event_id.to_string(), + status: "dispatched".to_string(), + message: format!("Build dispatched to worker {}", build_info.worker_id), + }; + } + scheduler.release_worker(&build_info.worker_id).await; + OrionBuildResult { + build_id: build_info.event_payload.build_event_id.to_string(), + status: "error".to_string(), + message: "Failed to dispatch task to worker".to_string(), + } +} + +async fn handle_immediate_task_dispatch_v2( + state: &AppState, + task_id: Uuid, + repo: &str, + cl_link: &str, + changes: Vec>, +) -> OrionBuildResult { + let build_event_id = Uuid::now_v7(); + let Some(chosen_id) = state + .scheduler + .search_and_claim_worker(&build_event_id.to_string()) + else { + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: "No available workers at the moment".to_string(), + }; + }; + + let target_id = Uuid::now_v7(); + let target_path = + match BuildTarget::insert_default_target(target_id, task_id, &state.conn).await { + Ok(default_path) => default_path, + Err(_) => { + state.scheduler.release_worker(&chosen_id).await; + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: format!("Failed to prepare target for task {}", task_id), + }; + } + }; + + let event = BuildEventPayload::new( + build_event_id, + task_id, + cl_link.to_string(), + repo.to_string(), + 0, + ); + let build_info = BuildInfo { + event_payload: event.clone(), + changes: changes.clone(), + target_id, + target_path, + worker_id: chosen_id, + auto_retry_judger: AutoRetryJudger::new(), + started_at: chrono::Utc::now(), + }; + + if let Err(e) = callisto::build_events::Model::insert_build( + build_event_id, + task_id, + repo.to_string(), + &state.conn, + ) + .await + { + state.scheduler.release_worker(&build_info.worker_id).await; + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: format!( + "Failed to insert build event into database for task {}: {}", + task_id, e + ), + }; + } + activate_worker(&build_info, &state.scheduler).await +} + +async fn handle_immediate_task_dispatch( + state: &AppState, + task_id: Uuid, + repo: &str, + cl_link: &str, + changes: Vec>, +) -> OrionBuildResult { + let idle_workers = state.scheduler.get_idle_workers(); + if idle_workers.is_empty() { + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: "No available workers at the moment".to_string(), + }; + } + let chosen_id = { + let mut rng = rand::rng(); + idle_workers[rng.random_range(0..idle_workers.len())].clone() + }; + let build_id = Uuid::now_v7(); + let target_model = match state.scheduler.ensure_target(task_id, "").await { + Ok(target) => target, + Err(_) => { + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: "Failed to prepare target ".to_string(), + }; + } + }; + let start_at = chrono::Utc::now(); + let _ = TargetRepository::update_state( + &state.conn, + target_model.id, + TargetState::Building, + Some( + start_at + .with_timezone(&chrono::FixedOffset::east_opt(0).unwrap_or_else(|| unreachable!())), + ), + None, + None, + ) + .await; + + let build_info = BuildInfo { + event_payload: BuildEventPayload::new( + build_id, + task_id, + cl_link.to_string(), + repo.to_string(), + 0, + ), + changes: changes.clone(), + target_id: target_model.id, + target_path: target_model.target_path.clone(), + worker_id: chosen_id.clone(), + auto_retry_judger: AutoRetryJudger::new(), + started_at: start_at, + }; + if BuildRepository::insert_build( + build_id, + task_id, + target_model.id, + repo.to_string(), + &state.conn, + ) + .await + .is_err() + { + return OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: "Failed to insert builds into database".to_string(), + }; + } + let msg = WSMessage::TaskBuild { + build_id: build_id.to_string(), + repo: repo.to_string(), + changes, + cl_link: cl_link.to_string(), + }; + if let Some(mut worker) = state.scheduler.workers.get_mut(&chosen_id) + && worker.sender.send(msg).is_ok() + { + worker.status = WorkerStatus::Busy { + build_id: build_id.to_string(), + phase: None, + }; + state + .scheduler + .active_builds + .insert(build_id.to_string(), build_info); + OrionBuildResult { + build_id: build_id.to_string(), + status: "dispatched".to_string(), + message: format!("Build dispatched to worker {}", chosen_id), + } + } else { + OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: "Failed to dispatch task to worker".to_string(), + } + } +} + +pub async fn task_handler_v2(state: &AppState, req: TaskBuildRequest) -> Response { + let task_id = Uuid::now_v7(); + if let Err(err) = + OrionTask::insert_task(task_id, &req.cl_link, &req.repo, &req.changes, &state.conn).await + { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": format!("Failed to insert task into database: {}", err)})), + ) + .into_response(); + } + let result = if state.scheduler.has_idle_workers() { + handle_immediate_task_dispatch_v2(state, task_id, &req.repo, &req.cl_link, req.changes) + .await + } else { + match state + .scheduler + .enqueue_task_v2(task_id, &req.cl_link, req.repo, req.changes, 0) + .await + { + Ok(build_id) => OrionBuildResult { + build_id: build_id.to_string(), + status: "queued".to_string(), + message: "Task queued for processing when workers become available".to_string(), + }, + Err(e) => { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"message": format!("Unable to queue task: {}", e)})), + ) + .into_response(); + } + } + }; + ( + StatusCode::OK, + Json(OrionServerResponse { + task_id: task_id.to_string(), + results: vec![result], + }), + ) + .into_response() +} + +pub async fn task_handler_v1(state: &AppState, req: TaskBuildRequest) -> Response { + let task_id = Uuid::now_v7(); + let task_name = format!("CL-{}-{}", req.cl_link, task_id); + if let Err(err) = TaskRepository::insert_task( + task_id, + req.cl_id, + Some(task_name), + None, + chrono::Utc::now().into(), + &state.conn, + ) + .await + { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": format!("Failed to insert task into database: {}", err)})), + ) + .into_response(); + } + let mut results = Vec::new(); + if state.scheduler.has_idle_workers() { + results.push( + handle_immediate_task_dispatch( + state, + task_id, + &req.repo, + &req.cl_link, + req.changes.clone(), + ) + .await, + ); + } else { + match state + .scheduler + .enqueue_task( + task_id, + &req.cl_link, + req.repo.clone(), + req.changes.clone(), + None, + 0, + ) + .await + { + Ok(build_id) => results.push(OrionBuildResult { + build_id: build_id.to_string(), + status: "queued".to_string(), + message: "Task queued for processing when workers become available".to_string(), + }), + Err(e) => results.push(OrionBuildResult { + build_id: "".to_string(), + status: "error".to_string(), + message: format!("Unable to queue task: {}", e), + }), + } + } + ( + StatusCode::OK, + Json(OrionServerResponse { + task_id: task_id.to_string(), + results, + }), + ) + .into_response() +} + +pub async fn task_build_list(state: &AppState, id: &str) -> Response { + let task_id = match id.parse::() { + Ok(uuid) => uuid, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(json!({"message": "Invalid task ID format"})), + ) + .into_response(); + } + }; + match TaskRepository::get_builds_by_task_id(task_id, &state.conn).await { + Some(build_ids) => { + let build_ids_str: Vec = + build_ids.into_iter().map(|uuid| uuid.to_string()).collect(); + (StatusCode::OK, Json(build_ids_str)).into_response() + } + None => ( + StatusCode::NOT_FOUND, + Json(json!({"message": "Task not found"})), + ) + .into_response(), + } +} + +pub async fn get_orion_clients_info( + state: &AppState, + params: PageParams, +) -> Result>, (StatusCode, Json)> { + let pagination = params.pagination; + let query = params.additional.clone(); + let page = pagination.page.max(1); + let per_page = pagination.per_page.clamp(1u64, 100); + let offset = (page - 1) * per_page; + let mut total: u64 = 0; + let mut items: Vec = Vec::with_capacity(per_page as usize); + + for entry in state.scheduler.workers.iter() { + let matches = query + .hostname + .as_ref() + .is_none_or(|h| entry.value().hostname.contains(h)) + && query + .status + .as_ref() + .is_none_or(|s| entry.value().status.status_type() == *s) + && query.phase.as_ref().is_none_or(|p| { + matches!( + entry.value().status, + WorkerStatus::Busy { phase: Some(ref x), .. } if *x == *p + ) + }); + if matches { + total += 1; + if total > offset && items.len() < per_page as usize { + items.push(OrionClientInfo { + client_id: entry.key().clone(), + hostname: entry.value().hostname.clone(), + orion_version: entry.value().orion_version.clone(), + start_time: entry.value().start_time, + last_heartbeat: entry.value().last_heartbeat, + }); + } + } + } + Ok(Json(CommonPage { total, items })) +} + +pub async fn get_orion_client_status_by_id( + state: &AppState, + id: &str, +) -> Result, (StatusCode, Json)> { + let worker = state.scheduler.workers.get(id).ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json(json!({"message": "Orion client not found"})), + ) + })?; + Ok(Json(OrionClientStatus::from_worker_status(&worker))) +} diff --git a/orion-server/src/service/mod.rs b/orion-server/src/service/mod.rs index 1e2ba0992..b5526cff9 100644 --- a/orion-server/src/service/mod.rs +++ b/orion-server/src/service/mod.rs @@ -1 +1,4 @@ +pub mod api_v2_service; pub mod target_build_status_service; +pub mod target_status_cache_service; +pub mod ws_service; diff --git a/orion-server/src/service/target_build_status_service.rs b/orion-server/src/service/target_build_status_service.rs index 289cb403c..b41e77fb1 100644 --- a/orion-server/src/service/target_build_status_service.rs +++ b/orion-server/src/service/target_build_status_service.rs @@ -6,32 +6,33 @@ use sea_orm::{ }; use uuid::Uuid; +pub struct NewTargetStatusInput { + pub id: Uuid, + pub task_id: Uuid, + pub target_package: String, + pub target_name: String, + pub target_configuration: String, + pub category: String, + pub identifier: String, + pub action: String, + pub status: OrionTargetStatusEnum, +} + pub struct TargetBuildStatusService; impl TargetBuildStatusService { - #[allow(clippy::too_many_arguments)] - pub fn new_active_model( - id: Uuid, - task_id: Uuid, - target_package: String, - target_name: String, - target_configuration: String, - category: String, - identifier: String, - action: String, - status: OrionTargetStatusEnum, - ) -> target_build_status::ActiveModel { + pub fn new_active_model(input: NewTargetStatusInput) -> target_build_status::ActiveModel { let now = Utc::now().into(); target_build_status::ActiveModel { - id: Set(id), - task_id: Set(task_id), - target_package: Set(target_package), - target_name: Set(target_name), - target_configuration: Set(target_configuration), - category: Set(category), - identifier: Set(identifier), - action: Set(action), - status: Set(status), + id: Set(input.id), + task_id: Set(input.task_id), + target_package: Set(input.target_package), + target_name: Set(input.target_name), + target_configuration: Set(input.target_configuration), + category: Set(input.category), + identifier: Set(input.identifier), + action: Set(input.action), + status: Set(input.status), created_at: Set(now), updated_at: Set(now), } diff --git a/orion-server/src/service/target_status_cache_service.rs b/orion-server/src/service/target_status_cache_service.rs new file mode 100644 index 000000000..b205a6c3d --- /dev/null +++ b/orion-server/src/service/target_status_cache_service.rs @@ -0,0 +1,222 @@ +use std::{collections::HashMap, sync::Arc}; + +use api_model::buck2::{types::TargetStatusResponse, ws::WSTargetBuildStatusEvent}; +use callisto::{sea_orm_active_enums::OrionTargetStatusEnum, target_build_status}; +use sea_orm::DatabaseConnection; +use tokio::sync::RwLock; +use uuid::Uuid; + +use crate::service::target_build_status_service::{NewTargetStatusInput, TargetBuildStatusService}; + +#[derive(Hash, Eq, PartialEq, Clone)] +struct ActionKey { + package: String, + name: String, + configuration: String, + category: String, + identifier: String, + action: String, +} + +#[derive(Clone)] +pub struct TargetStatusCache { + /// task_id -> (ActionKey -> ActiveModel) + inner: Arc>>>, +} + +impl TargetStatusCache { + pub fn new() -> Self { + Self { + inner: Arc::new(RwLock::new(HashMap::new())), + } + } + + pub async fn insert_event(&self, event: WSTargetBuildStatusEvent) { + let task_id = match Uuid::parse_str(&event.context.task_id) { + Ok(id) => id, + Err(_) => { + tracing::error!("Invalid task_id: {}", event.context.task_id); + return; + } + }; + + let status = OrionTargetStatusEnum::from_ws_status(&event.target.new_status); + let key = ActionKey { + package: event.target.configured_target_package.clone(), + name: event.target.configured_target_name.clone(), + configuration: event.target.configured_target_configuration.clone(), + category: event.target.category.clone(), + identifier: event.target.identifier.clone(), + action: event.target.action.clone(), + }; + let active_model = build_active_model(task_id, event, status); + + let mut guard = self.inner.write().await; + let task_map = guard.entry(task_id).or_default(); + task_map.insert(key, active_model); + } + + pub async fn flush_all(&self) -> Vec { + let mut guard = self.inner.write().await; + let mut result = Vec::new(); + for (_, action_map) in guard.drain() { + result.extend(action_map.into_values()); + } + result + } + + pub async fn auto_flush_loop( + self, + conn: DatabaseConnection, + mut shutdown: tokio::sync::watch::Receiver, + ) { + let mut ticker = tokio::time::interval(std::time::Duration::from_millis(500)); + loop { + tokio::select! { + _ = ticker.tick() => { + let models = self.flush_all().await; + if models.is_empty() { + continue; + } + if let Err(e) = TargetBuildStatusService::upsert_batch(&conn, models).await { + tracing::error!("Auto flush failed: {:?}", e); + } + } + _ = shutdown.changed() => { + tracing::info!("TargetStatusCache auto flush shutting down"); + let models = self.flush_all().await; + if !models.is_empty() { + let _ = TargetBuildStatusService::upsert_batch(&conn, models).await; + } + break; + } + } + } + } +} + +fn build_active_model( + task_id: Uuid, + event: WSTargetBuildStatusEvent, + status: OrionTargetStatusEnum, +) -> target_build_status::ActiveModel { + TargetBuildStatusService::new_active_model(NewTargetStatusInput { + id: Uuid::new_v4(), + task_id, + target_package: event.target.configured_target_package, + target_name: event.target.configured_target_name, + target_configuration: event.target.configured_target_configuration, + category: event.target.category, + identifier: event.target.identifier, + action: event.target.action, + status, + }) +} + +impl Default for TargetStatusCache { + fn default() -> Self { + Self::new() + } +} + +pub trait FromWsStatus { + fn from_ws_status(status: &str) -> Self; + fn as_str(&self) -> &str; +} + +impl FromWsStatus for OrionTargetStatusEnum { + fn from_ws_status(status: &str) -> Self { + match status.trim().to_ascii_lowercase().as_str() { + "pending" => Self::Pending, + "running" => Self::Running, + "success" | "succeeded" => Self::Success, + "failed" => Self::Failed, + _ => Self::Pending, + } + } + + fn as_str(&self) -> &str { + match self { + Self::Pending => "PENDING", + Self::Running => "RUNNING", + Self::Success => "SUCCESS", + Self::Failed => "FAILED", + } + } +} + +pub async fn targets_status_by_task_id( + conn: &DatabaseConnection, + task_id: &str, +) -> Result, (axum::http::StatusCode, String)> { + let task_uuid = Uuid::parse_str(task_id).map_err(|_| { + ( + axum::http::StatusCode::BAD_REQUEST, + "Invalid task_id".to_string(), + ) + })?; + let targets = TargetBuildStatusService::fetch_by_task_id(conn, task_uuid) + .await + .map_err(|_| { + ( + axum::http::StatusCode::INTERNAL_SERVER_ERROR, + "Database error".to_string(), + ) + })?; + if targets.is_empty() { + return Err(( + axum::http::StatusCode::NOT_FOUND, + "No target status found".to_string(), + )); + } + Ok(targets + .into_iter() + .map(|t| TargetStatusResponse { + id: t.id.to_string(), + task_id: t.task_id.to_string(), + package: t.target_package, + name: t.target_name, + configuration: t.target_configuration, + category: t.category, + identifier: t.identifier, + action: t.action, + status: t.status.as_str().to_owned(), + }) + .collect()) +} + +pub async fn target_status_by_id( + conn: &DatabaseConnection, + target_id: &str, +) -> Result { + let target_uuid = Uuid::parse_str(target_id).map_err(|e| { + ( + axum::http::StatusCode::BAD_REQUEST, + format!("Invalid target_id '{}': {}", target_id, e), + ) + })?; + let target = TargetBuildStatusService::find_by_id(conn, target_uuid) + .await + .map_err(|e| { + ( + axum::http::StatusCode::INTERNAL_SERVER_ERROR, + format!("Database query failed: {}", e), + ) + })? + .ok_or(( + axum::http::StatusCode::NOT_FOUND, + "Target not found".to_string(), + ))?; + + Ok(TargetStatusResponse { + id: target.id.to_string(), + task_id: target.task_id.to_string(), + package: target.target_package, + name: target.target_name, + configuration: target.target_configuration, + category: target.category, + identifier: target.identifier, + action: target.action, + status: target.status.as_str().to_string(), + }) +} diff --git a/orion-server/src/service/ws_service.rs b/orion-server/src/service/ws_service.rs new file mode 100644 index 000000000..8fc445f83 --- /dev/null +++ b/orion-server/src/service/ws_service.rs @@ -0,0 +1,332 @@ +use std::{net::SocketAddr, ops::ControlFlow}; + +use api_model::buck2::{types::LogEvent, ws::WSMessage}; +use axum::{ + extract::{ + ConnectInfo, State, WebSocketUpgrade, + ws::{Message, Utf8Bytes, WebSocket}, + }, + response::IntoResponse, +}; +use futures_util::{SinkExt, StreamExt}; +use sea_orm::{ActiveValue::Set, ColumnTrait, EntityTrait, QueryFilter as _}; +use tokio::sync::mpsc::{self, UnboundedSender}; +use uuid::Uuid; + +use crate::{ + app_state::AppState, + auto_retry::AutoRetryJudger, + entity::{builds, targets::TargetState}, + log::log_service::LogService, + repository::{build_events::BuildEvent, targets::TargetRepository}, + scheduler::{WorkerInfo, WorkerStatus}, +}; + +const RETRY_COUNT_MAX: i32 = 3; + +pub async fn ws_handler( + ws: WebSocketUpgrade, + ConnectInfo(addr): ConnectInfo, + State(state): State, +) -> impl IntoResponse { + tracing::info!("{addr} connected. Waiting for registration..."); + ws.on_upgrade(move |socket| handle_socket(socket, addr, state)) +} + +async fn handle_socket(socket: WebSocket, who: SocketAddr, state: AppState) { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let mut worker_id: Option = None; + let (mut sender, mut receiver) = socket.split(); + + let send_task = tokio::spawn(async move { + while let Some(msg) = rx.recv().await { + let msg_str = serde_json::to_string(&msg).unwrap_or_default(); + if sender + .send(Message::Text(Utf8Bytes::from(msg_str))) + .await + .is_err() + { + tracing::warn!("Failed to send message to {who}, client disconnected."); + break; + } + } + }); + + let state_clone = state.clone(); + let tx_clone = tx.clone(); + let recv_task = tokio::spawn(async move { + let mut worker_id_inner: Option = None; + while let Some(Ok(msg)) = receiver.next().await { + if process_message(msg, who, &state_clone, &mut worker_id_inner, &tx_clone) + .await + .is_break() + { + break; + } + } + worker_id_inner + }); + + tokio::select! { + _ = send_task => {}, + result = recv_task => { + if let Ok(final_worker_id) = result { + worker_id = final_worker_id; + } + } + } + + if let Some(id) = &worker_id { + tracing::info!("Cleaning up for worker: {id} from {who}."); + state.scheduler.workers.remove(id); + } else { + tracing::info!("Cleaning up unregistered connection from {who}."); + } + tracing::info!("Websocket context {who} destroyed"); +} + +async fn process_message( + msg: Message, + who: SocketAddr, + state: &AppState, + worker_id: &mut Option, + tx: &UnboundedSender, +) -> ControlFlow<(), ()> { + match msg { + Message::Text(t) => { + let ws_msg: Result = serde_json::from_str(&t); + if let Err(e) = ws_msg { + tracing::warn!("Failed to parse message from {who}: {e}"); + return ControlFlow::Continue(()); + } + let ws_msg = ws_msg.unwrap_or(WSMessage::Heartbeat); + + if worker_id.is_none() { + if let WSMessage::Register { + id, + hostname, + orion_version, + } = ws_msg + { + tracing::info!("Worker from {who} registered as: {id}"); + state.scheduler.workers.insert( + id.clone(), + WorkerInfo { + sender: tx.clone(), + status: WorkerStatus::Idle, + last_heartbeat: chrono::Utc::now(), + start_time: chrono::Utc::now(), + hostname, + orion_version, + }, + ); + *worker_id = Some(id); + state.scheduler.notify_task_available(); + } else { + tracing::error!( + "First message from {who} was not Register. Closing connection." + ); + return ControlFlow::Break(()); + } + return ControlFlow::Continue(()); + } + + let current_worker_id = worker_id.as_ref().unwrap_or(&String::new()).clone(); + match ws_msg { + WSMessage::Register { .. } => { + tracing::warn!( + "Worker {current_worker_id} sent Register message again. Ignoring." + ); + } + WSMessage::Heartbeat => { + if let Some(mut worker) = state.scheduler.workers.get_mut(¤t_worker_id) { + worker.last_heartbeat = chrono::Utc::now(); + if let WorkerStatus::Error(_) = worker.status { + worker.status = WorkerStatus::Idle; + } + } + } + WSMessage::TaskBuildOutput { build_id, output } => { + if let Some(build_info) = state.scheduler.active_builds.get(&build_id) { + let log_event = LogEvent { + task_id: build_info.event_payload.task_id.to_string(), + repo_name: LogService::last_segment(&build_info.event_payload.repo) + .to_string(), + build_id: build_id.clone(), + line: output.clone(), + is_end: false, + }; + state.log_service.publish(log_event); + } + if let Some(mut build_info) = state.scheduler.active_builds.get_mut(&build_id) { + build_info.auto_retry_judger.judge_by_output(&output); + } + } + WSMessage::TaskBuildCompleteV2 { + build_id, + success, + exit_code, + message, + } + | WSMessage::TaskBuildComplete { + build_id, + success, + exit_code, + message, + } => { + let ( + mut auto_retry_judger, + mut retry_count, + repo, + changes, + cl_link, + task_id, + target_id, + ) = if let Some(build_info) = state.scheduler.active_builds.get(&build_id) { + ( + build_info.auto_retry_judger.clone(), + build_info.event_payload.retry_count, + build_info.event_payload.repo.clone(), + build_info.changes.clone(), + build_info.event_payload.cl_link.clone(), + build_info.event_payload.task_id, + build_info.target_id, + ) + } else { + return ControlFlow::Continue(()); + }; + + auto_retry_judger.judge_by_exit_code(exit_code.unwrap_or(0)); + if auto_retry_judger.get_can_auto_retry() && retry_count < RETRY_COUNT_MAX { + retry_count += 1; + if let Some(mut build_info) = + state.scheduler.active_builds.get_mut(&build_id) + { + build_info.event_payload.retry_count = retry_count; + build_info.auto_retry_judger = AutoRetryJudger::new(); + } + let _ = builds::Entity::update_many() + .set(builds::ActiveModel { + retry_count: Set(retry_count), + ..Default::default() + }) + .filter( + builds::Column::Id.eq(build_id + .parse::() + .unwrap_or_else(|_| Uuid::nil())), + ) + .exec(&state.conn) + .await; + + let msg = WSMessage::TaskBuild { + build_id: build_id.clone(), + repo: repo.clone(), + cl_link, + changes, + }; + if let Some(worker) = state.scheduler.workers.get_mut(¤t_worker_id) + && worker.sender.send(msg).is_ok() + { + return ControlFlow::Continue(()); + } + } + + state.log_service.publish(LogEvent { + task_id: task_id.to_string(), + repo_name: LogService::last_segment(&repo).to_string(), + build_id: build_id.to_string(), + line: String::new(), + is_end: true, + }); + state.scheduler.active_builds.remove(&build_id); + + let _ = + BuildEvent::update_build_complete_result(&build_id, exit_code, &state.conn) + .await; + + let target_state = match (success, exit_code) { + (true, Some(0)) => TargetState::Completed, + (_, None) => TargetState::Interrupted, + _ => TargetState::Failed, + }; + let error_summary = if matches!(target_state, TargetState::Failed) { + match state + .log_service + .read_full_log( + &task_id.to_string(), + &LogService::last_segment(&repo), + &build_id.to_string(), + ) + .await + { + Ok(content) => find_caused_by_next_line_in_content(&content).await, + Err(_) => None, + } + } else { + None + }; + let _ = TargetRepository::update_state( + &state.conn, + target_id, + target_state, + None, + Some(chrono::Utc::now().with_timezone( + &chrono::FixedOffset::east_opt(0).unwrap_or_else(|| unreachable!()), + )), + error_summary, + ) + .await; + + if let Some(mut worker) = state.scheduler.workers.get_mut(¤t_worker_id) { + worker.status = if success { + WorkerStatus::Idle + } else { + WorkerStatus::Error(message) + }; + } + state.scheduler.notify_task_available(); + } + WSMessage::TaskPhaseUpdate { build_id, phase } => { + if let Some(mut worker) = state.scheduler.workers.get_mut(¤t_worker_id) + && let WorkerStatus::Busy { build_id: id, .. } = &worker.status + && &build_id == id + { + worker.status = WorkerStatus::Busy { + build_id, + phase: Some(phase), + }; + } + } + WSMessage::TargetBuildStatusBatch { events } => { + for update in events { + state.target_status_cache.insert_event(update).await; + } + } + _ => {} + } + } + Message::Close(_) => { + if let Some(id) = worker_id.take() + && let Some(mut worker) = state.scheduler.workers.get_mut(&id) + { + worker.status = WorkerStatus::Lost; + } + return ControlFlow::Break(()); + } + _ => {} + } + ControlFlow::Continue(()) +} + +async fn find_caused_by_next_line_in_content(content: &str) -> Option { + let mut last_was_caused = false; + for line in content.lines() { + if last_was_caused { + return Some(line.to_string()); + } + if line.trim() == "Caused by:" { + last_was_caused = true; + } + } + None +} diff --git a/saturn/src/context.rs b/saturn/src/context.rs index 6724b3274..87fb6c137 100644 --- a/saturn/src/context.rs +++ b/saturn/src/context.rs @@ -14,7 +14,6 @@ pub struct CedarContext { schema: Schema, } -#[allow(dead_code)] #[derive(Debug, Error)] pub enum ContextError { #[error("{0}")] @@ -33,7 +32,6 @@ pub enum ContextError { Json(#[from] serde_json::Error), } -#[allow(dead_code)] #[derive(Debug, Error)] pub enum SaturnContextError { #[error("Authorization Denied")] diff --git a/saturn/src/entitystore.rs b/saturn/src/entitystore.rs index b9066eca4..72e513fb7 100644 --- a/saturn/src/entitystore.rs +++ b/saturn/src/entitystore.rs @@ -30,7 +30,6 @@ impl EntityStore { } } - #[allow(dead_code)] pub fn as_entities(&self, schema: &Schema) -> Entities { let users = self.users.values().map(|user| user.clone().into()); let repos = self.repos.values().map(|repo| repo.clone().into());