Line data    Source code

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