diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index bed395518..9bf90999e 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -34,7 +34,7 @@ required-features = ["build-binary"] [features] build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"] -default = ["build-binary", "substrait"] +default = ["build-binary", "substrait", "rest-api"] # job info can cache stage plans, in some cases where # task plans can be re-computed, cache behavior may need to be disabled. disable-stage-plan-cache = [] diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index a7ade706f..292e829e6 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -10,6 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::scheduler_process::SchedulerErrorResponse; use crate::scheduler_server::SchedulerServer; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::state::execution_graph::ExecutionStage; @@ -122,15 +123,13 @@ pub async fn get_jobs< U: AsExecutionPlan + Send + Sync + 'static, >( State(data_server): State>>, -) -> Result { - // TODO: Display last seen information in UI +) -> Result { let state = &data_server.state; - let jobs = state - .task_manager - .get_jobs() - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let jobs = + state.task_manager.get_jobs().await.map_err(|_| { + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) + })?; let jobs: Vec = jobs .iter() @@ -166,17 +165,17 @@ pub async fn get_job< >( State(data_server): State>>, Path(job_id): Path, -) -> Result { +) -> Result { let graph = data_server .state .task_manager .get_job_execution_graph(&job_id) .await .map_err(|err| { - tracing::error!("Error occurred while getting the execution graph for job '{job_id}': {err:?}"); - StatusCode::INTERNAL_SERVER_ERROR + tracing::error!("Error occurred while getting the execution graph for job '{job_id}' reason: {err:?}"); + SchedulerErrorResponse::with_error(StatusCode::INTERNAL_SERVER_ERROR, format!("Error occurred while getting the execution graph for job '{job_id}' reason: {}", err)) })? - .ok_or(StatusCode::NOT_FOUND)?; + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; let stage_plan = format!("{:?}", graph); let job = graph.as_ref(); let (plain_status, job_status) = @@ -207,7 +206,7 @@ pub async fn cancel_job< >( State(data_server): State>>, Path(job_id): Path, -) -> Result { +) -> Result { // 404 if the job doesn't exist let job_status = data_server .state @@ -216,9 +215,12 @@ pub async fn cancel_job< .await .map_err(|err| { tracing::error!("Error getting job status: {err:?}"); - StatusCode::INTERNAL_SERVER_ERROR + SchedulerErrorResponse::with_error( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Error getting job status: {}", err), + ) })? - .ok_or(StatusCode::NOT_FOUND)?; + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; match &job_status.status { None | Some(Status::Queued(_)) | Some(Status::Running(_)) => { @@ -229,11 +231,13 @@ pub async fn cancel_job< tracing::error!( "Error getting query stage event loop sender: {err:?}" ); - StatusCode::INTERNAL_SERVER_ERROR + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) })? .post_event(QueryStageSchedulerEvent::JobCancel(job_id)) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + .map_err(|_| { + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) + })?; Ok(( StatusCode::OK, @@ -274,13 +278,18 @@ pub async fn get_query_stages< >( State(data_server): State>>, Path(job_id): Path, -) -> Result { +) -> Result { if let Some(graph) = data_server .state .task_manager .get_job_execution_graph(&job_id) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .map_err(|e| { + SchedulerErrorResponse::with_error( + StatusCode::INTERNAL_SERVER_ERROR, + e.to_string(), + ) + })? { let stages = graph .as_ref() @@ -409,16 +418,16 @@ pub async fn get_job_dot_graph< >( State(data_server): State>>, Path(job_id): Path, -) -> Result { +) -> Result { if let Some(graph) = data_server .state .task_manager .get_job_execution_graph(&job_id) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))? { ExecutionGraphDot::generate(graph.as_ref()) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) + .map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)) } else { Ok("Not Found".to_string()) } @@ -430,16 +439,16 @@ pub async fn get_query_stage_dot_graph< >( State(data_server): State>>, Path((job_id, stage_id)): Path<(String, usize)>, -) -> Result { +) -> Result { if let Some(graph) = data_server .state .task_manager .get_job_execution_graph(&job_id) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))? { ExecutionGraphDot::generate_for_query_stage(graph.as_ref(), stage_id) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) + .map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)) } else { Ok("Not Found".to_string()) } @@ -451,7 +460,7 @@ pub async fn get_job_svg_graph< >( State(data_server): State>>, Path(job_id): Path, -) -> Result { +) -> Result { let dot = get_job_dot_graph(State(data_server.clone()), Path(job_id)).await?; match graphviz_rust::parse(&dot) { Ok(graph) => { @@ -460,7 +469,9 @@ pub async fn get_job_svg_graph< &mut PrinterContext::default(), vec![CommandArg::Format(Format::Svg)], ) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + .map_err(|_| { + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) + })?; let svg = String::from_utf8_lossy(&result).to_string(); Ok(Response::builder() diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index 2662e3eea..feeb2f583 100644 --- a/ballista/scheduler/src/api/mod.rs +++ b/ballista/scheduler/src/api/mod.rs @@ -27,6 +27,7 @@ pub fn get_routes< ) -> Router { let router = Router::new() .route("/api/state", get(handlers::get_scheduler_state::)) + .route("/api/version", get(handlers::get_scheduler_state::)) .route("/api/executors", get(handlers::get_executors::)) .route("/api/jobs", get(handlers::get_jobs::)) .route( diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 4f520aff5..84bf6cf06 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -190,6 +190,11 @@ pub struct Config { help = "The interval to check expired or dead executors" )] pub expire_dead_executor_interval_seconds: u64, + + #[cfg(feature = "rest-api")] + /// Should the rest api be disabled + #[arg(long, default_value_t = false, help = "Should the rest api be disable")] + pub disable_rest: bool, } /// Configurations for the ballista scheduler of scheduling jobs and tasks @@ -245,6 +250,9 @@ pub struct SchedulerConfig { pub override_create_grpc_client_endpoint: Option, /// Whether to use TLS when connecting to executors (for flight proxy) pub use_tls: bool, + #[cfg(feature = "rest-api")] + /// Should the rest api be disabled + pub disable_rest: bool, } impl Default for SchedulerConfig { @@ -273,6 +281,8 @@ impl Default for SchedulerConfig { override_physical_codec: None, override_create_grpc_client_endpoint: None, use_tls: false, + #[cfg(feature = "rest-api")] + disable_rest: false, } } } @@ -520,6 +530,8 @@ impl TryFrom for SchedulerConfig { override_session_builder: None, override_create_grpc_client_endpoint: None, use_tls: false, + #[cfg(feature = "rest-api")] + disable_rest: opt.disable_rest, }; Ok(config) diff --git a/ballista/scheduler/src/scheduler_process.rs b/ballista/scheduler/src/scheduler_process.rs index 3e6a97b2f..5e87f559b 100644 --- a/ballista/scheduler/src/scheduler_process.rs +++ b/ballista/scheduler/src/scheduler_process.rs @@ -18,6 +18,8 @@ use crate::flight_proxy_service::BallistaFlightProxyService; use arrow_flight::flight_service_server::FlightServiceServer; +use axum::Json; +use axum::response::{IntoResponse, Response}; use ballista_core::BALLISTA_VERSION; use ballista_core::error::BallistaError; use ballista_core::extension::BallistaConfigGrpcEndpoint; @@ -130,17 +132,42 @@ pub async fn start_grpc_service< tonic_builder.add_service(ExternalScalerServer::new(scheduler.clone())); let tonic = tonic_builder.routes().into_axum_router(); - let tonic = tonic.fallback(|| async { (StatusCode::NOT_FOUND, "404 - Not Found") }); + + // registering default handler for unmatched requests + let tonic = + tonic.fallback(|| async { SchedulerErrorResponse::new(StatusCode::NOT_FOUND) }); #[cfg(feature = "rest-api")] - let axum = get_routes(Arc::new(scheduler)); - #[cfg(feature = "rest-api")] - let final_route = axum - .merge(tonic) - .into_make_service_with_connect_info::(); + let final_route = if config.disable_rest { + tonic + .route( + "/api/{*path}", + axum::routing::any(|| async { + SchedulerErrorResponse::with_error( + StatusCode::NOT_FOUND, + "Rest api has been disabled at startup".to_string(), + ) + }), + ) + .into_make_service_with_connect_info::() + } else { + let axum = get_routes(Arc::new(scheduler)); + axum.merge(tonic) + .into_make_service_with_connect_info::() + }; #[cfg(not(feature = "rest-api"))] - let final_route = tonic.into_make_service_with_connect_info::(); + let final_route = tonic + .route( + "/api/{*path}", + axum::routing::any(|| async { + SchedulerErrorResponse::with_error( + StatusCode::NOT_FOUND, + "Rest api has been disabled at compile time".to_string(), + ) + }), + ) + .into_make_service_with_connect_info::(); let listener = tokio::net::TcpListener::bind(&address) .await @@ -165,3 +192,39 @@ pub async fn start_server( start_grpc_service(address, scheduler).await } + +#[derive(Debug, serde::Serialize)] +pub(crate) struct SchedulerErrorResponse { + #[serde(skip)] + status_code: StatusCode, + http_code: u16, + reason: Option<&'static str>, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +impl SchedulerErrorResponse { + pub(crate) fn new(status_code: StatusCode) -> Self { + Self { + status_code, + reason: status_code.canonical_reason(), + http_code: status_code.as_u16(), + error: None, + } + } + pub(crate) fn with_error(status_code: StatusCode, error: String) -> Self { + Self { + status_code, + reason: status_code.canonical_reason(), + http_code: status_code.as_u16(), + error: Some(error), + } + } +} + +impl IntoResponse for SchedulerErrorResponse { + fn into_response(self) -> Response { + let status = self.status_code; + (status, Json(self)).into_response() + } +}