Line data    Source code

       1              : use crate::compute_prewarm::LfcPrewarmStateWithProgress;
       2              : use crate::http::JsonResponse;
       3              : use axum::response::{IntoResponse, Response};
       4              : use axum::{Json, http::StatusCode};
       5              : use axum_extra::extract::OptionalQuery;
       6              : use compute_api::responses::LfcOffloadState;
       7              : type Compute = axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>;
       8              : 
       9            0 : pub(in crate::http) async fn prewarm_state(compute: Compute) -> Json<LfcPrewarmStateWithProgress> {
      10            0 :     Json(compute.lfc_prewarm_state().await)
      11            0 : }
      12              : 
      13              : // Following functions are marked async for axum, as it's more convenient than wrapping these
      14              : // in async lambdas at call site
      15              : 
      16            0 : pub(in crate::http) async fn offload_state(compute: Compute) -> Json<LfcOffloadState> {
      17            0 :     Json(compute.lfc_offload_state())
      18            0 : }
      19              : 
      20            0 : #[derive(serde::Deserialize)]
      21              : pub struct PrewarmQuery {
      22              :     pub from_endpoint: String,
      23              : }
      24              : 
      25            0 : pub(in crate::http) async fn prewarm(
      26            0 :     compute: Compute,
      27            0 :     OptionalQuery(query): OptionalQuery<PrewarmQuery>,
      28            0 : ) -> Response {
      29            0 :     if compute.prewarm_lfc(query.map(|q| q.from_endpoint)) {
      30            0 :         StatusCode::ACCEPTED.into_response()
      31              :     } else {
      32            0 :         JsonResponse::error(
      33              :             StatusCode::TOO_MANY_REQUESTS,
      34              :             "Multiple requests for prewarm are not allowed",
      35              :         )
      36              :     }
      37            0 : }
      38              : 
      39            0 : pub(in crate::http) async fn offload(compute: Compute) -> Response {
      40            0 :     if compute.offload_lfc() {
      41            0 :         StatusCode::ACCEPTED.into_response()
      42              :     } else {
      43            0 :         JsonResponse::error(
      44              :             StatusCode::TOO_MANY_REQUESTS,
      45              :             "Multiple requests for prewarm offload are not allowed",
      46              :         )
      47              :     }
      48            0 : }