Line data    Source code

       1              : use crate::compute::{ComputeNode, forward_termination_signal};
       2              : use crate::http::JsonResponse;
       3              : use axum::extract::State;
       4              : use axum::response::{IntoResponse, Response};
       5              : use axum_extra::extract::OptionalQuery;
       6              : use compute_api::responses::{ComputeStatus, TerminateMode, TerminateResponse};
       7              : use http::StatusCode;
       8              : use serde::Deserialize;
       9              : use std::sync::Arc;
      10              : use tokio::task;
      11              : use tracing::info;
      12              : 
      13            0 : #[derive(Deserialize, Default)]
      14              : pub struct TerminateQuery {
      15              :     mode: TerminateMode,
      16              : }
      17              : 
      18              : /// Terminate the compute.
      19            0 : pub(in crate::http) async fn terminate(
      20            0 :     State(compute): State<Arc<ComputeNode>>,
      21            0 :     OptionalQuery(terminate): OptionalQuery<TerminateQuery>,
      22            0 : ) -> Response {
      23            0 :     let mode = terminate.unwrap_or_default().mode;
      24              :     {
      25            0 :         let mut state = compute.state.lock().unwrap();
      26            0 :         if state.status == ComputeStatus::Terminated {
      27            0 :             let response = TerminateResponse {
      28            0 :                 lsn: state.terminate_flush_lsn,
      29            0 :             };
      30            0 :             return JsonResponse::success(StatusCode::CREATED, response);
      31            0 :         }
      32              : 
      33            0 :         if !matches!(state.status, ComputeStatus::Empty | ComputeStatus::Running) {
      34            0 :             return JsonResponse::invalid_status(state.status);
      35            0 :         }
      36              : 
      37              :         // If compute is Empty, there's no Postgres to terminate. The regular compute_ctl termination path
      38              :         // assumes Postgres to be configured and running, so we just special-handle this case by exiting
      39              :         // the process directly.
      40            0 :         if compute.params.lakebase_mode && state.status == ComputeStatus::Empty {
      41            0 :             drop(state);
      42            0 :             info!("terminating empty compute - will exit process");
      43              : 
      44              :             // Queue a task to exit the process after 5 seconds. The 5-second delay aims to
      45              :             // give enough time for the HTTP response to be sent so that HCM doesn't get an abrupt
      46              :             // connection termination.
      47            0 :             tokio::spawn(async {
      48            0 :                 tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
      49            0 :                 info!("exiting process after terminating empty compute");
      50            0 :                 std::process::exit(0);
      51              :             });
      52              : 
      53            0 :             return StatusCode::OK.into_response();
      54            0 :         }
      55              : 
      56              :         // For Running status, proceed with normal termination
      57            0 :         state.set_status(mode.into(), &compute.state_changed);
      58            0 :         drop(state);
      59              :     }
      60              : 
      61            0 :     forward_termination_signal(false);
      62            0 :     info!("sent signal and notified waiters");
      63              : 
      64              :     // Spawn a blocking thread to wait for compute to become Terminated.
      65              :     // This is needed to do not block the main pool of workers and
      66              :     // be able to serve other requests while some particular request
      67              :     // is waiting for compute to finish configuration.
      68            0 :     let c = compute.clone();
      69            0 :     let lsn = task::spawn_blocking(move || {
      70            0 :         let mut state = c.state.lock().unwrap();
      71            0 :         while state.status != ComputeStatus::Terminated {
      72            0 :             state = c.state_changed.wait(state).unwrap();
      73            0 :             info!(
      74            0 :                 "waiting for compute to become {}, current status: {:?}",
      75              :                 ComputeStatus::Terminated,
      76            0 :                 state.status
      77              :             );
      78              :         }
      79            0 :         state.terminate_flush_lsn
      80            0 :     })
      81            0 :     .await
      82            0 :     .unwrap();
      83            0 :     info!("terminated Postgres");
      84            0 :     JsonResponse::success(StatusCode::OK, TerminateResponse { lsn })
      85            0 : }