//! Worker Service - Business logic layer for worker operations //! //! This module provides a clean separation between HTTP concerns (in routers) //! and business logic for worker management. The service orchestrates //! WorkerRegistry and JobQueue operations. use std::sync::Arc; use axum::{ http::StatusCode, response::{IntoResponse, Response}, Json, }; use serde_json::json; use tracing::warn; use crate::{ config::RouterConfig, core::{worker::worker_to_info, worker_registry::WorkerId, Job, JobQueue, WorkerRegistry}, protocols::worker_spec::{ WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest, }, }; /// Error types for worker service operations #[derive(Debug)] pub enum WorkerServiceError { /// Worker with given ID was not found NotFound { worker_id: String }, /// Invalid worker ID format (expected UUID) InvalidId { raw: String, message: String }, /// Job queue not initialized QueueNotInitialized, /// Failed to submit job to queue QueueSubmitFailed { message: String }, } impl WorkerServiceError { pub fn error_code(&self) -> &'static str { match self { Self::NotFound { .. } => "WORKER_NOT_FOUND", Self::InvalidId { .. } => "BAD_REQUEST", Self::QueueNotInitialized => "INTERNAL_SERVER_ERROR", Self::QueueSubmitFailed { .. } => "INTERNAL_SERVER_ERROR", } } pub fn status_code(&self) -> StatusCode { match self { Self::NotFound { .. } => StatusCode::NOT_FOUND, Self::InvalidId { .. } => StatusCode::BAD_REQUEST, Self::QueueNotInitialized => StatusCode::INTERNAL_SERVER_ERROR, Self::QueueSubmitFailed { .. } => StatusCode::INTERNAL_SERVER_ERROR, } } } impl std::fmt::Display for WorkerServiceError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::NotFound { worker_id } => write!(f, "Worker {} not found", worker_id), Self::InvalidId { raw, message } => { write!( f, "Invalid worker_id '{}' (expected UUID). Error: {}", raw, message ) } Self::QueueNotInitialized => write!(f, "Job queue not initialized"), Self::QueueSubmitFailed { message } => write!(f, "{}", message), } } } impl std::error::Error for WorkerServiceError {} impl IntoResponse for WorkerServiceError { fn into_response(self) -> Response { let error = WorkerErrorResponse { error: self.to_string(), code: self.error_code().to_string(), }; (self.status_code(), Json(error)).into_response() } } /// Result of creating a worker (async job submission) #[derive(Debug)] pub struct CreateWorkerResult { pub worker_id: WorkerId, pub url: String, pub location: String, } impl IntoResponse for CreateWorkerResult { fn into_response(self) -> Response { let response = json!({ "status": "accepted", "worker_id": self.worker_id.as_str(), "url": self.url, "location": self.location, "message": "Worker addition queued for background processing" }); ( StatusCode::ACCEPTED, [(http::header::LOCATION, self.location)], Json(response), ) .into_response() } } /// Result of deleting a worker (async job submission) #[derive(Debug)] pub struct DeleteWorkerResult { pub worker_id: WorkerId, pub url: String, } impl IntoResponse for DeleteWorkerResult { fn into_response(self) -> Response { let response = json!({ "status": "accepted", "worker_id": self.worker_id.as_str(), "message": "Worker removal queued for background processing" }); (StatusCode::ACCEPTED, Json(response)).into_response() } } /// Result of updating a worker (async job submission) #[derive(Debug)] pub struct UpdateWorkerResult { pub worker_id: WorkerId, pub url: String, } impl IntoResponse for UpdateWorkerResult { fn into_response(self) -> Response { let response = json!({ "status": "accepted", "worker_id": self.worker_id.as_str(), "message": "Worker update queued for background processing" }); (StatusCode::ACCEPTED, Json(response)).into_response() } } /// Result of listing workers #[derive(Debug)] pub struct ListWorkersResult { pub workers: Vec, pub total: usize, pub prefill_count: usize, pub decode_count: usize, pub regular_count: usize, } impl IntoResponse for ListWorkersResult { fn into_response(self) -> Response { let response = json!({ "workers": self.workers, "total": self.total, "stats": { "prefill_count": self.prefill_count, "decode_count": self.decode_count, "regular_count": self.regular_count, } }); Json(response).into_response() } } /// Wrapper for WorkerInfo to implement IntoResponse pub struct GetWorkerResponse(pub WorkerInfo); impl IntoResponse for GetWorkerResponse { fn into_response(self) -> Response { Json(self.0).into_response() } } /// Worker Service - Orchestrates worker business logic /// /// This service provides a clean API for worker operations, separating /// business logic from HTTP concerns. Handlers in server.rs become thin /// wrappers that translate between HTTP and this service. pub struct WorkerService { worker_registry: Arc, job_queue: Arc>>, router_config: RouterConfig, } impl WorkerService { /// Create a new WorkerService pub fn new( worker_registry: Arc, job_queue: Arc>>, router_config: RouterConfig, ) -> Self { Self { worker_registry, job_queue, router_config, } } /// Parse and validate a worker ID string pub fn parse_worker_id(&self, raw: &str) -> Result { uuid::Uuid::parse_str(raw) .map(|_| WorkerId::from_string(raw.to_string())) .map_err(|e| WorkerServiceError::InvalidId { raw: raw.to_string(), message: e.to_string(), }) } /// Get the job queue, returning an error if not initialized fn get_job_queue(&self) -> Result<&Arc, WorkerServiceError> { self.job_queue .get() .ok_or(WorkerServiceError::QueueNotInitialized) } pub async fn create_worker( &self, mut config: WorkerConfigRequest, ) -> Result { if self.router_config.api_key.is_some() && config.api_key.is_none() { warn!( "Adding worker {} without API key while router has API key configured. \ Worker will be accessible without authentication. \ If the worker requires the same API key as the router, please specify it explicitly.", config.url ); } config.dp_aware = self.router_config.dp_aware; let worker_url = config.url.clone(); let worker_id = self.worker_registry.reserve_id_for_url(&worker_url); let job = Job::AddWorker { config: Box::new(config), }; self.get_job_queue()? .submit(job) .await .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?; let location = format!("/workers/{}", worker_id.as_str()); Ok(CreateWorkerResult { worker_id, url: worker_url, location, }) } /// List all workers with their info pub fn list_workers(&self) -> ListWorkersResult { let workers = self.worker_registry.get_all_with_ids(); let worker_infos: Vec = workers .iter() .map(|(worker_id, worker)| { let mut info = worker_to_info(worker); info.id = worker_id.as_str().to_string(); info }) .collect(); let stats = self.worker_registry.stats(); ListWorkersResult { workers: worker_infos, total: stats.total_workers, prefill_count: stats.prefill_workers, decode_count: stats.decode_workers, regular_count: stats.regular_workers, } } pub fn get_worker(&self, worker_id_raw: &str) -> Result { let worker_id = self.parse_worker_id(worker_id_raw)?; let job_queue = self.get_job_queue()?; if let Some(worker) = self.worker_registry.get(&worker_id) { let worker_url = worker.url().to_string(); let mut worker_info = worker_to_info(&worker); worker_info.id = worker_id.as_str().to_string(); if let Some(status) = job_queue.get_status(&worker_url) { worker_info.job_status = Some(status); } return Ok(GetWorkerResponse(worker_info)); } if let Some(worker_url) = self.worker_registry.get_url_by_id(&worker_id) { if let Some(status) = job_queue.get_status(&worker_url) { return Ok(GetWorkerResponse(WorkerInfo::pending( worker_id.as_str(), worker_url, Some(status), ))); } } Err(WorkerServiceError::NotFound { worker_id: worker_id_raw.to_string(), }) } /// Delete a worker by ID (submits async job) pub async fn delete_worker( &self, worker_id_raw: &str, ) -> Result { let worker_id = self.parse_worker_id(worker_id_raw)?; let url = self .worker_registry .get_url_by_id(&worker_id) .ok_or_else(|| WorkerServiceError::NotFound { worker_id: worker_id_raw.to_string(), })?; let job = Job::RemoveWorker { url: url.clone() }; let job_queue = self.get_job_queue()?; job_queue .submit(job) .await .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?; Ok(DeleteWorkerResult { worker_id, url }) } /// Update a worker by ID (submits async job) pub async fn update_worker( &self, worker_id_raw: &str, update: WorkerUpdateRequest, ) -> Result { let worker_id = self.parse_worker_id(worker_id_raw)?; let url = self .worker_registry .get_url_by_id(&worker_id) .ok_or_else(|| WorkerServiceError::NotFound { worker_id: worker_id_raw.to_string(), })?; let job = Job::UpdateWorker { url: url.clone(), update: Box::new(update), }; let job_queue = self.get_job_queue()?; job_queue .submit(job) .await .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?; Ok(UpdateWorkerResult { worker_id, url }) } }