LCOV - code coverage report
Current view: top level - compute_tools/src - compute.rs (source / functions) Coverage Total Hit
Test: bd421c520964a8037e429c13001c8b5f00fa42ec.info Lines: 5.7 % 1262 72
Test Date: 2025-07-24 16:21:20 Functions: 7.9 % 114 9

            Line data    Source code
       1              : use anyhow::{Context, Result};
       2              : use chrono::{DateTime, Utc};
       3              : use compute_api::privilege::Privilege;
       4              : use compute_api::responses::{
       5              :     ComputeConfig, ComputeCtlConfig, ComputeMetrics, ComputeStatus, LfcOffloadState,
       6              :     LfcPrewarmState, PromoteState, TlsConfig,
       7              : };
       8              : use compute_api::spec::{
       9              :     ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PageserverProtocol, PgIdent,
      10              : };
      11              : use futures::StreamExt;
      12              : use futures::future::join_all;
      13              : use futures::stream::FuturesUnordered;
      14              : use itertools::Itertools;
      15              : use nix::sys::signal::{Signal, kill};
      16              : use nix::unistd::Pid;
      17              : use once_cell::sync::Lazy;
      18              : use pageserver_page_api::{self as page_api, BaseBackupCompression};
      19              : use postgres;
      20              : use postgres::NoTls;
      21              : use postgres::error::SqlState;
      22              : use remote_storage::{DownloadError, RemotePath};
      23              : use std::collections::{HashMap, HashSet};
      24              : use std::ffi::OsString;
      25              : use std::os::unix::fs::{PermissionsExt, symlink};
      26              : use std::path::Path;
      27              : use std::process::{Command, Stdio};
      28              : use std::str::FromStr;
      29              : use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
      30              : use std::sync::{Arc, Condvar, Mutex, RwLock};
      31              : use std::time::{Duration, Instant};
      32              : use std::{env, fs};
      33              : use tokio::{spawn, sync::watch, task::JoinHandle, time};
      34              : use tracing::{Instrument, debug, error, info, instrument, warn};
      35              : use url::Url;
      36              : use utils::id::{TenantId, TimelineId};
      37              : use utils::lsn::Lsn;
      38              : use utils::measured_stream::MeasuredReader;
      39              : use utils::pid_file;
      40              : use utils::shard::{ShardCount, ShardIndex, ShardNumber};
      41              : 
      42              : use crate::configurator::launch_configurator;
      43              : use crate::disk_quota::set_disk_quota;
      44              : use crate::installed_extensions::get_installed_extensions;
      45              : use crate::logger::startup_context_from_env;
      46              : use crate::lsn_lease::launch_lsn_lease_bg_task_for_static;
      47              : use crate::metrics::COMPUTE_CTL_UP;
      48              : use crate::monitor::launch_monitor;
      49              : use crate::pg_helpers::*;
      50              : use crate::pgbouncer::*;
      51              : use crate::rsyslog::{
      52              :     PostgresLogsRsyslogConfig, configure_audit_rsyslog, configure_postgres_logs_export,
      53              :     launch_pgaudit_gc,
      54              : };
      55              : use crate::spec::*;
      56              : use crate::swap::resize_swap;
      57              : use crate::sync_sk::{check_if_synced, ping_safekeeper};
      58              : use crate::tls::watch_cert_for_changes;
      59              : use crate::{config, extension_server, local_proxy};
      60              : 
      61              : pub static SYNC_SAFEKEEPERS_PID: AtomicU32 = AtomicU32::new(0);
      62              : pub static PG_PID: AtomicU32 = AtomicU32::new(0);
      63              : // This is an arbitrary build tag. Fine as a default / for testing purposes
      64              : // in-case of not-set environment var
      65              : const BUILD_TAG_DEFAULT: &str = "latest";
      66              : /// Build tag/version of the compute node binaries/image. It's tricky and ugly
      67              : /// to pass it everywhere as a part of `ComputeNodeParams`, so we use a
      68              : /// global static variable.
      69            0 : pub static BUILD_TAG: Lazy<String> = Lazy::new(|| {
      70            0 :     option_env!("BUILD_TAG")
      71            0 :         .unwrap_or(BUILD_TAG_DEFAULT)
      72            0 :         .to_string()
      73            0 : });
      74              : const DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL: u64 = 3600;
      75              : 
      76              : /// Static configuration params that don't change after startup. These mostly
      77              : /// come from the CLI args, or are derived from them.
      78              : #[derive(Clone, Debug)]
      79              : pub struct ComputeNodeParams {
      80              :     /// The ID of the compute
      81              :     pub compute_id: String,
      82              : 
      83              :     /// Url type maintains proper escaping
      84              :     pub connstr: url::Url,
      85              : 
      86              :     /// The name of the 'weak' superuser role, which we give to the users.
      87              :     /// It follows the allow list approach, i.e., we take a standard role
      88              :     /// and grant it extra permissions with explicit GRANTs here and there,
      89              :     /// and core patches.
      90              :     pub privileged_role_name: String,
      91              : 
      92              :     pub resize_swap_on_bind: bool,
      93              :     pub set_disk_quota_for_fs: Option<String>,
      94              : 
      95              :     // VM monitor parameters
      96              :     #[cfg(target_os = "linux")]
      97              :     pub filecache_connstr: String,
      98              :     #[cfg(target_os = "linux")]
      99              :     pub cgroup: String,
     100              :     #[cfg(target_os = "linux")]
     101              :     pub vm_monitor_addr: String,
     102              : 
     103              :     pub pgdata: String,
     104              :     pub pgbin: String,
     105              :     pub pgversion: String,
     106              : 
     107              :     /// The port that the compute's external HTTP server listens on
     108              :     pub external_http_port: u16,
     109              :     /// The port that the compute's internal HTTP server listens on
     110              :     pub internal_http_port: u16,
     111              : 
     112              :     /// the address of extension storage proxy gateway
     113              :     pub remote_ext_base_url: Option<Url>,
     114              : 
     115              :     /// Interval for installed extensions collection
     116              :     pub installed_extensions_collection_interval: Arc<AtomicU64>,
     117              :     /// Hadron instance ID of the compute node.
     118              :     pub instance_id: Option<String>,
     119              :     /// Timeout of PG compute startup in the Init state.
     120              :     pub pg_init_timeout: Option<Duration>,
     121              :     // Path to the `pg_isready` binary.
     122              :     pub pg_isready_bin: String,
     123              :     pub lakebase_mode: bool,
     124              : 
     125              :     pub build_tag: String,
     126              :     pub control_plane_uri: Option<String>,
     127              :     pub config_path_test_only: Option<OsString>,
     128              : }
     129              : 
     130              : type TaskHandle = Mutex<Option<JoinHandle<()>>>;
     131              : 
     132              : /// Compute node info shared across several `compute_ctl` threads.
     133              : pub struct ComputeNode {
     134              :     pub params: ComputeNodeParams,
     135              : 
     136              :     // We connect to Postgres from many different places, so build configs once
     137              :     // and reuse them where needed. These are derived from 'params.connstr'
     138              :     pub conn_conf: postgres::config::Config,
     139              :     pub tokio_conn_conf: tokio_postgres::config::Config,
     140              : 
     141              :     /// Volatile part of the `ComputeNode`, which should be used under `Mutex`.
     142              :     /// To allow HTTP API server to serving status requests, while configuration
     143              :     /// is in progress, lock should be held only for short periods of time to do
     144              :     /// read/write, not the whole configuration process.
     145              :     pub state: Mutex<ComputeState>,
     146              :     /// `Condvar` to allow notifying waiters about state changes.
     147              :     pub state_changed: Condvar,
     148              : 
     149              :     // key: ext_archive_name, value: started download time, download_completed?
     150              :     pub ext_download_progress: RwLock<HashMap<String, (DateTime<Utc>, bool)>>,
     151              :     pub compute_ctl_config: ComputeCtlConfig,
     152              : 
     153              :     /// Handle to the extension stats collection task
     154              :     extension_stats_task: TaskHandle,
     155              :     lfc_offload_task: TaskHandle,
     156              : }
     157              : 
     158              : // store some metrics about download size that might impact startup time
     159              : #[derive(Clone, Debug)]
     160              : pub struct RemoteExtensionMetrics {
     161              :     num_ext_downloaded: u64,
     162              :     largest_ext_size: u64,
     163              :     total_ext_download_size: u64,
     164              : }
     165              : 
     166              : #[derive(Clone, Debug)]
     167              : pub struct ComputeState {
     168              :     pub start_time: DateTime<Utc>,
     169              :     pub pg_start_time: Option<DateTime<Utc>>,
     170              :     pub status: ComputeStatus,
     171              :     /// Timestamp of the last Postgres activity. It could be `None` if
     172              :     /// compute wasn't used since start.
     173              :     pub last_active: Option<DateTime<Utc>>,
     174              :     pub error: Option<String>,
     175              : 
     176              :     /// Compute spec. This can be received from the CLI or - more likely -
     177              :     /// passed by the control plane with a /configure HTTP request.
     178              :     pub pspec: Option<ParsedSpec>,
     179              : 
     180              :     /// If the spec is passed by a /configure request, 'startup_span' is the
     181              :     /// /configure request's tracing span. The main thread enters it when it
     182              :     /// processes the compute startup, so that the compute startup is considered
     183              :     /// to be part of the /configure request for tracing purposes.
     184              :     ///
     185              :     /// If the request handling thread/task called startup_compute() directly,
     186              :     /// it would automatically be a child of the request handling span, and we
     187              :     /// wouldn't need this. But because we use the main thread to perform the
     188              :     /// startup, and the /configure task just waits for it to finish, we need to
     189              :     /// set up the span relationship ourselves.
     190              :     pub startup_span: Option<tracing::span::Span>,
     191              : 
     192              :     pub lfc_prewarm_state: LfcPrewarmState,
     193              :     pub lfc_offload_state: LfcOffloadState,
     194              : 
     195              :     /// WAL flush LSN that is set after terminating Postgres and syncing safekeepers if
     196              :     /// mode == ComputeMode::Primary. None otherwise
     197              :     pub terminate_flush_lsn: Option<Lsn>,
     198              :     pub promote_state: Option<watch::Receiver<PromoteState>>,
     199              : 
     200              :     pub metrics: ComputeMetrics,
     201              : }
     202              : 
     203              : impl ComputeState {
     204            0 :     pub fn new() -> Self {
     205            0 :         Self {
     206            0 :             start_time: Utc::now(),
     207            0 :             pg_start_time: None,
     208            0 :             status: ComputeStatus::Empty,
     209            0 :             last_active: None,
     210            0 :             error: None,
     211            0 :             pspec: None,
     212            0 :             startup_span: None,
     213            0 :             metrics: ComputeMetrics::default(),
     214            0 :             lfc_prewarm_state: LfcPrewarmState::default(),
     215            0 :             lfc_offload_state: LfcOffloadState::default(),
     216            0 :             terminate_flush_lsn: None,
     217            0 :             promote_state: None,
     218            0 :         }
     219            0 :     }
     220              : 
     221            0 :     pub fn set_status(&mut self, status: ComputeStatus, state_changed: &Condvar) {
     222            0 :         let prev = self.status;
     223            0 :         info!("Changing compute status from {} to {}", prev, status);
     224            0 :         self.status = status;
     225            0 :         state_changed.notify_all();
     226              : 
     227            0 :         COMPUTE_CTL_UP.reset();
     228            0 :         COMPUTE_CTL_UP
     229            0 :             .with_label_values(&[&BUILD_TAG, status.to_string().as_str()])
     230            0 :             .set(1);
     231            0 :     }
     232              : 
     233            0 :     pub fn set_failed_status(&mut self, err: anyhow::Error, state_changed: &Condvar) {
     234            0 :         self.error = Some(format!("{err:?}"));
     235            0 :         self.set_status(ComputeStatus::Failed, state_changed);
     236            0 :     }
     237              : }
     238              : 
     239              : impl Default for ComputeState {
     240            0 :     fn default() -> Self {
     241            0 :         Self::new()
     242            0 :     }
     243              : }
     244              : 
     245              : #[derive(Clone, Debug)]
     246              : pub struct ParsedSpec {
     247              :     pub spec: ComputeSpec,
     248              :     pub tenant_id: TenantId,
     249              :     pub timeline_id: TimelineId,
     250              :     pub pageserver_connstr: String,
     251              :     pub safekeeper_connstrings: Vec<String>,
     252              :     pub storage_auth_token: Option<String>,
     253              :     /// k8s dns name and port
     254              :     pub endpoint_storage_addr: Option<String>,
     255              :     pub endpoint_storage_token: Option<String>,
     256              : }
     257              : 
     258              : impl ParsedSpec {
     259            1 :     pub fn validate(&self) -> Result<(), String> {
     260              :         // Only Primary nodes are using safekeeper_connstrings, and at the moment
     261              :         // this method only validates that part of the specs.
     262            1 :         if self.spec.mode != ComputeMode::Primary {
     263            0 :             return Ok(());
     264            1 :         }
     265              : 
     266              :         // While it seems like a good idea to check for an odd number of entries in
     267              :         // the safekeepers connection string, changes to the list of safekeepers might
     268              :         // incur appending a new server to a list of 3, in which case a list of 4
     269              :         // entries is okay in production.
     270              :         //
     271              :         // Still we want unique entries, and at least one entry in the vector
     272            1 :         if self.safekeeper_connstrings.is_empty() {
     273            0 :             return Err(String::from("safekeeper_connstrings is empty"));
     274            1 :         }
     275              : 
     276              :         // check for uniqueness of the connection strings in the set
     277            1 :         let mut connstrings = self.safekeeper_connstrings.clone();
     278              : 
     279            1 :         connstrings.sort();
     280            1 :         let mut previous = &connstrings[0];
     281              : 
     282            2 :         for current in connstrings.iter().skip(1) {
     283              :             // duplicate entry?
     284            2 :             if current == previous {
     285            1 :                 return Err(format!(
     286            1 :                     "duplicate entry in safekeeper_connstrings: {current}!",
     287            1 :                 ));
     288            1 :             }
     289              : 
     290            1 :             previous = current;
     291              :         }
     292              : 
     293            0 :         Ok(())
     294            1 :     }
     295              : }
     296              : 
     297              : impl TryFrom<ComputeSpec> for ParsedSpec {
     298              :     type Error = String;
     299            1 :     fn try_from(spec: ComputeSpec) -> Result<Self, String> {
     300              :         // Extract the options from the spec file that are needed to connect to
     301              :         // the storage system.
     302              :         //
     303              :         // For backwards-compatibility, the top-level fields in the spec file
     304              :         // may be empty. In that case, we need to dig them from the GUCs in the
     305              :         // cluster.settings field.
     306            1 :         let pageserver_connstr = spec
     307            1 :             .pageserver_connstring
     308            1 :             .clone()
     309            1 :             .or_else(|| spec.cluster.settings.find("neon.pageserver_connstring"))
     310            1 :             .ok_or("pageserver connstr should be provided")?;
     311            1 :         let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
     312            1 :             if matches!(spec.mode, ComputeMode::Primary) {
     313            1 :                 spec.cluster
     314            1 :                     .settings
     315            1 :                     .find("neon.safekeepers")
     316            1 :                     .ok_or("safekeeper connstrings should be provided")?
     317            1 :                     .split(',')
     318            4 :                     .map(|str| str.to_string())
     319            1 :                     .collect()
     320              :             } else {
     321            0 :                 vec![]
     322              :             }
     323              :         } else {
     324            0 :             spec.safekeeper_connstrings.clone()
     325              :         };
     326              : 
     327            1 :         let storage_auth_token = spec.storage_auth_token.clone();
     328            1 :         let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
     329            0 :             tenant_id
     330              :         } else {
     331            1 :             spec.cluster
     332            1 :                 .settings
     333            1 :                 .find("neon.tenant_id")
     334            1 :                 .ok_or("tenant id should be provided")
     335            1 :                 .map(|s| TenantId::from_str(&s))?
     336            1 :                 .or(Err("invalid tenant id"))?
     337              :         };
     338            1 :         let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
     339            0 :             timeline_id
     340              :         } else {
     341            1 :             spec.cluster
     342            1 :                 .settings
     343            1 :                 .find("neon.timeline_id")
     344            1 :                 .ok_or("timeline id should be provided")
     345            1 :                 .map(|s| TimelineId::from_str(&s))?
     346            1 :                 .or(Err("invalid timeline id"))?
     347              :         };
     348              : 
     349            1 :         let endpoint_storage_addr: Option<String> = spec
     350            1 :             .endpoint_storage_addr
     351            1 :             .clone()
     352            1 :             .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_addr"));
     353            1 :         let endpoint_storage_token = spec
     354            1 :             .endpoint_storage_token
     355            1 :             .clone()
     356            1 :             .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_token"));
     357              : 
     358            1 :         let res = ParsedSpec {
     359            1 :             spec,
     360            1 :             pageserver_connstr,
     361            1 :             safekeeper_connstrings,
     362            1 :             storage_auth_token,
     363            1 :             tenant_id,
     364            1 :             timeline_id,
     365            1 :             endpoint_storage_addr,
     366            1 :             endpoint_storage_token,
     367            1 :         };
     368              : 
     369              :         // Now check validity of the parsed specification
     370            1 :         res.validate()?;
     371            0 :         Ok(res)
     372            1 :     }
     373              : }
     374              : 
     375              : /// If we are a VM, returns a [`Command`] that will run in the `neon-postgres`
     376              : /// cgroup. Otherwise returns the default `Command::new(cmd)`
     377              : ///
     378              : /// This function should be used to start postgres, as it will start it in the
     379              : /// neon-postgres cgroup if we are a VM. This allows autoscaling to control
     380              : /// postgres' resource usage. The cgroup will exist in VMs because vm-builder
     381              : /// creates it during the sysinit phase of its inittab.
     382            0 : fn maybe_cgexec(cmd: &str) -> Command {
     383              :     // The cplane sets this env var for autoscaling computes.
     384              :     // use `var_os` so we don't have to worry about the variable being valid
     385              :     // unicode. Should never be an concern . . . but just in case
     386            0 :     if env::var_os("AUTOSCALING").is_some() {
     387            0 :         let mut command = Command::new("cgexec");
     388            0 :         command.args(["-g", "memory:neon-postgres"]);
     389            0 :         command.arg(cmd);
     390            0 :         command
     391              :     } else {
     392            0 :         Command::new(cmd)
     393              :     }
     394            0 : }
     395              : 
     396              : struct PostgresHandle {
     397              :     postgres: std::process::Child,
     398              :     log_collector: JoinHandle<Result<()>>,
     399              : }
     400              : 
     401              : impl PostgresHandle {
     402              :     /// Return PID of the postgres (postmaster) process
     403            0 :     fn pid(&self) -> Pid {
     404            0 :         Pid::from_raw(self.postgres.id() as i32)
     405            0 :     }
     406              : }
     407              : 
     408              : struct StartVmMonitorResult {
     409              :     #[cfg(target_os = "linux")]
     410              :     token: tokio_util::sync::CancellationToken,
     411              :     #[cfg(target_os = "linux")]
     412              :     vm_monitor: Option<JoinHandle<Result<()>>>,
     413              : }
     414              : 
     415              : impl ComputeNode {
     416            0 :     pub fn new(params: ComputeNodeParams, config: ComputeConfig) -> Result<Self> {
     417            0 :         let connstr = params.connstr.as_str();
     418            0 :         let mut conn_conf = postgres::config::Config::from_str(connstr)
     419            0 :             .context("cannot build postgres config from connstr")?;
     420            0 :         let mut tokio_conn_conf = tokio_postgres::config::Config::from_str(connstr)
     421            0 :             .context("cannot build tokio postgres config from connstr")?;
     422              : 
     423              :         // Users can set some configuration parameters per database with
     424              :         //   ALTER DATABASE ... SET ...
     425              :         //
     426              :         // There are at least these parameters:
     427              :         //
     428              :         //   - role=some_other_role
     429              :         //   - default_transaction_read_only=on
     430              :         //   - statement_timeout=1, i.e., 1ms, which will cause most of the queries to fail
     431              :         //   - search_path=non_public_schema, this should be actually safe because
     432              :         //     we don't call any functions in user databases, but better to always reset
     433              :         //     it to public.
     434              :         //
     435              :         // that can affect `compute_ctl` and prevent it from properly configuring the database schema.
     436              :         // Unset them via connection string options before connecting to the database.
     437              :         // N.B. keep it in sync with `ZENITH_OPTIONS` in `get_maintenance_client()`.
     438              :         const EXTRA_OPTIONS: &str = "-c role=cloud_admin -c default_transaction_read_only=off -c search_path=public -c statement_timeout=0 -c pgaudit.log=none";
     439            0 :         let options = match conn_conf.get_options() {
     440              :             // Allow the control plane to override any options set by the
     441              :             // compute
     442            0 :             Some(options) => format!("{EXTRA_OPTIONS} {options}"),
     443            0 :             None => EXTRA_OPTIONS.to_string(),
     444              :         };
     445            0 :         conn_conf.options(&options);
     446            0 :         tokio_conn_conf.options(&options);
     447              : 
     448            0 :         let mut new_state = ComputeState::new();
     449            0 :         if let Some(spec) = config.spec {
     450            0 :             let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?;
     451            0 :             new_state.pspec = Some(pspec);
     452            0 :         }
     453              : 
     454            0 :         Ok(ComputeNode {
     455            0 :             params,
     456            0 :             conn_conf,
     457            0 :             tokio_conn_conf,
     458            0 :             state: Mutex::new(new_state),
     459            0 :             state_changed: Condvar::new(),
     460            0 :             ext_download_progress: RwLock::new(HashMap::new()),
     461            0 :             compute_ctl_config: config.compute_ctl_config,
     462            0 :             extension_stats_task: Mutex::new(None),
     463            0 :             lfc_offload_task: Mutex::new(None),
     464            0 :         })
     465            0 :     }
     466              : 
     467              :     /// Top-level control flow of compute_ctl. Returns a process exit code we should
     468              :     /// exit with.
     469            0 :     pub fn run(self) -> Result<Option<i32>> {
     470            0 :         let this = Arc::new(self);
     471              : 
     472            0 :         let cli_spec = this.state.lock().unwrap().pspec.clone();
     473              : 
     474              :         // If this is a pooled VM, prewarm before starting HTTP server and becoming
     475              :         // available for binding. Prewarming helps Postgres start quicker later,
     476              :         // because QEMU will already have its memory allocated from the host, and
     477              :         // the necessary binaries will already be cached.
     478            0 :         if cli_spec.is_none() {
     479            0 :             this.prewarm_postgres_vm_memory()?;
     480            0 :         }
     481              : 
     482              :         // Set the up metric with Empty status before starting the HTTP server.
     483              :         // That way on the first metric scrape, an external observer will see us
     484              :         // as 'up' and 'empty' (unless the compute was started with a spec or
     485              :         // already configured by control plane).
     486            0 :         COMPUTE_CTL_UP
     487            0 :             .with_label_values(&[&BUILD_TAG, ComputeStatus::Empty.to_string().as_str()])
     488            0 :             .set(1);
     489              : 
     490              :         // Launch the external HTTP server first, so that we can serve control plane
     491              :         // requests while configuration is still in progress.
     492            0 :         crate::http::server::Server::External {
     493            0 :             port: this.params.external_http_port,
     494            0 :             config: this.compute_ctl_config.clone(),
     495            0 :             compute_id: this.params.compute_id.clone(),
     496            0 :             instance_id: this.params.instance_id.clone(),
     497            0 :         }
     498            0 :         .launch(&this);
     499              : 
     500              :         // The internal HTTP server could be launched later, but there isn't much
     501              :         // sense in waiting.
     502            0 :         crate::http::server::Server::Internal {
     503            0 :             port: this.params.internal_http_port,
     504            0 :         }
     505            0 :         .launch(&this);
     506              : 
     507              :         // If we got a spec from the CLI already, use that. Otherwise wait for the
     508              :         // control plane to pass it to us with a /configure HTTP request
     509            0 :         let pspec = if let Some(cli_spec) = cli_spec {
     510            0 :             cli_spec
     511              :         } else {
     512            0 :             this.wait_spec()?
     513              :         };
     514              : 
     515            0 :         launch_lsn_lease_bg_task_for_static(&this);
     516              : 
     517              :         // We have a spec, start the compute
     518            0 :         let mut delay_exit = false;
     519            0 :         let mut vm_monitor = None;
     520            0 :         let mut pg_process: Option<PostgresHandle> = None;
     521              : 
     522            0 :         match this.start_compute(&mut pg_process) {
     523            0 :             Ok(()) => {
     524            0 :                 // Success! Launch remaining services (just vm-monitor currently)
     525            0 :                 vm_monitor =
     526            0 :                     Some(this.start_vm_monitor(pspec.spec.disable_lfc_resizing.unwrap_or(false)));
     527            0 :             }
     528            0 :             Err(err) => {
     529              :                 // Something went wrong with the startup. Log it and expose the error to
     530              :                 // HTTP status requests.
     531            0 :                 error!("could not start the compute node: {:#}", err);
     532            0 :                 this.set_failed_status(err);
     533            0 :                 delay_exit = true;
     534              : 
     535              :                 // If the error happened after starting PostgreSQL, kill it
     536            0 :                 if let Some(ref pg_process) = pg_process {
     537            0 :                     kill(pg_process.pid(), Signal::SIGQUIT).ok();
     538            0 :                 }
     539              :             }
     540              :         }
     541              : 
     542              :         // If startup was successful, or it failed in the late stages,
     543              :         // PostgreSQL is now running. Wait until it exits.
     544            0 :         let exit_code = if let Some(pg_handle) = pg_process {
     545            0 :             let exit_status = this.wait_postgres(pg_handle);
     546            0 :             info!("Postgres exited with code {}, shutting down", exit_status);
     547            0 :             exit_status.code()
     548              :         } else {
     549            0 :             None
     550              :         };
     551              : 
     552            0 :         this.terminate_extension_stats_task();
     553            0 :         this.terminate_lfc_offload_task();
     554              : 
     555              :         // Terminate the vm_monitor so it releases the file watcher on
     556              :         // /sys/fs/cgroup/neon-postgres.
     557              :         // Note: the vm-monitor only runs on linux because it requires cgroups.
     558            0 :         if let Some(vm_monitor) = vm_monitor {
     559              :             cfg_if::cfg_if! {
     560              :                 if #[cfg(target_os = "linux")] {
     561              :                     // Kills all threads spawned by the monitor
     562            0 :                     vm_monitor.token.cancel();
     563            0 :                     if let Some(handle) = vm_monitor.vm_monitor {
     564            0 :                         // Kills the actual task running the monitor
     565            0 :                         handle.abort();
     566            0 :                     }
     567              :                 } else {
     568              :                     _ = vm_monitor; // appease unused lint on macOS
     569              :                 }
     570              :             }
     571            0 :         }
     572              : 
     573              :         // Reap the postgres process
     574            0 :         delay_exit |= this.cleanup_after_postgres_exit()?;
     575              : 
     576              :         // /terminate returns LSN. If we don't sleep at all, connection will break and we
     577              :         // won't get result. If we sleep too much, tests will take significantly longer
     578              :         // and Github Action run will error out
     579            0 :         let sleep_duration = if delay_exit {
     580            0 :             Duration::from_secs(30)
     581              :         } else {
     582            0 :             Duration::from_millis(300)
     583              :         };
     584              : 
     585              :         // If launch failed, keep serving HTTP requests for a while, so the cloud
     586              :         // control plane can get the actual error.
     587            0 :         if delay_exit {
     588            0 :             info!("giving control plane 30s to collect the error before shutdown");
     589            0 :         }
     590            0 :         std::thread::sleep(sleep_duration);
     591            0 :         Ok(exit_code)
     592            0 :     }
     593              : 
     594            0 :     pub fn wait_spec(&self) -> Result<ParsedSpec> {
     595            0 :         info!("no compute spec provided, waiting");
     596            0 :         let mut state = self.state.lock().unwrap();
     597            0 :         while state.status != ComputeStatus::ConfigurationPending {
     598            0 :             state = self.state_changed.wait(state).unwrap();
     599            0 :         }
     600              : 
     601            0 :         info!("got spec, continue configuration");
     602            0 :         let spec = state.pspec.as_ref().unwrap().clone();
     603              : 
     604              :         // Record for how long we slept waiting for the spec.
     605            0 :         let now = Utc::now();
     606            0 :         state.metrics.wait_for_spec_ms = now
     607            0 :             .signed_duration_since(state.start_time)
     608            0 :             .to_std()
     609            0 :             .unwrap()
     610            0 :             .as_millis() as u64;
     611              : 
     612              :         // Reset start time, so that the total startup time that is calculated later will
     613              :         // not include the time that we waited for the spec.
     614            0 :         state.start_time = now;
     615              : 
     616            0 :         Ok(spec)
     617            0 :     }
     618              : 
     619              :     /// Start compute.
     620              :     ///
     621              :     /// Prerequisites:
     622              :     /// - the compute spec has been placed in self.state.pspec
     623              :     ///
     624              :     /// On success:
     625              :     /// - status is set to ComputeStatus::Running
     626              :     /// - self.running_postgres is set
     627              :     ///
     628              :     /// On error:
     629              :     /// - status is left in ComputeStatus::Init. The caller is responsible for setting it to Failed
     630              :     /// - if Postgres was started before the fatal error happened, self.running_postgres is
     631              :     ///   set. The caller is responsible for killing it.
     632              :     ///
     633              :     /// Note that this is in the critical path of a compute cold start. Keep this fast.
     634              :     /// Try to do things concurrently, to hide the latencies.
     635            0 :     fn start_compute(self: &Arc<Self>, pg_handle: &mut Option<PostgresHandle>) -> Result<()> {
     636              :         let compute_state: ComputeState;
     637              : 
     638              :         let start_compute_span;
     639              :         let _this_entered;
     640              :         {
     641            0 :             let mut state_guard = self.state.lock().unwrap();
     642              : 
     643              :             // Create a tracing span for the startup operation.
     644              :             //
     645              :             // We could otherwise just annotate the function with #[instrument], but if
     646              :             // we're being configured from a /configure HTTP request, we want the
     647              :             // startup to be considered part of the /configure request.
     648              :             //
     649              :             // Similarly, if a trace ID was passed in env variables, attach it to the span.
     650            0 :             start_compute_span = {
     651              :                 // Temporarily enter the parent span, so that the new span becomes its child.
     652            0 :                 if let Some(p) = state_guard.startup_span.take() {
     653            0 :                     let _parent_entered = p.entered();
     654            0 :                     tracing::info_span!("start_compute")
     655            0 :                 } else if let Some(otel_context) = startup_context_from_env() {
     656              :                     use tracing_opentelemetry::OpenTelemetrySpanExt;
     657            0 :                     let span = tracing::info_span!("start_compute");
     658            0 :                     span.set_parent(otel_context);
     659            0 :                     span
     660              :                 } else {
     661            0 :                     tracing::info_span!("start_compute")
     662              :                 }
     663              :             };
     664            0 :             _this_entered = start_compute_span.enter();
     665              : 
     666              :             // Hadron: Record postgres start time (used to enforce pg_init_timeout).
     667            0 :             state_guard.pg_start_time.replace(Utc::now());
     668              : 
     669            0 :             state_guard.set_status(ComputeStatus::Init, &self.state_changed);
     670            0 :             compute_state = state_guard.clone()
     671              :         }
     672              : 
     673            0 :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
     674            0 :         info!(
     675            0 :             "starting compute for project {}, operation {}, tenant {}, timeline {}, project {}, branch {}, endpoint {}, features {:?}, spec.remote_extensions {:?}",
     676            0 :             pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
     677            0 :             pspec.spec.operation_uuid.as_deref().unwrap_or("None"),
     678              :             pspec.tenant_id,
     679              :             pspec.timeline_id,
     680            0 :             pspec.spec.project_id.as_deref().unwrap_or("None"),
     681            0 :             pspec.spec.branch_id.as_deref().unwrap_or("None"),
     682            0 :             pspec.spec.endpoint_id.as_deref().unwrap_or("None"),
     683              :             pspec.spec.features,
     684              :             pspec.spec.remote_extensions,
     685              :         );
     686              : 
     687              :         ////// PRE-STARTUP PHASE: things that need to be finished before we start the Postgres process
     688              : 
     689              :         // Collect all the tasks that must finish here
     690            0 :         let mut pre_tasks = tokio::task::JoinSet::new();
     691              : 
     692              :         // Make sure TLS certificates are properly loaded and in the right place.
     693            0 :         if self.compute_ctl_config.tls.is_some() {
     694            0 :             let this = self.clone();
     695            0 :             pre_tasks.spawn(async move {
     696            0 :                 this.watch_cert_for_changes().await;
     697              : 
     698            0 :                 Ok::<(), anyhow::Error>(())
     699            0 :             });
     700            0 :         }
     701              : 
     702            0 :         let tls_config = self.tls_config(&pspec.spec);
     703              : 
     704              :         // If there are any remote extensions in shared_preload_libraries, start downloading them
     705            0 :         if pspec.spec.remote_extensions.is_some() {
     706            0 :             let (this, spec) = (self.clone(), pspec.spec.clone());
     707            0 :             pre_tasks.spawn(async move {
     708            0 :                 this.download_preload_extensions(&spec)
     709            0 :                     .in_current_span()
     710            0 :                     .await
     711            0 :             });
     712            0 :         }
     713              : 
     714              :         // Prepare pgdata directory. This downloads the basebackup, among other things.
     715              :         {
     716            0 :             let (this, cs) = (self.clone(), compute_state.clone());
     717            0 :             pre_tasks.spawn_blocking_child(move || this.prepare_pgdata(&cs));
     718              :         }
     719              : 
     720              :         // Resize swap to the desired size if the compute spec says so
     721            0 :         if let (Some(size_bytes), true) =
     722            0 :             (pspec.spec.swap_size_bytes, self.params.resize_swap_on_bind)
     723              :         {
     724            0 :             pre_tasks.spawn_blocking_child(move || {
     725              :                 // To avoid 'swapoff' hitting postgres startup, we need to run resize-swap to completion
     726              :                 // *before* starting postgres.
     727              :                 //
     728              :                 // In theory, we could do this asynchronously if SkipSwapon was enabled for VMs, but this
     729              :                 // carries a risk of introducing hard-to-debug issues - e.g. if postgres sometimes gets
     730              :                 // OOM-killed during startup because swap wasn't available yet.
     731            0 :                 resize_swap(size_bytes).context("failed to resize swap")?;
     732            0 :                 let size_mib = size_bytes as f32 / (1 << 20) as f32; // just for more coherent display.
     733            0 :                 info!(%size_bytes, %size_mib, "resized swap");
     734              : 
     735            0 :                 Ok::<(), anyhow::Error>(())
     736            0 :             });
     737            0 :         }
     738              : 
     739              :         // Set disk quota if the compute spec says so
     740            0 :         if let (Some(disk_quota_bytes), Some(disk_quota_fs_mountpoint)) = (
     741            0 :             pspec.spec.disk_quota_bytes,
     742            0 :             self.params.set_disk_quota_for_fs.as_ref(),
     743              :         ) {
     744            0 :             let disk_quota_fs_mountpoint = disk_quota_fs_mountpoint.clone();
     745            0 :             pre_tasks.spawn_blocking_child(move || {
     746            0 :                 set_disk_quota(disk_quota_bytes, &disk_quota_fs_mountpoint)
     747            0 :                     .context("failed to set disk quota")?;
     748            0 :                 let size_mib = disk_quota_bytes as f32 / (1 << 20) as f32; // just for more coherent display.
     749            0 :                 info!(%disk_quota_bytes, %size_mib, "set disk quota");
     750              : 
     751            0 :                 Ok::<(), anyhow::Error>(())
     752            0 :             });
     753            0 :         }
     754              : 
     755              :         // tune pgbouncer
     756            0 :         if let Some(pgbouncer_settings) = &pspec.spec.pgbouncer_settings {
     757            0 :             info!("tuning pgbouncer");
     758              : 
     759            0 :             let pgbouncer_settings = pgbouncer_settings.clone();
     760            0 :             let tls_config = tls_config.clone();
     761              : 
     762              :             // Spawn a background task to do the tuning,
     763              :             // so that we don't block the main thread that starts Postgres.
     764            0 :             let _handle = tokio::spawn(async move {
     765            0 :                 let res = tune_pgbouncer(pgbouncer_settings, tls_config).await;
     766            0 :                 if let Err(err) = res {
     767            0 :                     error!("error while tuning pgbouncer: {err:?}");
     768              :                     // Continue with the startup anyway
     769            0 :                 }
     770            0 :             });
     771            0 :         }
     772              : 
     773              :         // configure local_proxy
     774            0 :         if let Some(local_proxy) = &pspec.spec.local_proxy_config {
     775            0 :             info!("configuring local_proxy");
     776              : 
     777              :             // Spawn a background task to do the configuration,
     778              :             // so that we don't block the main thread that starts Postgres.
     779              : 
     780            0 :             let mut local_proxy = local_proxy.clone();
     781            0 :             local_proxy.tls = tls_config.clone();
     782              : 
     783            0 :             let _handle = tokio::spawn(async move {
     784            0 :                 if let Err(err) = local_proxy::configure(&local_proxy) {
     785            0 :                     error!("error while configuring local_proxy: {err:?}");
     786              :                     // Continue with the startup anyway
     787            0 :                 }
     788            0 :             });
     789            0 :         }
     790              : 
     791              :         // Configure and start rsyslog for compliance audit logging
     792            0 :         match pspec.spec.audit_log_level {
     793              :             ComputeAudit::Hipaa | ComputeAudit::Extended | ComputeAudit::Full => {
     794            0 :                 let remote_tls_endpoint =
     795            0 :                     std::env::var("AUDIT_LOGGING_TLS_ENDPOINT").unwrap_or("".to_string());
     796            0 :                 let remote_plain_endpoint =
     797            0 :                     std::env::var("AUDIT_LOGGING_ENDPOINT").unwrap_or("".to_string());
     798              : 
     799            0 :                 if remote_plain_endpoint.is_empty() && remote_tls_endpoint.is_empty() {
     800            0 :                     anyhow::bail!(
     801            0 :                         "AUDIT_LOGGING_ENDPOINT and AUDIT_LOGGING_TLS_ENDPOINT are both empty"
     802              :                     );
     803            0 :                 }
     804              : 
     805            0 :                 let log_directory_path = Path::new(&self.params.pgdata).join("log");
     806            0 :                 let log_directory_path = log_directory_path.to_string_lossy().to_string();
     807              : 
     808              :                 // Add project_id,endpoint_id to identify the logs.
     809              :                 //
     810              :                 // These ids are passed from cplane,
     811            0 :                 let endpoint_id = pspec.spec.endpoint_id.as_deref().unwrap_or("");
     812            0 :                 let project_id = pspec.spec.project_id.as_deref().unwrap_or("");
     813              : 
     814            0 :                 configure_audit_rsyslog(
     815            0 :                     log_directory_path.clone(),
     816            0 :                     endpoint_id,
     817            0 :                     project_id,
     818            0 :                     &remote_plain_endpoint,
     819            0 :                     &remote_tls_endpoint,
     820            0 :                 )?;
     821              : 
     822              :                 // Launch a background task to clean up the audit logs
     823            0 :                 launch_pgaudit_gc(log_directory_path);
     824              :             }
     825            0 :             _ => {}
     826              :         }
     827              : 
     828              :         // Configure and start rsyslog for Postgres logs export
     829            0 :         let conf = PostgresLogsRsyslogConfig::new(pspec.spec.logs_export_host.as_deref());
     830            0 :         configure_postgres_logs_export(conf)?;
     831              : 
     832              :         // Launch remaining service threads
     833            0 :         let _monitor_handle = launch_monitor(self);
     834            0 :         let _configurator_handle = launch_configurator(self);
     835              : 
     836              :         // Wait for all the pre-tasks to finish before starting postgres
     837            0 :         let rt = tokio::runtime::Handle::current();
     838            0 :         while let Some(res) = rt.block_on(pre_tasks.join_next()) {
     839            0 :             res??;
     840              :         }
     841              : 
     842              :         ////// START POSTGRES
     843            0 :         let start_time = Utc::now();
     844            0 :         let pg_process = self.start_postgres(pspec.storage_auth_token.clone())?;
     845            0 :         let postmaster_pid = pg_process.pid();
     846            0 :         *pg_handle = Some(pg_process);
     847              : 
     848              :         // If this is a primary endpoint, perform some post-startup configuration before
     849              :         // opening it up for the world.
     850            0 :         let config_time = Utc::now();
     851            0 :         if pspec.spec.mode == ComputeMode::Primary {
     852            0 :             self.configure_as_primary(&compute_state)?;
     853              : 
     854            0 :             let conf = self.get_tokio_conn_conf(None);
     855            0 :             tokio::task::spawn(async {
     856            0 :                 let _ = installed_extensions(conf).await;
     857            0 :             });
     858            0 :         }
     859              : 
     860              :         // All done!
     861            0 :         let startup_end_time = Utc::now();
     862            0 :         let metrics = {
     863            0 :             let mut state = self.state.lock().unwrap();
     864            0 :             state.metrics.start_postgres_ms = config_time
     865            0 :                 .signed_duration_since(start_time)
     866            0 :                 .to_std()
     867            0 :                 .unwrap()
     868            0 :                 .as_millis() as u64;
     869            0 :             state.metrics.config_ms = startup_end_time
     870            0 :                 .signed_duration_since(config_time)
     871            0 :                 .to_std()
     872            0 :                 .unwrap()
     873            0 :                 .as_millis() as u64;
     874            0 :             state.metrics.total_startup_ms = startup_end_time
     875            0 :                 .signed_duration_since(compute_state.start_time)
     876            0 :                 .to_std()
     877            0 :                 .unwrap()
     878            0 :                 .as_millis() as u64;
     879            0 :             state.metrics.clone()
     880              :         };
     881            0 :         self.set_status(ComputeStatus::Running);
     882              : 
     883              :         // Log metrics so that we can search for slow operations in logs
     884            0 :         info!(?metrics, postmaster_pid = %postmaster_pid, "compute start finished");
     885              : 
     886            0 :         self.spawn_extension_stats_task();
     887              : 
     888            0 :         if pspec.spec.autoprewarm {
     889            0 :             info!("autoprewarming on startup as requested");
     890            0 :             self.prewarm_lfc(None);
     891            0 :         }
     892            0 :         if let Some(seconds) = pspec.spec.offload_lfc_interval_seconds {
     893            0 :             self.spawn_lfc_offload_task(Duration::from_secs(seconds.into()));
     894            0 :         };
     895            0 :         Ok(())
     896            0 :     }
     897              : 
     898              :     #[instrument(skip_all)]
     899              :     async fn download_preload_extensions(&self, spec: &ComputeSpec) -> Result<()> {
     900              :         let remote_extensions = if let Some(remote_extensions) = &spec.remote_extensions {
     901              :             remote_extensions
     902              :         } else {
     903              :             return Ok(());
     904              :         };
     905              : 
     906              :         // First, create control files for all available extensions
     907              :         extension_server::create_control_files(remote_extensions, &self.params.pgbin);
     908              : 
     909              :         let library_load_start_time = Utc::now();
     910              :         let remote_ext_metrics = self.prepare_preload_libraries(spec).await?;
     911              : 
     912              :         let library_load_time = Utc::now()
     913              :             .signed_duration_since(library_load_start_time)
     914              :             .to_std()
     915              :             .unwrap()
     916              :             .as_millis() as u64;
     917              :         let mut state = self.state.lock().unwrap();
     918              :         state.metrics.load_ext_ms = library_load_time;
     919              :         state.metrics.num_ext_downloaded = remote_ext_metrics.num_ext_downloaded;
     920              :         state.metrics.largest_ext_size = remote_ext_metrics.largest_ext_size;
     921              :         state.metrics.total_ext_download_size = remote_ext_metrics.total_ext_download_size;
     922              :         info!(
     923              :             "Loading shared_preload_libraries took {:?}ms",
     924              :             library_load_time
     925              :         );
     926              :         info!("{:?}", remote_ext_metrics);
     927              : 
     928              :         Ok(())
     929              :     }
     930              : 
     931              :     /// Start the vm-monitor if directed to. The vm-monitor only runs on linux
     932              :     /// because it requires cgroups.
     933            0 :     fn start_vm_monitor(&self, disable_lfc_resizing: bool) -> StartVmMonitorResult {
     934              :         cfg_if::cfg_if! {
     935              :             if #[cfg(target_os = "linux")] {
     936              :                 use std::env;
     937              :                 use tokio_util::sync::CancellationToken;
     938              : 
     939              :                 // This token is used internally by the monitor to clean up all threads
     940            0 :                 let token = CancellationToken::new();
     941              : 
     942              :                 // don't pass postgres connection string to vm-monitor if we don't want it to resize LFC
     943            0 :                 let pgconnstr = if disable_lfc_resizing {
     944            0 :                     None
     945              :                 } else {
     946            0 :                     Some(self.params.filecache_connstr.clone())
     947              :                 };
     948              : 
     949            0 :                 let vm_monitor = if env::var_os("AUTOSCALING").is_some() {
     950            0 :                     let vm_monitor = tokio::spawn(vm_monitor::start(
     951            0 :                         Box::leak(Box::new(vm_monitor::Args {
     952            0 :                             cgroup: Some(self.params.cgroup.clone()),
     953            0 :                             pgconnstr,
     954            0 :                             addr: self.params.vm_monitor_addr.clone(),
     955            0 :                         })),
     956            0 :                         token.clone(),
     957              :                     ));
     958            0 :                     Some(vm_monitor)
     959              :                 } else {
     960            0 :                     None
     961              :                 };
     962            0 :                 StartVmMonitorResult { token, vm_monitor }
     963              :             } else {
     964              :                 _ = disable_lfc_resizing; // appease unused lint on macOS
     965              :                 StartVmMonitorResult { }
     966              :             }
     967              :         }
     968            0 :     }
     969              : 
     970            0 :     fn cleanup_after_postgres_exit(&self) -> Result<bool> {
     971              :         // Maybe sync safekeepers again, to speed up next startup
     972            0 :         let compute_state = self.state.lock().unwrap().clone();
     973            0 :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
     974            0 :         let lsn = if matches!(pspec.spec.mode, compute_api::spec::ComputeMode::Primary) {
     975            0 :             info!("syncing safekeepers on shutdown");
     976            0 :             let storage_auth_token = pspec.storage_auth_token.clone();
     977            0 :             let lsn = self.sync_safekeepers(storage_auth_token)?;
     978            0 :             info!(%lsn, "synced safekeepers");
     979            0 :             Some(lsn)
     980              :         } else {
     981            0 :             info!("not primary, not syncing safekeepers");
     982            0 :             None
     983              :         };
     984              : 
     985            0 :         let mut state = self.state.lock().unwrap();
     986            0 :         state.terminate_flush_lsn = lsn;
     987              : 
     988            0 :         let delay_exit = state.status == ComputeStatus::TerminationPendingFast;
     989            0 :         if state.status == ComputeStatus::TerminationPendingFast
     990            0 :             || state.status == ComputeStatus::TerminationPendingImmediate
     991              :         {
     992            0 :             info!(
     993            0 :                 "Changing compute status from {} to {}",
     994            0 :                 state.status,
     995              :                 ComputeStatus::Terminated
     996              :             );
     997            0 :             state.status = ComputeStatus::Terminated;
     998            0 :             self.state_changed.notify_all();
     999            0 :         }
    1000            0 :         drop(state);
    1001              : 
    1002            0 :         if let Err(err) = self.check_for_core_dumps() {
    1003            0 :             error!("error while checking for core dumps: {err:?}");
    1004            0 :         }
    1005              : 
    1006            0 :         Ok(delay_exit)
    1007            0 :     }
    1008              : 
    1009              :     /// Check that compute node has corresponding feature enabled.
    1010            0 :     pub fn has_feature(&self, feature: ComputeFeature) -> bool {
    1011            0 :         let state = self.state.lock().unwrap();
    1012              : 
    1013            0 :         if let Some(s) = state.pspec.as_ref() {
    1014            0 :             s.spec.features.contains(&feature)
    1015              :         } else {
    1016            0 :             false
    1017              :         }
    1018            0 :     }
    1019              : 
    1020            0 :     pub fn set_status(&self, status: ComputeStatus) {
    1021            0 :         let mut state = self.state.lock().unwrap();
    1022            0 :         state.set_status(status, &self.state_changed);
    1023            0 :     }
    1024              : 
    1025            0 :     pub fn set_failed_status(&self, err: anyhow::Error) {
    1026            0 :         let mut state = self.state.lock().unwrap();
    1027            0 :         state.set_failed_status(err, &self.state_changed);
    1028            0 :     }
    1029              : 
    1030            0 :     pub fn get_status(&self) -> ComputeStatus {
    1031            0 :         self.state.lock().unwrap().status
    1032            0 :     }
    1033              : 
    1034            0 :     pub fn get_timeline_id(&self) -> Option<TimelineId> {
    1035            0 :         self.state
    1036            0 :             .lock()
    1037            0 :             .unwrap()
    1038            0 :             .pspec
    1039            0 :             .as_ref()
    1040            0 :             .map(|s| s.timeline_id)
    1041            0 :     }
    1042              : 
    1043              :     // Remove `pgdata` directory and create it again with right permissions.
    1044            0 :     fn create_pgdata(&self) -> Result<()> {
    1045              :         // Ignore removal error, likely it is a 'No such file or directory (os error 2)'.
    1046              :         // If it is something different then create_dir() will error out anyway.
    1047            0 :         let pgdata = &self.params.pgdata;
    1048            0 :         let _ok = fs::remove_dir_all(pgdata);
    1049            0 :         fs::create_dir(pgdata)?;
    1050            0 :         fs::set_permissions(pgdata, fs::Permissions::from_mode(0o700))?;
    1051              : 
    1052            0 :         Ok(())
    1053            0 :     }
    1054              : 
    1055              :     /// Fetches a basebackup from the Pageserver using the compute state's Pageserver connstring and
    1056              :     /// unarchives it to `pgdata` directory, replacing any existing contents.
    1057              :     #[instrument(skip_all, fields(%lsn))]
    1058              :     fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
    1059              :         let spec = compute_state.pspec.as_ref().expect("spec must be set");
    1060              : 
    1061              :         let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
    1062              :         let started = Instant::now();
    1063              : 
    1064              :         let (connected, size) = match PageserverProtocol::from_connstring(shard0_connstr)? {
    1065              :             PageserverProtocol::Libpq => self.try_get_basebackup_libpq(spec, lsn)?,
    1066              :             PageserverProtocol::Grpc => self.try_get_basebackup_grpc(spec, lsn)?,
    1067              :         };
    1068              : 
    1069              :         self.fix_zenith_signal_neon_signal()?;
    1070              : 
    1071              :         let mut state = self.state.lock().unwrap();
    1072              :         state.metrics.pageserver_connect_micros =
    1073              :             connected.duration_since(started).as_micros() as u64;
    1074              :         state.metrics.basebackup_bytes = size as u64;
    1075              :         state.metrics.basebackup_ms = started.elapsed().as_millis() as u64;
    1076              : 
    1077              :         Ok(())
    1078              :     }
    1079              : 
    1080              :     /// Move the Zenith signal file to Neon signal file location.
    1081              :     /// This makes Compute compatible with older PageServers that don't yet
    1082              :     /// know about the Zenith->Neon rename.
    1083            0 :     fn fix_zenith_signal_neon_signal(&self) -> Result<()> {
    1084            0 :         let datadir = Path::new(&self.params.pgdata);
    1085              : 
    1086            0 :         let neonsig = datadir.join("neon.signal");
    1087              : 
    1088            0 :         if neonsig.is_file() {
    1089            0 :             return Ok(());
    1090            0 :         }
    1091              : 
    1092            0 :         let zenithsig = datadir.join("zenith.signal");
    1093              : 
    1094            0 :         if zenithsig.is_file() {
    1095            0 :             fs::copy(zenithsig, neonsig)?;
    1096            0 :         }
    1097              : 
    1098            0 :         Ok(())
    1099            0 :     }
    1100              : 
    1101              :     /// Fetches a basebackup via gRPC. The connstring must use grpc://. Returns the timestamp when
    1102              :     /// the connection was established, and the (compressed) size of the basebackup.
    1103            0 :     fn try_get_basebackup_grpc(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
    1104            0 :         let shard0_connstr = spec
    1105            0 :             .pageserver_connstr
    1106            0 :             .split(',')
    1107            0 :             .next()
    1108            0 :             .unwrap()
    1109            0 :             .to_string();
    1110            0 :         let shard_index = match spec.pageserver_connstr.split(',').count() as u8 {
    1111            0 :             0 | 1 => ShardIndex::unsharded(),
    1112            0 :             count => ShardIndex::new(ShardNumber(0), ShardCount(count)),
    1113              :         };
    1114              : 
    1115            0 :         let (reader, connected) = tokio::runtime::Handle::current().block_on(async move {
    1116            0 :             let mut client = page_api::Client::connect(
    1117            0 :                 shard0_connstr,
    1118            0 :                 spec.tenant_id,
    1119            0 :                 spec.timeline_id,
    1120            0 :                 shard_index,
    1121            0 :                 spec.storage_auth_token.clone(),
    1122            0 :                 None, // NB: base backups use payload compression
    1123            0 :             )
    1124            0 :             .await?;
    1125            0 :             let connected = Instant::now();
    1126            0 :             let reader = client
    1127            0 :                 .get_base_backup(page_api::GetBaseBackupRequest {
    1128            0 :                     lsn: (lsn != Lsn(0)).then_some(lsn),
    1129            0 :                     compression: BaseBackupCompression::Gzip,
    1130            0 :                     replica: spec.spec.mode != ComputeMode::Primary,
    1131            0 :                     full: false,
    1132            0 :                 })
    1133            0 :                 .await?;
    1134            0 :             anyhow::Ok((reader, connected))
    1135            0 :         })?;
    1136              : 
    1137            0 :         let mut reader = MeasuredReader::new(tokio_util::io::SyncIoBridge::new(reader));
    1138              : 
    1139              :         // Set `ignore_zeros` so that unpack() reads the entire stream and doesn't just stop at the
    1140              :         // end-of-archive marker. If the server errors, the tar::Builder drop handler will write an
    1141              :         // end-of-archive marker before the error is emitted, and we would not see the error.
    1142            0 :         let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut reader));
    1143            0 :         ar.set_ignore_zeros(true);
    1144            0 :         ar.unpack(&self.params.pgdata)?;
    1145              : 
    1146            0 :         Ok((connected, reader.get_byte_count()))
    1147            0 :     }
    1148              : 
    1149              :     /// Fetches a basebackup via libpq. The connstring must use postgresql://. Returns the timestamp
    1150              :     /// when the connection was established, and the (compressed) size of the basebackup.
    1151            0 :     fn try_get_basebackup_libpq(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
    1152            0 :         let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
    1153            0 :         let mut config = postgres::Config::from_str(shard0_connstr)?;
    1154              : 
    1155              :         // Use the storage auth token from the config file, if given.
    1156              :         // Note: this overrides any password set in the connection string.
    1157            0 :         if let Some(storage_auth_token) = &spec.storage_auth_token {
    1158            0 :             info!("Got storage auth token from spec file");
    1159            0 :             config.password(storage_auth_token);
    1160              :         } else {
    1161            0 :             info!("Storage auth token not set");
    1162              :         }
    1163              : 
    1164            0 :         config.application_name("compute_ctl");
    1165            0 :         config.options(&format!(
    1166            0 :             "-c neon.compute_mode={}",
    1167            0 :             spec.spec.mode.to_type_str()
    1168            0 :         ));
    1169              : 
    1170              :         // Connect to pageserver
    1171            0 :         let mut client = config.connect(NoTls)?;
    1172            0 :         let connected = Instant::now();
    1173              : 
    1174            0 :         let basebackup_cmd = match lsn {
    1175              :             Lsn(0) => {
    1176            0 :                 if spec.spec.mode != ComputeMode::Primary {
    1177            0 :                     format!(
    1178            0 :                         "basebackup {} {} --gzip --replica",
    1179              :                         spec.tenant_id, spec.timeline_id
    1180              :                     )
    1181              :                 } else {
    1182            0 :                     format!("basebackup {} {} --gzip", spec.tenant_id, spec.timeline_id)
    1183              :                 }
    1184              :             }
    1185              :             _ => {
    1186            0 :                 if spec.spec.mode != ComputeMode::Primary {
    1187            0 :                     format!(
    1188            0 :                         "basebackup {} {} {} --gzip --replica",
    1189              :                         spec.tenant_id, spec.timeline_id, lsn
    1190              :                     )
    1191              :                 } else {
    1192            0 :                     format!(
    1193            0 :                         "basebackup {} {} {} --gzip",
    1194              :                         spec.tenant_id, spec.timeline_id, lsn
    1195              :                     )
    1196              :                 }
    1197              :             }
    1198              :         };
    1199              : 
    1200            0 :         let copyreader = client.copy_out(basebackup_cmd.as_str())?;
    1201            0 :         let mut measured_reader = MeasuredReader::new(copyreader);
    1202            0 :         let mut bufreader = std::io::BufReader::new(&mut measured_reader);
    1203              : 
    1204              :         // Read the archive directly from the `CopyOutReader`
    1205              :         //
    1206              :         // Set `ignore_zeros` so that unpack() reads all the Copy data and
    1207              :         // doesn't stop at the end-of-archive marker. Otherwise, if the server
    1208              :         // sends an Error after finishing the tarball, we will not notice it.
    1209              :         // The tar::Builder drop handler will write an end-of-archive marker
    1210              :         // before emitting the error, and we would not see it otherwise.
    1211            0 :         let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
    1212            0 :         ar.set_ignore_zeros(true);
    1213            0 :         ar.unpack(&self.params.pgdata)?;
    1214              : 
    1215            0 :         Ok((connected, measured_reader.get_byte_count()))
    1216            0 :     }
    1217              : 
    1218              :     // Gets the basebackup in a retry loop
    1219              :     #[instrument(skip_all, fields(%lsn))]
    1220              :     pub fn get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
    1221              :         let mut retry_period_ms = 500.0;
    1222              :         let mut attempts = 0;
    1223              :         const DEFAULT_ATTEMPTS: u16 = 10;
    1224              :         #[cfg(feature = "testing")]
    1225              :         let max_attempts = if let Ok(v) = env::var("NEON_COMPUTE_TESTING_BASEBACKUP_RETRIES") {
    1226              :             u16::from_str(&v).unwrap()
    1227              :         } else {
    1228              :             DEFAULT_ATTEMPTS
    1229              :         };
    1230              :         #[cfg(not(feature = "testing"))]
    1231              :         let max_attempts = DEFAULT_ATTEMPTS;
    1232              :         loop {
    1233              :             let result = self.try_get_basebackup(compute_state, lsn);
    1234              :             match result {
    1235              :                 Ok(_) => {
    1236              :                     return result;
    1237              :                 }
    1238              :                 Err(ref e) if attempts < max_attempts => {
    1239              :                     warn!(
    1240              :                         "Failed to get basebackup: {} (attempt {}/{})",
    1241              :                         e, attempts, max_attempts
    1242              :                     );
    1243              :                     std::thread::sleep(std::time::Duration::from_millis(retry_period_ms as u64));
    1244              :                     retry_period_ms *= 1.5;
    1245              :                 }
    1246              :                 Err(_) => {
    1247              :                     return result;
    1248              :                 }
    1249              :             }
    1250              :             attempts += 1;
    1251              :         }
    1252              :     }
    1253              : 
    1254            0 :     pub async fn check_safekeepers_synced_async(
    1255            0 :         &self,
    1256            0 :         compute_state: &ComputeState,
    1257            0 :     ) -> Result<Option<Lsn>> {
    1258              :         // Construct a connection config for each safekeeper
    1259            0 :         let pspec: ParsedSpec = compute_state
    1260            0 :             .pspec
    1261            0 :             .as_ref()
    1262            0 :             .expect("spec must be set")
    1263            0 :             .clone();
    1264            0 :         let sk_connstrs: Vec<String> = pspec.safekeeper_connstrings.clone();
    1265            0 :         let sk_configs = sk_connstrs.into_iter().map(|connstr| {
    1266              :             // Format connstr
    1267            0 :             let id = connstr.clone();
    1268            0 :             let connstr = format!("postgresql://no_user@{connstr}");
    1269            0 :             let options = format!(
    1270            0 :                 "-c timeline_id={} tenant_id={}",
    1271              :                 pspec.timeline_id, pspec.tenant_id
    1272              :             );
    1273              : 
    1274              :             // Construct client
    1275            0 :             let mut config = tokio_postgres::Config::from_str(&connstr).unwrap();
    1276            0 :             config.options(&options);
    1277            0 :             if let Some(storage_auth_token) = pspec.storage_auth_token.clone() {
    1278            0 :                 config.password(storage_auth_token);
    1279            0 :             }
    1280              : 
    1281            0 :             (id, config)
    1282            0 :         });
    1283              : 
    1284              :         // Create task set to query all safekeepers
    1285            0 :         let mut tasks = FuturesUnordered::new();
    1286            0 :         let quorum = sk_configs.len() / 2 + 1;
    1287            0 :         for (id, config) in sk_configs {
    1288            0 :             let timeout = tokio::time::Duration::from_millis(100);
    1289            0 :             let task = tokio::time::timeout(timeout, ping_safekeeper(id, config));
    1290            0 :             tasks.push(tokio::spawn(task));
    1291            0 :         }
    1292              : 
    1293              :         // Get a quorum of responses or errors
    1294            0 :         let mut responses = Vec::new();
    1295            0 :         let mut join_errors = Vec::new();
    1296            0 :         let mut task_errors = Vec::new();
    1297            0 :         let mut timeout_errors = Vec::new();
    1298            0 :         while let Some(response) = tasks.next().await {
    1299            0 :             match response {
    1300            0 :                 Ok(Ok(Ok(r))) => responses.push(r),
    1301            0 :                 Ok(Ok(Err(e))) => task_errors.push(e),
    1302            0 :                 Ok(Err(e)) => timeout_errors.push(e),
    1303            0 :                 Err(e) => join_errors.push(e),
    1304              :             };
    1305            0 :             if responses.len() >= quorum {
    1306            0 :                 break;
    1307            0 :             }
    1308            0 :             if join_errors.len() + task_errors.len() + timeout_errors.len() >= quorum {
    1309            0 :                 break;
    1310            0 :             }
    1311              :         }
    1312              : 
    1313              :         // In case of error, log and fail the check, but don't crash.
    1314              :         // We're playing it safe because these errors could be transient
    1315              :         // and we don't yet retry.
    1316            0 :         if responses.len() < quorum {
    1317            0 :             error!(
    1318            0 :                 "failed sync safekeepers check {:?} {:?} {:?}",
    1319              :                 join_errors, task_errors, timeout_errors
    1320              :             );
    1321            0 :             return Ok(None);
    1322            0 :         }
    1323              : 
    1324            0 :         Ok(check_if_synced(responses))
    1325            0 :     }
    1326              : 
    1327              :     // Fast path for sync_safekeepers. If they're already synced we get the lsn
    1328              :     // in one roundtrip. If not, we should do a full sync_safekeepers.
    1329              :     #[instrument(skip_all)]
    1330              :     pub fn check_safekeepers_synced(&self, compute_state: &ComputeState) -> Result<Option<Lsn>> {
    1331              :         let start_time = Utc::now();
    1332              : 
    1333              :         let rt = tokio::runtime::Handle::current();
    1334              :         let result = rt.block_on(self.check_safekeepers_synced_async(compute_state));
    1335              : 
    1336              :         // Record runtime
    1337              :         self.state.lock().unwrap().metrics.sync_sk_check_ms = Utc::now()
    1338              :             .signed_duration_since(start_time)
    1339              :             .to_std()
    1340              :             .unwrap()
    1341              :             .as_millis() as u64;
    1342              :         result
    1343              :     }
    1344              : 
    1345              :     // Run `postgres` in a special mode with `--sync-safekeepers` argument
    1346              :     // and return the reported LSN back to the caller.
    1347              :     #[instrument(skip_all)]
    1348              :     pub fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
    1349              :         let start_time = Utc::now();
    1350              : 
    1351              :         let mut sync_handle = maybe_cgexec(&self.params.pgbin)
    1352              :             .args(["--sync-safekeepers"])
    1353              :             .env("PGDATA", &self.params.pgdata) // we cannot use -D in this mode
    1354              :             .envs(if let Some(storage_auth_token) = &storage_auth_token {
    1355              :                 vec![("NEON_AUTH_TOKEN", storage_auth_token)]
    1356              :             } else {
    1357              :                 vec![]
    1358              :             })
    1359              :             .stdout(Stdio::piped())
    1360              :             .stderr(Stdio::piped())
    1361              :             .spawn()
    1362              :             .expect("postgres --sync-safekeepers failed to start");
    1363              :         SYNC_SAFEKEEPERS_PID.store(sync_handle.id(), Ordering::SeqCst);
    1364              : 
    1365              :         // `postgres --sync-safekeepers` will print all log output to stderr and
    1366              :         // final LSN to stdout. So we leave stdout to collect LSN, while stderr logs
    1367              :         // will be collected in a child thread.
    1368              :         let stderr = sync_handle
    1369              :             .stderr
    1370              :             .take()
    1371              :             .expect("stderr should be captured");
    1372              :         let logs_handle = handle_postgres_logs(stderr);
    1373              : 
    1374              :         let sync_output = sync_handle
    1375              :             .wait_with_output()
    1376              :             .expect("postgres --sync-safekeepers failed");
    1377              :         SYNC_SAFEKEEPERS_PID.store(0, Ordering::SeqCst);
    1378              : 
    1379              :         // Process has exited, so we can join the logs thread.
    1380              :         let _ = tokio::runtime::Handle::current()
    1381              :             .block_on(logs_handle)
    1382            0 :             .map_err(|e| tracing::error!("log task panicked: {:?}", e));
    1383              : 
    1384              :         if !sync_output.status.success() {
    1385              :             anyhow::bail!(
    1386              :                 "postgres --sync-safekeepers exited with non-zero status: {}. stdout: {}",
    1387              :                 sync_output.status,
    1388              :                 String::from_utf8(sync_output.stdout)
    1389              :                     .expect("postgres --sync-safekeepers exited, and stdout is not utf-8"),
    1390              :             );
    1391              :         }
    1392              : 
    1393              :         self.state.lock().unwrap().metrics.sync_safekeepers_ms = Utc::now()
    1394              :             .signed_duration_since(start_time)
    1395              :             .to_std()
    1396              :             .unwrap()
    1397              :             .as_millis() as u64;
    1398              : 
    1399              :         let lsn = Lsn::from_str(String::from_utf8(sync_output.stdout)?.trim())?;
    1400              : 
    1401              :         Ok(lsn)
    1402              :     }
    1403              : 
    1404              :     /// Do all the preparations like PGDATA directory creation, configuration,
    1405              :     /// safekeepers sync, basebackup, etc.
    1406              :     #[instrument(skip_all)]
    1407              :     pub fn prepare_pgdata(&self, compute_state: &ComputeState) -> Result<()> {
    1408              :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
    1409              :         let spec = &pspec.spec;
    1410              :         let pgdata_path = Path::new(&self.params.pgdata);
    1411              : 
    1412              :         let tls_config = self.tls_config(&pspec.spec);
    1413              : 
    1414              :         // Remove/create an empty pgdata directory and put configuration there.
    1415              :         self.create_pgdata()?;
    1416              :         config::write_postgres_conf(
    1417              :             pgdata_path,
    1418              :             &self.params,
    1419              :             &pspec.spec,
    1420              :             self.params.internal_http_port,
    1421              :             tls_config,
    1422              :         )?;
    1423              : 
    1424              :         // Syncing safekeepers is only safe with primary nodes: if a primary
    1425              :         // is already connected it will be kicked out, so a secondary (standby)
    1426              :         // cannot sync safekeepers.
    1427              :         let lsn = match spec.mode {
    1428              :             ComputeMode::Primary => {
    1429              :                 info!("checking if safekeepers are synced");
    1430              :                 let lsn = if let Ok(Some(lsn)) = self.check_safekeepers_synced(compute_state) {
    1431              :                     lsn
    1432              :                 } else {
    1433              :                     info!("starting safekeepers syncing");
    1434              :                     self.sync_safekeepers(pspec.storage_auth_token.clone())
    1435              :                         .with_context(|| "failed to sync safekeepers")?
    1436              :                 };
    1437              :                 info!("safekeepers synced at LSN {}", lsn);
    1438              :                 lsn
    1439              :             }
    1440              :             ComputeMode::Static(lsn) => {
    1441              :                 info!("Starting read-only node at static LSN {}", lsn);
    1442              :                 lsn
    1443              :             }
    1444              :             ComputeMode::Replica => {
    1445              :                 info!("Initializing standby from latest Pageserver LSN");
    1446              :                 Lsn(0)
    1447              :             }
    1448              :         };
    1449              : 
    1450              :         info!(
    1451              :             "getting basebackup@{} from pageserver {}",
    1452              :             lsn, &pspec.pageserver_connstr
    1453              :         );
    1454            0 :         self.get_basebackup(compute_state, lsn).with_context(|| {
    1455            0 :             format!(
    1456            0 :                 "failed to get basebackup@{} from pageserver {}",
    1457            0 :                 lsn, &pspec.pageserver_connstr
    1458              :             )
    1459            0 :         })?;
    1460              : 
    1461              :         // Update pg_hba.conf received with basebackup.
    1462              :         update_pg_hba(pgdata_path, None)?;
    1463              : 
    1464              :         // Place pg_dynshmem under /dev/shm. This allows us to use
    1465              :         // 'dynamic_shared_memory_type = mmap' so that the files are placed in
    1466              :         // /dev/shm, similar to how 'dynamic_shared_memory_type = posix' works.
    1467              :         //
    1468              :         // Why on earth don't we just stick to the 'posix' default, you might
    1469              :         // ask.  It turns out that making large allocations with 'posix' doesn't
    1470              :         // work very well with autoscaling. The behavior we want is that:
    1471              :         //
    1472              :         // 1. You can make large DSM allocations, larger than the current RAM
    1473              :         //    size of the VM, without errors
    1474              :         //
    1475              :         // 2. If the allocated memory is really used, the VM is scaled up
    1476              :         //    automatically to accommodate that
    1477              :         //
    1478              :         // We try to make that possible by having swap in the VM. But with the
    1479              :         // default 'posix' DSM implementation, we fail step 1, even when there's
    1480              :         // plenty of swap available. PostgreSQL uses posix_fallocate() to create
    1481              :         // the shmem segment, which is really just a file in /dev/shm in Linux,
    1482              :         // but posix_fallocate() on tmpfs returns ENOMEM if the size is larger
    1483              :         // than available RAM.
    1484              :         //
    1485              :         // Using 'dynamic_shared_memory_type = mmap' works around that, because
    1486              :         // the Postgres 'mmap' DSM implementation doesn't use
    1487              :         // posix_fallocate(). Instead, it uses repeated calls to write(2) to
    1488              :         // fill the file with zeros. It's weird that that differs between
    1489              :         // 'posix' and 'mmap', but we take advantage of it. When the file is
    1490              :         // filled slowly with write(2), the kernel allows it to grow larger, as
    1491              :         // long as there's swap available.
    1492              :         //
    1493              :         // In short, using 'dynamic_shared_memory_type = mmap' allows us one DSM
    1494              :         // segment to be larger than currently available RAM. But because we
    1495              :         // don't want to store it on a real file, which the kernel would try to
    1496              :         // flush to disk, so symlink pg_dynshm to /dev/shm.
    1497              :         //
    1498              :         // We don't set 'dynamic_shared_memory_type = mmap' here, we let the
    1499              :         // control plane control that option. If 'mmap' is not used, this
    1500              :         // symlink doesn't affect anything.
    1501              :         //
    1502              :         // See https://github.com/neondatabase/autoscaling/issues/800
    1503              :         std::fs::remove_dir(pgdata_path.join("pg_dynshmem"))?;
    1504              :         symlink("/dev/shm/", pgdata_path.join("pg_dynshmem"))?;
    1505              : 
    1506              :         match spec.mode {
    1507              :             ComputeMode::Primary => {}
    1508              :             ComputeMode::Replica | ComputeMode::Static(..) => {
    1509              :                 add_standby_signal(pgdata_path)?;
    1510              :             }
    1511              :         }
    1512              : 
    1513              :         Ok(())
    1514              :     }
    1515              : 
    1516              :     /// Start and stop a postgres process to warm up the VM for startup.
    1517            0 :     pub fn prewarm_postgres_vm_memory(&self) -> Result<()> {
    1518            0 :         info!("prewarming VM memory");
    1519              : 
    1520              :         // Create pgdata
    1521            0 :         let pgdata = &format!("{}.warmup", self.params.pgdata);
    1522            0 :         create_pgdata(pgdata)?;
    1523              : 
    1524              :         // Run initdb to completion
    1525            0 :         info!("running initdb");
    1526            0 :         let initdb_bin = Path::new(&self.params.pgbin)
    1527            0 :             .parent()
    1528            0 :             .unwrap()
    1529            0 :             .join("initdb");
    1530            0 :         Command::new(initdb_bin)
    1531            0 :             .args(["--pgdata", pgdata])
    1532            0 :             .output()
    1533            0 :             .expect("cannot start initdb process");
    1534              : 
    1535              :         // Write conf
    1536              :         use std::io::Write;
    1537            0 :         let conf_path = Path::new(pgdata).join("postgresql.conf");
    1538            0 :         let mut file = std::fs::File::create(conf_path)?;
    1539            0 :         writeln!(file, "shared_buffers=65536")?;
    1540            0 :         writeln!(file, "port=51055")?; // Nobody should be connecting
    1541            0 :         writeln!(file, "shared_preload_libraries = 'neon'")?;
    1542              : 
    1543              :         // Start postgres
    1544            0 :         info!("starting postgres");
    1545            0 :         let mut pg = maybe_cgexec(&self.params.pgbin)
    1546            0 :             .args(["-D", pgdata])
    1547            0 :             .spawn()
    1548            0 :             .expect("cannot start postgres process");
    1549              : 
    1550              :         // Stop it when it's ready
    1551            0 :         info!("waiting for postgres");
    1552            0 :         wait_for_postgres(&mut pg, Path::new(pgdata))?;
    1553              :         // SIGQUIT orders postgres to exit immediately. We don't want to SIGKILL
    1554              :         // it to avoid orphaned processes prowling around while datadir is
    1555              :         // wiped.
    1556            0 :         let pm_pid = Pid::from_raw(pg.id() as i32);
    1557            0 :         kill(pm_pid, Signal::SIGQUIT)?;
    1558            0 :         info!("sent SIGQUIT signal");
    1559            0 :         pg.wait()?;
    1560            0 :         info!("done prewarming vm memory");
    1561              : 
    1562              :         // clean up
    1563            0 :         let _ok = fs::remove_dir_all(pgdata);
    1564            0 :         Ok(())
    1565            0 :     }
    1566              : 
    1567              :     /// Start Postgres as a child process and wait for it to start accepting
    1568              :     /// connections.
    1569              :     ///
    1570              :     /// Returns a handle to the child process and a handle to the logs thread.
    1571              :     #[instrument(skip_all)]
    1572              :     pub fn start_postgres(&self, storage_auth_token: Option<String>) -> Result<PostgresHandle> {
    1573              :         let pgdata_path = Path::new(&self.params.pgdata);
    1574              : 
    1575              :         // Run postgres as a child process.
    1576              :         let mut pg = maybe_cgexec(&self.params.pgbin)
    1577              :             .args(["-D", &self.params.pgdata])
    1578              :             .envs(if let Some(storage_auth_token) = &storage_auth_token {
    1579              :                 vec![("NEON_AUTH_TOKEN", storage_auth_token)]
    1580              :             } else {
    1581              :                 vec![]
    1582              :             })
    1583              :             .stderr(Stdio::piped())
    1584              :             .spawn()
    1585              :             .expect("cannot start postgres process");
    1586              :         PG_PID.store(pg.id(), Ordering::SeqCst);
    1587              : 
    1588              :         // Start a task to collect logs from stderr.
    1589              :         let stderr = pg.stderr.take().expect("stderr should be captured");
    1590              :         let logs_handle = handle_postgres_logs(stderr);
    1591              : 
    1592              :         wait_for_postgres(&mut pg, pgdata_path)?;
    1593              : 
    1594              :         Ok(PostgresHandle {
    1595              :             postgres: pg,
    1596              :             log_collector: logs_handle,
    1597              :         })
    1598              :     }
    1599              : 
    1600              :     /// Wait for the child Postgres process forever. In this state Ctrl+C will
    1601              :     /// propagate to Postgres and it will be shut down as well.
    1602            0 :     fn wait_postgres(&self, mut pg_handle: PostgresHandle) -> std::process::ExitStatus {
    1603            0 :         info!(postmaster_pid = %pg_handle.postgres.id(), "Waiting for Postgres to exit");
    1604              : 
    1605            0 :         let ecode = pg_handle
    1606            0 :             .postgres
    1607            0 :             .wait()
    1608            0 :             .expect("failed to start waiting on Postgres process");
    1609            0 :         PG_PID.store(0, Ordering::SeqCst);
    1610              : 
    1611              :         // Process has exited. Wait for the log collecting task to finish.
    1612            0 :         let _ = tokio::runtime::Handle::current()
    1613            0 :             .block_on(pg_handle.log_collector)
    1614            0 :             .map_err(|e| tracing::error!("log task panicked: {:?}", e));
    1615              : 
    1616            0 :         ecode
    1617            0 :     }
    1618              : 
    1619              :     /// Do post configuration of the already started Postgres. This function spawns a background task to
    1620              :     /// configure the database after applying the compute spec. Currently, it upgrades the neon extension
    1621              :     /// version. In the future, it may upgrade all 3rd-party extensions.
    1622              :     #[instrument(skip_all)]
    1623              :     pub fn post_apply_config(&self) -> Result<()> {
    1624              :         let conf = self.get_tokio_conn_conf(Some("compute_ctl:post_apply_config"));
    1625            0 :         tokio::spawn(async move {
    1626            0 :             let res = async {
    1627            0 :                 let (mut client, connection) = conf.connect(NoTls).await?;
    1628            0 :                 tokio::spawn(async move {
    1629            0 :                     if let Err(e) = connection.await {
    1630            0 :                         eprintln!("connection error: {e}");
    1631            0 :                     }
    1632            0 :                 });
    1633              : 
    1634            0 :                 handle_neon_extension_upgrade(&mut client)
    1635            0 :                     .await
    1636            0 :                     .context("handle_neon_extension_upgrade")?;
    1637            0 :                 Ok::<_, anyhow::Error>(())
    1638            0 :             }
    1639            0 :             .await;
    1640            0 :             if let Err(err) = res {
    1641            0 :                 error!("error while post_apply_config: {err:#}");
    1642            0 :             }
    1643            0 :         });
    1644              :         Ok(())
    1645              :     }
    1646              : 
    1647            0 :     pub fn get_conn_conf(&self, application_name: Option<&str>) -> postgres::Config {
    1648            0 :         let mut conf = self.conn_conf.clone();
    1649            0 :         if let Some(application_name) = application_name {
    1650            0 :             conf.application_name(application_name);
    1651            0 :         }
    1652            0 :         conf
    1653            0 :     }
    1654              : 
    1655            0 :     pub fn get_tokio_conn_conf(&self, application_name: Option<&str>) -> tokio_postgres::Config {
    1656            0 :         let mut conf = self.tokio_conn_conf.clone();
    1657            0 :         if let Some(application_name) = application_name {
    1658            0 :             conf.application_name(application_name);
    1659            0 :         }
    1660            0 :         conf
    1661            0 :     }
    1662              : 
    1663            0 :     pub async fn get_maintenance_client(
    1664            0 :         conf: &tokio_postgres::Config,
    1665            0 :     ) -> Result<tokio_postgres::Client> {
    1666            0 :         let mut conf = conf.clone();
    1667            0 :         conf.application_name("compute_ctl:apply_config");
    1668              : 
    1669            0 :         let (client, conn) = match conf.connect(NoTls).await {
    1670              :             // If connection fails, it may be the old node with `zenith_admin` superuser.
    1671              :             //
    1672              :             // In this case we need to connect with old `zenith_admin` name
    1673              :             // and create new user. We cannot simply rename connected user,
    1674              :             // but we can create a new one and grant it all privileges.
    1675            0 :             Err(e) => match e.code() {
    1676              :                 Some(&SqlState::INVALID_PASSWORD)
    1677              :                 | Some(&SqlState::INVALID_AUTHORIZATION_SPECIFICATION) => {
    1678              :                     // Connect with `zenith_admin` if `cloud_admin` could not authenticate
    1679            0 :                     info!(
    1680            0 :                         "cannot connect to Postgres: {}, retrying with 'zenith_admin' username",
    1681              :                         e
    1682              :                     );
    1683            0 :                     let mut zenith_admin_conf = postgres::config::Config::from(conf.clone());
    1684            0 :                     zenith_admin_conf.application_name("compute_ctl:apply_config");
    1685            0 :                     zenith_admin_conf.user("zenith_admin");
    1686              : 
    1687              :                     // It doesn't matter what were the options before, here we just want
    1688              :                     // to connect and create a new superuser role.
    1689              :                     const ZENITH_OPTIONS: &str = "-c role=zenith_admin -c default_transaction_read_only=off -c search_path=public -c statement_timeout=0";
    1690            0 :                     zenith_admin_conf.options(ZENITH_OPTIONS);
    1691              : 
    1692            0 :                     let mut client =
    1693            0 :                         zenith_admin_conf.connect(NoTls)
    1694            0 :                             .context("broken cloud_admin credential: tried connecting with cloud_admin but could not authenticate, and zenith_admin does not work either")?;
    1695              : 
    1696              :                     // Disable forwarding so that users don't get a cloud_admin role
    1697            0 :                     let mut func = || {
    1698            0 :                         client.simple_query("SET neon.forward_ddl = false")?;
    1699            0 :                         client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?;
    1700            0 :                         client.simple_query("GRANT zenith_admin TO cloud_admin")?;
    1701            0 :                         Ok::<_, anyhow::Error>(())
    1702            0 :                     };
    1703            0 :                     func().context("apply_config setup cloud_admin")?;
    1704              : 
    1705            0 :                     drop(client);
    1706              : 
    1707              :                     // Reconnect with connstring with expected name
    1708            0 :                     conf.connect(NoTls).await?
    1709              :                 }
    1710            0 :                 _ => return Err(e.into()),
    1711              :             },
    1712            0 :             Ok((client, conn)) => (client, conn),
    1713              :         };
    1714              : 
    1715            0 :         spawn(async move {
    1716            0 :             if let Err(e) = conn.await {
    1717            0 :                 error!("maintenance client connection error: {}", e);
    1718            0 :             }
    1719            0 :         });
    1720              : 
    1721              :         // Disable DDL forwarding because control plane already knows about the roles/databases
    1722              :         // we're about to modify.
    1723            0 :         client
    1724            0 :             .simple_query("SET neon.forward_ddl = false")
    1725            0 :             .await
    1726            0 :             .context("apply_config SET neon.forward_ddl = false")?;
    1727              : 
    1728            0 :         Ok(client)
    1729            0 :     }
    1730              : 
    1731              :     /// Do initial configuration of the already started Postgres.
    1732              :     #[instrument(skip_all)]
    1733              :     pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
    1734              :         let conf = self.get_tokio_conn_conf(Some("compute_ctl:apply_config"));
    1735              : 
    1736              :         let conf = Arc::new(conf);
    1737              :         let spec = Arc::new(
    1738              :             compute_state
    1739              :                 .pspec
    1740              :                 .as_ref()
    1741              :                 .expect("spec must be set")
    1742              :                 .spec
    1743              :                 .clone(),
    1744              :         );
    1745              : 
    1746              :         let mut tls_config = None::<TlsConfig>;
    1747              :         if spec.features.contains(&ComputeFeature::TlsExperimental) {
    1748              :             tls_config = self.compute_ctl_config.tls.clone();
    1749              :         }
    1750              : 
    1751              :         self.update_installed_extensions_collection_interval(&spec);
    1752              : 
    1753              :         let max_concurrent_connections = self.max_service_connections(compute_state, &spec);
    1754              : 
    1755              :         // Merge-apply spec & changes to PostgreSQL state.
    1756              :         self.apply_spec_sql(spec.clone(), conf.clone(), max_concurrent_connections)?;
    1757              : 
    1758              :         if let Some(local_proxy) = &spec.clone().local_proxy_config {
    1759              :             let mut local_proxy = local_proxy.clone();
    1760              :             local_proxy.tls = tls_config.clone();
    1761              : 
    1762              :             info!("configuring local_proxy");
    1763              :             local_proxy::configure(&local_proxy).context("apply_config local_proxy")?;
    1764              :         }
    1765              : 
    1766              :         // Run migrations separately to not hold up cold starts
    1767              :         let lakebase_mode = self.params.lakebase_mode;
    1768              :         let params = self.params.clone();
    1769            0 :         tokio::spawn(async move {
    1770            0 :             let mut conf = conf.as_ref().clone();
    1771            0 :             conf.application_name("compute_ctl:migrations");
    1772              : 
    1773            0 :             match conf.connect(NoTls).await {
    1774            0 :                 Ok((mut client, connection)) => {
    1775            0 :                     tokio::spawn(async move {
    1776            0 :                         if let Err(e) = connection.await {
    1777            0 :                             eprintln!("connection error: {e}");
    1778            0 :                         }
    1779            0 :                     });
    1780            0 :                     if let Err(e) = handle_migrations(params, &mut client, lakebase_mode).await {
    1781            0 :                         error!("Failed to run migrations: {}", e);
    1782            0 :                     }
    1783              :                 }
    1784            0 :                 Err(e) => {
    1785            0 :                     error!(
    1786            0 :                         "Failed to connect to the compute for running migrations: {}",
    1787              :                         e
    1788              :                     );
    1789              :                 }
    1790              :             };
    1791            0 :         });
    1792              : 
    1793              :         Ok::<(), anyhow::Error>(())
    1794              :     }
    1795              : 
    1796              :     // Signal to the configurator to refresh the configuration by pulling a new spec from the HCC.
    1797              :     // Note that this merely triggers a notification on a condition variable the configurator thread
    1798              :     // waits on. The configurator thread (in configurator.rs) pulls the new spec from the HCC and
    1799              :     // applies it.
    1800            0 :     pub async fn signal_refresh_configuration(&self) -> Result<()> {
    1801            0 :         let states_allowing_configuration_refresh = [
    1802            0 :             ComputeStatus::Running,
    1803            0 :             ComputeStatus::Failed,
    1804            0 :             ComputeStatus::RefreshConfigurationPending,
    1805            0 :         ];
    1806              : 
    1807            0 :         let mut state = self.state.lock().expect("state lock poisoned");
    1808            0 :         if states_allowing_configuration_refresh.contains(&state.status) {
    1809            0 :             state.status = ComputeStatus::RefreshConfigurationPending;
    1810            0 :             self.state_changed.notify_all();
    1811            0 :             Ok(())
    1812            0 :         } else if state.status == ComputeStatus::Init {
    1813              :             // If the compute is in Init state, we can't refresh the configuration immediately,
    1814              :             // but we should be able to do that soon.
    1815            0 :             Ok(())
    1816              :         } else {
    1817            0 :             Err(anyhow::anyhow!(
    1818            0 :                 "Cannot refresh compute configuration in state {:?}",
    1819            0 :                 state.status
    1820            0 :             ))
    1821              :         }
    1822            0 :     }
    1823              : 
    1824              :     // Wrapped this around `pg_ctl reload`, but right now we don't use
    1825              :     // `pg_ctl` for start / stop.
    1826              :     #[instrument(skip_all)]
    1827              :     fn pg_reload_conf(&self) -> Result<()> {
    1828              :         let pgctl_bin = Path::new(&self.params.pgbin)
    1829              :             .parent()
    1830              :             .unwrap()
    1831              :             .join("pg_ctl");
    1832              :         Command::new(pgctl_bin)
    1833              :             .args(["reload", "-D", &self.params.pgdata])
    1834              :             .output()
    1835              :             .expect("cannot run pg_ctl process");
    1836              :         Ok(())
    1837              :     }
    1838              : 
    1839              :     /// Similar to `apply_config()`, but does a bit different sequence of operations,
    1840              :     /// as it's used to reconfigure a previously started and configured Postgres node.
    1841              :     #[instrument(skip_all)]
    1842              :     pub fn reconfigure(&self) -> Result<()> {
    1843              :         let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
    1844              : 
    1845              :         let tls_config = self.tls_config(&spec);
    1846              : 
    1847              :         self.update_installed_extensions_collection_interval(&spec);
    1848              : 
    1849              :         if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
    1850              :             info!("tuning pgbouncer");
    1851              : 
    1852              :             let pgbouncer_settings = pgbouncer_settings.clone();
    1853              :             let tls_config = tls_config.clone();
    1854              : 
    1855              :             // Spawn a background task to do the tuning,
    1856              :             // so that we don't block the main thread that starts Postgres.
    1857            0 :             tokio::spawn(async move {
    1858            0 :                 let res = tune_pgbouncer(pgbouncer_settings, tls_config).await;
    1859            0 :                 if let Err(err) = res {
    1860            0 :                     error!("error while tuning pgbouncer: {err:?}");
    1861            0 :                 }
    1862            0 :             });
    1863              :         }
    1864              : 
    1865              :         if let Some(ref local_proxy) = spec.local_proxy_config {
    1866              :             info!("configuring local_proxy");
    1867              : 
    1868              :             // Spawn a background task to do the configuration,
    1869              :             // so that we don't block the main thread that starts Postgres.
    1870              :             let mut local_proxy = local_proxy.clone();
    1871              :             local_proxy.tls = tls_config.clone();
    1872            0 :             tokio::spawn(async move {
    1873            0 :                 if let Err(err) = local_proxy::configure(&local_proxy) {
    1874            0 :                     error!("error while configuring local_proxy: {err:?}");
    1875            0 :                 }
    1876            0 :             });
    1877              :         }
    1878              : 
    1879              :         // Reconfigure rsyslog for Postgres logs export
    1880              :         let conf = PostgresLogsRsyslogConfig::new(spec.logs_export_host.as_deref());
    1881              :         configure_postgres_logs_export(conf)?;
    1882              : 
    1883              :         // Write new config
    1884              :         let pgdata_path = Path::new(&self.params.pgdata);
    1885              :         config::write_postgres_conf(
    1886              :             pgdata_path,
    1887              :             &self.params,
    1888              :             &spec,
    1889              :             self.params.internal_http_port,
    1890              :             tls_config,
    1891              :         )?;
    1892              : 
    1893              :         self.pg_reload_conf()?;
    1894              : 
    1895              :         if !spec.skip_pg_catalog_updates {
    1896              :             let max_concurrent_connections = spec.reconfigure_concurrency;
    1897              :             // Temporarily reset max_cluster_size in config
    1898              :             // to avoid the possibility of hitting the limit, while we are reconfiguring:
    1899              :             // creating new extensions, roles, etc.
    1900            0 :             config::with_compute_ctl_tmp_override(pgdata_path, "neon.max_cluster_size=-1", || {
    1901            0 :                 self.pg_reload_conf()?;
    1902              : 
    1903            0 :                 if spec.mode == ComputeMode::Primary {
    1904            0 :                     let conf = self.get_tokio_conn_conf(Some("compute_ctl:reconfigure"));
    1905            0 :                     let conf = Arc::new(conf);
    1906              : 
    1907            0 :                     let spec = Arc::new(spec.clone());
    1908              : 
    1909            0 :                     self.apply_spec_sql(spec, conf, max_concurrent_connections)?;
    1910            0 :                 }
    1911              : 
    1912            0 :                 Ok(())
    1913            0 :             })?;
    1914              :             self.pg_reload_conf()?;
    1915              :         }
    1916              : 
    1917              :         let unknown_op = "unknown".to_string();
    1918              :         let op_id = spec.operation_uuid.as_ref().unwrap_or(&unknown_op);
    1919              :         info!(
    1920              :             "finished reconfiguration of compute node for operation {}",
    1921              :             op_id
    1922              :         );
    1923              : 
    1924              :         Ok(())
    1925              :     }
    1926              : 
    1927              :     #[instrument(skip_all)]
    1928              :     pub fn configure_as_primary(&self, compute_state: &ComputeState) -> Result<()> {
    1929              :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
    1930              : 
    1931              :         assert!(pspec.spec.mode == ComputeMode::Primary);
    1932              :         if !pspec.spec.skip_pg_catalog_updates {
    1933              :             let pgdata_path = Path::new(&self.params.pgdata);
    1934              :             // temporarily reset max_cluster_size in config
    1935              :             // to avoid the possibility of hitting the limit, while we are applying config:
    1936              :             // creating new extensions, roles, etc...
    1937            0 :             config::with_compute_ctl_tmp_override(pgdata_path, "neon.max_cluster_size=-1", || {
    1938            0 :                 self.pg_reload_conf()?;
    1939              : 
    1940            0 :                 self.apply_config(compute_state)?;
    1941              : 
    1942            0 :                 Ok(())
    1943            0 :             })?;
    1944              : 
    1945              :             let postgresql_conf_path = pgdata_path.join("postgresql.conf");
    1946              :             if config::line_in_file(
    1947              :                 &postgresql_conf_path,
    1948              :                 "neon.disable_logical_replication_subscribers=false",
    1949              :             )? {
    1950              :                 info!(
    1951              :                     "updated postgresql.conf to set neon.disable_logical_replication_subscribers=false"
    1952              :                 );
    1953              :             }
    1954              :             self.pg_reload_conf()?;
    1955              :         }
    1956              :         self.post_apply_config()?;
    1957              : 
    1958              :         Ok(())
    1959              :     }
    1960              : 
    1961            0 :     pub async fn watch_cert_for_changes(self: Arc<Self>) {
    1962              :         // update status on cert renewal
    1963            0 :         if let Some(tls_config) = &self.compute_ctl_config.tls {
    1964            0 :             let tls_config = tls_config.clone();
    1965              : 
    1966              :             // wait until the cert exists.
    1967            0 :             let mut cert_watch = watch_cert_for_changes(tls_config.cert_path.clone()).await;
    1968              : 
    1969            0 :             tokio::task::spawn_blocking(move || {
    1970            0 :                 let handle = tokio::runtime::Handle::current();
    1971              :                 'cert_update: loop {
    1972              :                     // let postgres/pgbouncer/local_proxy know the new cert/key exists.
    1973              :                     // we need to wait until it's configurable first.
    1974              : 
    1975            0 :                     let mut state = self.state.lock().unwrap();
    1976              :                     'status_update: loop {
    1977            0 :                         match state.status {
    1978              :                             // let's update the state to config pending
    1979              :                             ComputeStatus::ConfigurationPending | ComputeStatus::Running => {
    1980            0 :                                 state.set_status(
    1981            0 :                                     ComputeStatus::ConfigurationPending,
    1982            0 :                                     &self.state_changed,
    1983            0 :                                 );
    1984            0 :                                 break 'status_update;
    1985              :                             }
    1986              : 
    1987              :                             // exit loop
    1988              :                             ComputeStatus::Failed
    1989              :                             | ComputeStatus::TerminationPendingFast
    1990              :                             | ComputeStatus::TerminationPendingImmediate
    1991            0 :                             | ComputeStatus::Terminated => break 'cert_update,
    1992              : 
    1993              :                             // wait
    1994              :                             ComputeStatus::Init
    1995              :                             | ComputeStatus::Configuration
    1996              :                             | ComputeStatus::RefreshConfigurationPending
    1997            0 :                             | ComputeStatus::Empty => {
    1998            0 :                                 state = self.state_changed.wait(state).unwrap();
    1999            0 :                             }
    2000              :                         }
    2001              :                     }
    2002            0 :                     drop(state);
    2003              : 
    2004              :                     // wait for a new certificate update
    2005            0 :                     if handle.block_on(cert_watch.changed()).is_err() {
    2006            0 :                         break;
    2007            0 :                     }
    2008              :                 }
    2009            0 :             });
    2010            0 :         }
    2011            0 :     }
    2012              : 
    2013            0 :     pub fn tls_config(&self, spec: &ComputeSpec) -> &Option<TlsConfig> {
    2014            0 :         if spec.features.contains(&ComputeFeature::TlsExperimental) {
    2015            0 :             &self.compute_ctl_config.tls
    2016              :         } else {
    2017            0 :             &None::<TlsConfig>
    2018              :         }
    2019            0 :     }
    2020              : 
    2021              :     /// Update the `last_active` in the shared state, but ensure that it's a more recent one.
    2022            0 :     pub fn update_last_active(&self, last_active: Option<DateTime<Utc>>) {
    2023            0 :         let mut state = self.state.lock().unwrap();
    2024              :         // NB: `Some(<DateTime>)` is always greater than `None`.
    2025            0 :         if last_active > state.last_active {
    2026            0 :             state.last_active = last_active;
    2027            0 :             debug!("set the last compute activity time to: {:?}", last_active);
    2028            0 :         }
    2029            0 :     }
    2030              : 
    2031              :     // Look for core dumps and collect backtraces.
    2032              :     //
    2033              :     // EKS worker nodes have following core dump settings:
    2034              :     //   /proc/sys/kernel/core_pattern -> core
    2035              :     //   /proc/sys/kernel/core_uses_pid -> 1
    2036              :     //   ulimit -c -> unlimited
    2037              :     // which results in core dumps being written to postgres data directory as core.<pid>.
    2038              :     //
    2039              :     // Use that as a default location and pattern, except macos where core dumps are written
    2040              :     // to /cores/ directory by default.
    2041              :     //
    2042              :     // With default Linux settings, the core dump file is called just "core", so check for
    2043              :     // that too.
    2044            0 :     pub fn check_for_core_dumps(&self) -> Result<()> {
    2045            0 :         let core_dump_dir = match std::env::consts::OS {
    2046            0 :             "macos" => Path::new("/cores/"),
    2047            0 :             _ => Path::new(&self.params.pgdata),
    2048              :         };
    2049              : 
    2050              :         // Collect core dump paths if any
    2051            0 :         info!("checking for core dumps in {}", core_dump_dir.display());
    2052            0 :         let files = fs::read_dir(core_dump_dir)?;
    2053            0 :         let cores = files.filter_map(|entry| {
    2054            0 :             let entry = entry.ok()?;
    2055              : 
    2056            0 :             let is_core_dump = match entry.file_name().to_str()? {
    2057            0 :                 n if n.starts_with("core.") => true,
    2058            0 :                 "core" => true,
    2059            0 :                 _ => false,
    2060              :             };
    2061            0 :             if is_core_dump {
    2062            0 :                 Some(entry.path())
    2063              :             } else {
    2064            0 :                 None
    2065              :             }
    2066            0 :         });
    2067              : 
    2068              :         // Print backtrace for each core dump
    2069            0 :         for core_path in cores {
    2070            0 :             warn!(
    2071            0 :                 "core dump found: {}, collecting backtrace",
    2072            0 :                 core_path.display()
    2073              :             );
    2074              : 
    2075              :             // Try first with gdb
    2076            0 :             let backtrace = Command::new("gdb")
    2077            0 :                 .args(["--batch", "-q", "-ex", "bt", &self.params.pgbin])
    2078            0 :                 .arg(&core_path)
    2079            0 :                 .output();
    2080              : 
    2081              :             // Try lldb if no gdb is found -- that is handy for local testing on macOS
    2082            0 :             let backtrace = match backtrace {
    2083            0 :                 Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
    2084            0 :                     warn!("cannot find gdb, trying lldb");
    2085            0 :                     Command::new("lldb")
    2086            0 :                         .arg("-c")
    2087            0 :                         .arg(&core_path)
    2088            0 :                         .args(["--batch", "-o", "bt all", "-o", "quit"])
    2089            0 :                         .output()
    2090              :                 }
    2091            0 :                 _ => backtrace,
    2092            0 :             }?;
    2093              : 
    2094            0 :             warn!(
    2095            0 :                 "core dump backtrace: {}",
    2096            0 :                 String::from_utf8_lossy(&backtrace.stdout)
    2097              :             );
    2098            0 :             warn!(
    2099            0 :                 "debugger stderr: {}",
    2100            0 :                 String::from_utf8_lossy(&backtrace.stderr)
    2101              :             );
    2102              :         }
    2103              : 
    2104            0 :         Ok(())
    2105            0 :     }
    2106              : 
    2107              :     /// Select `pg_stat_statements` data and return it as a stringified JSON
    2108            0 :     pub async fn collect_insights(&self) -> String {
    2109            0 :         let mut result_rows: Vec<String> = Vec::new();
    2110            0 :         let conf = self.get_tokio_conn_conf(Some("compute_ctl:collect_insights"));
    2111            0 :         let connect_result = conf.connect(NoTls).await;
    2112            0 :         let (client, connection) = connect_result.unwrap();
    2113            0 :         tokio::spawn(async move {
    2114            0 :             if let Err(e) = connection.await {
    2115            0 :                 eprintln!("connection error: {e}");
    2116            0 :             }
    2117            0 :         });
    2118            0 :         let result = client
    2119            0 :             .simple_query(
    2120            0 :                 "SELECT
    2121            0 :     row_to_json(pg_stat_statements)
    2122            0 : FROM
    2123            0 :     pg_stat_statements
    2124            0 : WHERE
    2125            0 :     userid != 'cloud_admin'::regrole::oid
    2126            0 : ORDER BY
    2127            0 :     (mean_exec_time + mean_plan_time) DESC
    2128            0 : LIMIT 100",
    2129            0 :             )
    2130            0 :             .await;
    2131              : 
    2132            0 :         if let Ok(raw_rows) = result {
    2133            0 :             for message in raw_rows.iter() {
    2134            0 :                 if let postgres::SimpleQueryMessage::Row(row) = message {
    2135            0 :                     if let Some(json) = row.get(0) {
    2136            0 :                         result_rows.push(json.to_string());
    2137            0 :                     }
    2138            0 :                 }
    2139              :             }
    2140              : 
    2141            0 :             format!("{{\"pg_stat_statements\": [{}]}}", result_rows.join(","))
    2142              :         } else {
    2143            0 :             "{{\"pg_stat_statements\": []}}".to_string()
    2144              :         }
    2145            0 :     }
    2146              : 
    2147              :     // download an archive, unzip and place files in correct locations
    2148            0 :     pub async fn download_extension(
    2149            0 :         &self,
    2150            0 :         real_ext_name: String,
    2151            0 :         ext_path: RemotePath,
    2152            0 :     ) -> Result<u64, DownloadError> {
    2153            0 :         let remote_ext_base_url =
    2154            0 :             self.params
    2155            0 :                 .remote_ext_base_url
    2156            0 :                 .as_ref()
    2157            0 :                 .ok_or(DownloadError::BadInput(anyhow::anyhow!(
    2158            0 :                     "Remote extensions storage is not configured",
    2159            0 :                 )))?;
    2160              : 
    2161            0 :         let ext_archive_name = ext_path.object_name().expect("bad path");
    2162              : 
    2163            0 :         let mut first_try = false;
    2164            0 :         if !self
    2165            0 :             .ext_download_progress
    2166            0 :             .read()
    2167            0 :             .expect("lock err")
    2168            0 :             .contains_key(ext_archive_name)
    2169            0 :         {
    2170            0 :             self.ext_download_progress
    2171            0 :                 .write()
    2172            0 :                 .expect("lock err")
    2173            0 :                 .insert(ext_archive_name.to_string(), (Utc::now(), false));
    2174            0 :             first_try = true;
    2175            0 :         }
    2176            0 :         let (download_start, download_completed) =
    2177            0 :             self.ext_download_progress.read().expect("lock err")[ext_archive_name];
    2178            0 :         let start_time_delta = Utc::now()
    2179            0 :             .signed_duration_since(download_start)
    2180            0 :             .to_std()
    2181            0 :             .unwrap()
    2182            0 :             .as_millis() as u64;
    2183              : 
    2184              :         // how long to wait for extension download if it was started by another process
    2185              :         const HANG_TIMEOUT: u64 = 3000; // milliseconds
    2186              : 
    2187            0 :         if download_completed {
    2188            0 :             info!("extension already downloaded, skipping re-download");
    2189            0 :             return Ok(0);
    2190            0 :         } else if start_time_delta < HANG_TIMEOUT && !first_try {
    2191            0 :             info!(
    2192            0 :                 "download {ext_archive_name} already started by another process, hanging untill completion or timeout"
    2193              :             );
    2194            0 :             let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500));
    2195              :             loop {
    2196            0 :                 info!("waiting for download");
    2197            0 :                 interval.tick().await;
    2198            0 :                 let (_, download_completed_now) =
    2199            0 :                     self.ext_download_progress.read().expect("lock")[ext_archive_name];
    2200            0 :                 if download_completed_now {
    2201            0 :                     info!("download finished by whoever else downloaded it");
    2202            0 :                     return Ok(0);
    2203            0 :                 }
    2204              :             }
    2205              :             // NOTE: the above loop will get terminated
    2206              :             // based on the timeout of the download function
    2207            0 :         }
    2208              : 
    2209              :         // if extension hasn't been downloaded before or the previous
    2210              :         // attempt to download was at least HANG_TIMEOUT ms ago
    2211              :         // then we try to download it here
    2212            0 :         info!("downloading new extension {ext_archive_name}");
    2213              : 
    2214            0 :         let download_size = extension_server::download_extension(
    2215            0 :             &real_ext_name,
    2216            0 :             &ext_path,
    2217            0 :             remote_ext_base_url,
    2218            0 :             &self.params.pgbin,
    2219            0 :         )
    2220            0 :         .await
    2221            0 :         .map_err(DownloadError::Other);
    2222              : 
    2223            0 :         if download_size.is_ok() {
    2224            0 :             self.ext_download_progress
    2225            0 :                 .write()
    2226            0 :                 .expect("bad lock")
    2227            0 :                 .insert(ext_archive_name.to_string(), (download_start, true));
    2228            0 :         }
    2229              : 
    2230            0 :         download_size
    2231            0 :     }
    2232              : 
    2233            0 :     pub async fn set_role_grants(
    2234            0 :         &self,
    2235            0 :         db_name: &PgIdent,
    2236            0 :         schema_name: &PgIdent,
    2237            0 :         privileges: &[Privilege],
    2238            0 :         role_name: &PgIdent,
    2239            0 :     ) -> Result<()> {
    2240              :         use tokio_postgres::NoTls;
    2241              : 
    2242            0 :         let mut conf = self.get_tokio_conn_conf(Some("compute_ctl:set_role_grants"));
    2243            0 :         conf.dbname(db_name);
    2244              : 
    2245            0 :         let (db_client, conn) = conf
    2246            0 :             .connect(NoTls)
    2247            0 :             .await
    2248            0 :             .context("Failed to connect to the database")?;
    2249            0 :         tokio::spawn(conn);
    2250              : 
    2251              :         // TODO: support other types of grants apart from schemas?
    2252              : 
    2253              :         // check the role grants first - to gracefully handle read-replicas.
    2254            0 :         let select = "SELECT privilege_type
    2255            0 :             FROM pg_namespace
    2256            0 :                 JOIN LATERAL (SELECT * FROM aclexplode(nspacl) AS x) acl ON true
    2257            0 :                 JOIN pg_user users ON acl.grantee = users.usesysid
    2258            0 :             WHERE users.usename = $1
    2259            0 :                 AND nspname = $2";
    2260            0 :         let rows = db_client
    2261            0 :             .query(select, &[role_name, schema_name])
    2262            0 :             .await
    2263            0 :             .with_context(|| format!("Failed to execute query: {select}"))?;
    2264              : 
    2265            0 :         let already_granted: HashSet<String> = rows.into_iter().map(|row| row.get(0)).collect();
    2266              : 
    2267            0 :         let grants = privileges
    2268            0 :             .iter()
    2269            0 :             .filter(|p| !already_granted.contains(p.as_str()))
    2270              :             // should not be quoted as it's part of the command.
    2271              :             // is already sanitized so it's ok
    2272            0 :             .map(|p| p.as_str())
    2273            0 :             .join(", ");
    2274              : 
    2275            0 :         if !grants.is_empty() {
    2276              :             // quote the schema and role name as identifiers to sanitize them.
    2277            0 :             let schema_name = schema_name.pg_quote();
    2278            0 :             let role_name = role_name.pg_quote();
    2279              : 
    2280            0 :             let query = format!("GRANT {grants} ON SCHEMA {schema_name} TO {role_name}",);
    2281            0 :             db_client
    2282            0 :                 .simple_query(&query)
    2283            0 :                 .await
    2284            0 :                 .with_context(|| format!("Failed to execute query: {query}"))?;
    2285            0 :         }
    2286              : 
    2287            0 :         Ok(())
    2288            0 :     }
    2289              : 
    2290            0 :     pub async fn install_extension(
    2291            0 :         &self,
    2292            0 :         ext_name: &PgIdent,
    2293            0 :         db_name: &PgIdent,
    2294            0 :         ext_version: ExtVersion,
    2295            0 :     ) -> Result<ExtVersion> {
    2296              :         use tokio_postgres::NoTls;
    2297              : 
    2298            0 :         let mut conf = self.get_tokio_conn_conf(Some("compute_ctl:install_extension"));
    2299            0 :         conf.dbname(db_name);
    2300              : 
    2301            0 :         let (db_client, conn) = conf
    2302            0 :             .connect(NoTls)
    2303            0 :             .await
    2304            0 :             .context("Failed to connect to the database")?;
    2305            0 :         tokio::spawn(conn);
    2306              : 
    2307            0 :         let version_query = "SELECT extversion FROM pg_extension WHERE extname = $1";
    2308            0 :         let version: Option<ExtVersion> = db_client
    2309            0 :             .query_opt(version_query, &[&ext_name])
    2310            0 :             .await
    2311            0 :             .with_context(|| format!("Failed to execute query: {version_query}"))?
    2312            0 :             .map(|row| row.get(0));
    2313              : 
    2314              :         // sanitize the inputs as postgres idents.
    2315            0 :         let ext_name: String = ext_name.pg_quote();
    2316            0 :         let quoted_version: String = ext_version.pg_quote();
    2317              : 
    2318            0 :         if let Some(installed_version) = version {
    2319            0 :             if installed_version == ext_version {
    2320            0 :                 return Ok(installed_version);
    2321            0 :             }
    2322            0 :             let query = format!("ALTER EXTENSION {ext_name} UPDATE TO {quoted_version}");
    2323            0 :             db_client
    2324            0 :                 .simple_query(&query)
    2325            0 :                 .await
    2326            0 :                 .with_context(|| format!("Failed to execute query: {query}"))?;
    2327              :         } else {
    2328            0 :             let query =
    2329            0 :                 format!("CREATE EXTENSION IF NOT EXISTS {ext_name} WITH VERSION {quoted_version}");
    2330            0 :             db_client
    2331            0 :                 .simple_query(&query)
    2332            0 :                 .await
    2333            0 :                 .with_context(|| format!("Failed to execute query: {query}"))?;
    2334              :         }
    2335              : 
    2336            0 :         Ok(ext_version)
    2337            0 :     }
    2338              : 
    2339            0 :     pub async fn prepare_preload_libraries(
    2340            0 :         &self,
    2341            0 :         spec: &ComputeSpec,
    2342            0 :     ) -> Result<RemoteExtensionMetrics> {
    2343            0 :         if self.params.remote_ext_base_url.is_none() {
    2344            0 :             return Ok(RemoteExtensionMetrics {
    2345            0 :                 num_ext_downloaded: 0,
    2346            0 :                 largest_ext_size: 0,
    2347            0 :                 total_ext_download_size: 0,
    2348            0 :             });
    2349            0 :         }
    2350            0 :         let remote_extensions = spec
    2351            0 :             .remote_extensions
    2352            0 :             .as_ref()
    2353            0 :             .ok_or(anyhow::anyhow!("Remote extensions are not configured"))?;
    2354              : 
    2355            0 :         info!("parse shared_preload_libraries from spec.cluster.settings");
    2356            0 :         let mut libs_vec = Vec::new();
    2357            0 :         if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
    2358            0 :             libs_vec = libs
    2359            0 :                 .split(&[',', '\'', ' '])
    2360            0 :                 .filter(|s| *s != "neon" && !s.is_empty())
    2361            0 :                 .map(str::to_string)
    2362            0 :                 .collect();
    2363            0 :         }
    2364            0 :         info!("parse shared_preload_libraries from provided postgresql.conf");
    2365              : 
    2366              :         // that is used in neon_local and python tests
    2367            0 :         if let Some(conf) = &spec.cluster.postgresql_conf {
    2368            0 :             let conf_lines = conf.split('\n').collect::<Vec<&str>>();
    2369            0 :             let mut shared_preload_libraries_line = "";
    2370            0 :             for line in conf_lines {
    2371            0 :                 if line.starts_with("shared_preload_libraries") {
    2372            0 :                     shared_preload_libraries_line = line;
    2373            0 :                 }
    2374              :             }
    2375            0 :             let mut preload_libs_vec = Vec::new();
    2376            0 :             if let Some(libs) = shared_preload_libraries_line.split("='").nth(1) {
    2377            0 :                 preload_libs_vec = libs
    2378            0 :                     .split(&[',', '\'', ' '])
    2379            0 :                     .filter(|s| *s != "neon" && !s.is_empty())
    2380            0 :                     .map(str::to_string)
    2381            0 :                     .collect();
    2382            0 :             }
    2383            0 :             libs_vec.extend(preload_libs_vec);
    2384            0 :         }
    2385              : 
    2386              :         // Don't try to download libraries that are not in the index.
    2387              :         // Assume that they are already present locally.
    2388            0 :         libs_vec.retain(|lib| remote_extensions.library_index.contains_key(lib));
    2389              : 
    2390            0 :         info!("Downloading to shared preload libraries: {:?}", &libs_vec);
    2391              : 
    2392            0 :         let mut download_tasks = Vec::new();
    2393            0 :         for library in &libs_vec {
    2394            0 :             let (ext_name, ext_path) =
    2395            0 :                 remote_extensions.get_ext(library, true, &BUILD_TAG, &self.params.pgversion)?;
    2396            0 :             download_tasks.push(self.download_extension(ext_name, ext_path));
    2397              :         }
    2398            0 :         let results = join_all(download_tasks).await;
    2399              : 
    2400            0 :         let mut remote_ext_metrics = RemoteExtensionMetrics {
    2401            0 :             num_ext_downloaded: 0,
    2402            0 :             largest_ext_size: 0,
    2403            0 :             total_ext_download_size: 0,
    2404            0 :         };
    2405            0 :         for result in results {
    2406            0 :             let download_size = match result {
    2407            0 :                 Ok(res) => {
    2408            0 :                     remote_ext_metrics.num_ext_downloaded += 1;
    2409            0 :                     res
    2410              :                 }
    2411            0 :                 Err(err) => {
    2412              :                     // if we failed to download an extension, we don't want to fail the whole
    2413              :                     // process, but we do want to log the error
    2414            0 :                     error!("Failed to download extension: {}", err);
    2415            0 :                     0
    2416              :                 }
    2417              :             };
    2418              : 
    2419            0 :             remote_ext_metrics.largest_ext_size =
    2420            0 :                 std::cmp::max(remote_ext_metrics.largest_ext_size, download_size);
    2421            0 :             remote_ext_metrics.total_ext_download_size += download_size;
    2422              :         }
    2423            0 :         Ok(remote_ext_metrics)
    2424            0 :     }
    2425              : 
    2426              :     /// Waits until current thread receives a state changed notification and
    2427              :     /// the pageserver connection strings has changed.
    2428              :     ///
    2429              :     /// The operation will time out after a specified duration.
    2430            0 :     pub fn wait_timeout_while_pageserver_connstr_unchanged(&self, duration: Duration) {
    2431            0 :         let state = self.state.lock().unwrap();
    2432            0 :         let old_pageserver_connstr = state
    2433            0 :             .pspec
    2434            0 :             .as_ref()
    2435            0 :             .expect("spec must be set")
    2436            0 :             .pageserver_connstr
    2437            0 :             .clone();
    2438            0 :         let mut unchanged = true;
    2439            0 :         let _ = self
    2440            0 :             .state_changed
    2441            0 :             .wait_timeout_while(state, duration, |s| {
    2442            0 :                 let pageserver_connstr = &s
    2443            0 :                     .pspec
    2444            0 :                     .as_ref()
    2445            0 :                     .expect("spec must be set")
    2446            0 :                     .pageserver_connstr;
    2447            0 :                 unchanged = pageserver_connstr == &old_pageserver_connstr;
    2448            0 :                 unchanged
    2449            0 :             })
    2450            0 :             .unwrap();
    2451            0 :         if !unchanged {
    2452            0 :             info!("Pageserver config changed");
    2453            0 :         }
    2454            0 :     }
    2455              : 
    2456            0 :     pub fn spawn_extension_stats_task(&self) {
    2457            0 :         self.terminate_extension_stats_task();
    2458              : 
    2459            0 :         let conf = self.tokio_conn_conf.clone();
    2460            0 :         let atomic_interval = self.params.installed_extensions_collection_interval.clone();
    2461            0 :         let mut installed_extensions_collection_interval =
    2462            0 :             2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst);
    2463            0 :         info!(
    2464            0 :             "[NEON_EXT_SPAWN] Spawning background installed extensions worker with Timeout: {}",
    2465              :             installed_extensions_collection_interval
    2466              :         );
    2467            0 :         let handle = tokio::spawn(async move {
    2468              :             loop {
    2469            0 :                 info!(
    2470            0 :                     "[NEON_EXT_INT_SLEEP]: Interval: {}",
    2471              :                     installed_extensions_collection_interval
    2472              :                 );
    2473              :                 // Sleep at the start of the loop to ensure that two collections don't happen at the same time.
    2474              :                 // The first collection happens during compute startup.
    2475            0 :                 tokio::time::sleep(tokio::time::Duration::from_secs(
    2476            0 :                     installed_extensions_collection_interval,
    2477            0 :                 ))
    2478            0 :                 .await;
    2479            0 :                 let _ = installed_extensions(conf.clone()).await;
    2480              :                 // Acquire a read lock on the compute spec and then update the interval if necessary
    2481            0 :                 installed_extensions_collection_interval = std::cmp::max(
    2482            0 :                     installed_extensions_collection_interval,
    2483            0 :                     2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst),
    2484            0 :                 );
    2485              :             }
    2486              :         });
    2487              : 
    2488              :         // Store the new task handle
    2489            0 :         *self.extension_stats_task.lock().unwrap() = Some(handle);
    2490            0 :     }
    2491              : 
    2492            0 :     fn terminate_extension_stats_task(&self) {
    2493            0 :         if let Some(h) = self.extension_stats_task.lock().unwrap().take() {
    2494            0 :             h.abort()
    2495            0 :         }
    2496            0 :     }
    2497              : 
    2498            0 :     pub fn spawn_lfc_offload_task(self: &Arc<Self>, interval: Duration) {
    2499            0 :         self.terminate_lfc_offload_task();
    2500            0 :         let secs = interval.as_secs();
    2501            0 :         let this = self.clone();
    2502              : 
    2503            0 :         info!("spawning LFC offload worker with {secs}s interval");
    2504            0 :         let handle = spawn(async move {
    2505            0 :             let mut interval = time::interval(interval);
    2506            0 :             interval.tick().await; // returns immediately
    2507              :             loop {
    2508            0 :                 interval.tick().await;
    2509              : 
    2510            0 :                 let prewarm_state = this.state.lock().unwrap().lfc_prewarm_state.clone();
    2511              :                 // Do not offload LFC state if we are currently prewarming or any issue occurred.
    2512              :                 // If we'd do that, we might override the LFC state in endpoint storage with some
    2513              :                 // incomplete state. Imagine a situation:
    2514              :                 // 1. Endpoint started with `autoprewarm: true`
    2515              :                 // 2. While prewarming is not completed, we upload the new incomplete state
    2516              :                 // 3. Compute gets interrupted and restarts
    2517              :                 // 4. We start again and try to prewarm with the state from 2. instead of the previous complete state
    2518            0 :                 if matches!(
    2519            0 :                     prewarm_state,
    2520              :                     LfcPrewarmState::Completed
    2521              :                         | LfcPrewarmState::NotPrewarmed
    2522              :                         | LfcPrewarmState::Skipped
    2523              :                 ) {
    2524            0 :                     this.offload_lfc_async().await;
    2525            0 :                 }
    2526              :             }
    2527              :         });
    2528            0 :         *self.lfc_offload_task.lock().unwrap() = Some(handle);
    2529            0 :     }
    2530              : 
    2531            0 :     fn terminate_lfc_offload_task(&self) {
    2532            0 :         if let Some(h) = self.lfc_offload_task.lock().unwrap().take() {
    2533            0 :             h.abort()
    2534            0 :         }
    2535            0 :     }
    2536              : 
    2537            0 :     fn update_installed_extensions_collection_interval(&self, spec: &ComputeSpec) {
    2538              :         // Update the interval for collecting installed extensions statistics
    2539              :         // If the value is -1, we never suspend so set the value to default collection.
    2540              :         // If the value is 0, it means default, we will just continue to use the default.
    2541            0 :         if spec.suspend_timeout_seconds == -1 || spec.suspend_timeout_seconds == 0 {
    2542            0 :             self.params.installed_extensions_collection_interval.store(
    2543            0 :                 DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL,
    2544            0 :                 std::sync::atomic::Ordering::SeqCst,
    2545            0 :             );
    2546            0 :         } else {
    2547            0 :             self.params.installed_extensions_collection_interval.store(
    2548            0 :                 spec.suspend_timeout_seconds as u64,
    2549            0 :                 std::sync::atomic::Ordering::SeqCst,
    2550            0 :             );
    2551            0 :         }
    2552            0 :     }
    2553              : }
    2554              : 
    2555            0 : pub async fn installed_extensions(conf: tokio_postgres::Config) -> Result<()> {
    2556            0 :     let res = get_installed_extensions(conf).await;
    2557            0 :     match res {
    2558            0 :         Ok(extensions) => {
    2559            0 :             info!(
    2560            0 :                 "[NEON_EXT_STAT] {}",
    2561            0 :                 serde_json::to_string(&extensions).expect("failed to serialize extensions list")
    2562              :             );
    2563              :         }
    2564            0 :         Err(err) => error!("could not get installed extensions: {err}"),
    2565              :     }
    2566            0 :     Ok(())
    2567            0 : }
    2568              : 
    2569            0 : pub fn forward_termination_signal(dev_mode: bool) {
    2570            0 :     let ss_pid = SYNC_SAFEKEEPERS_PID.load(Ordering::SeqCst);
    2571            0 :     if ss_pid != 0 {
    2572            0 :         let ss_pid = nix::unistd::Pid::from_raw(ss_pid as i32);
    2573            0 :         kill(ss_pid, Signal::SIGTERM).ok();
    2574            0 :     }
    2575              : 
    2576            0 :     if !dev_mode {
    2577              :         //  Terminate pgbouncer with SIGKILL
    2578            0 :         match pid_file::read(PGBOUNCER_PIDFILE.into()) {
    2579            0 :             Ok(pid_file::PidFileRead::LockedByOtherProcess(pid)) => {
    2580            0 :                 info!("sending SIGKILL to pgbouncer process pid: {}", pid);
    2581            0 :                 if let Err(e) = kill(pid, Signal::SIGKILL) {
    2582            0 :                     error!("failed to terminate pgbouncer: {}", e);
    2583            0 :                 }
    2584              :             }
    2585              :             // pgbouncer does not lock the pid file, so we read and kill the process directly
    2586              :             Ok(pid_file::PidFileRead::NotHeldByAnyProcess(_)) => {
    2587            0 :                 if let Ok(pid_str) = std::fs::read_to_string(PGBOUNCER_PIDFILE) {
    2588            0 :                     if let Ok(pid) = pid_str.trim().parse::<i32>() {
    2589            0 :                         info!(
    2590            0 :                             "sending SIGKILL to pgbouncer process pid: {} (from unlocked pid file)",
    2591              :                             pid
    2592              :                         );
    2593            0 :                         if let Err(e) = kill(Pid::from_raw(pid), Signal::SIGKILL) {
    2594            0 :                             error!("failed to terminate pgbouncer: {}", e);
    2595            0 :                         }
    2596            0 :                     }
    2597              :                 } else {
    2598            0 :                     info!("pgbouncer pid file exists but process not running");
    2599              :                 }
    2600              :             }
    2601              :             Ok(pid_file::PidFileRead::NotExist) => {
    2602            0 :                 info!("pgbouncer pid file not found, process may not be running");
    2603              :             }
    2604            0 :             Err(e) => {
    2605            0 :                 error!("error reading pgbouncer pid file: {}", e);
    2606              :             }
    2607              :         }
    2608              : 
    2609              :         // Terminate local_proxy
    2610            0 :         match pid_file::read("/etc/local_proxy/pid".into()) {
    2611            0 :             Ok(pid_file::PidFileRead::LockedByOtherProcess(pid)) => {
    2612            0 :                 info!("sending SIGTERM to local_proxy process pid: {}", pid);
    2613            0 :                 if let Err(e) = kill(pid, Signal::SIGTERM) {
    2614            0 :                     error!("failed to terminate local_proxy: {}", e);
    2615            0 :                 }
    2616              :             }
    2617              :             Ok(pid_file::PidFileRead::NotHeldByAnyProcess(_)) => {
    2618            0 :                 info!("local_proxy PID file exists but process not running");
    2619              :             }
    2620              :             Ok(pid_file::PidFileRead::NotExist) => {
    2621            0 :                 info!("local_proxy PID file not found, process may not be running");
    2622              :             }
    2623            0 :             Err(e) => {
    2624            0 :                 error!("error reading local_proxy PID file: {}", e);
    2625              :             }
    2626              :         }
    2627              :     } else {
    2628            0 :         info!("Skipping pgbouncer and local_proxy termination because in dev mode");
    2629              :     }
    2630              : 
    2631            0 :     let pg_pid = PG_PID.load(Ordering::SeqCst);
    2632            0 :     if pg_pid != 0 {
    2633            0 :         let pg_pid = nix::unistd::Pid::from_raw(pg_pid as i32);
    2634            0 :         // Use 'fast' shutdown (SIGINT) because it also creates a shutdown checkpoint, which is important for
    2635            0 :         // ROs to get a list of running xacts faster instead of going through the CLOG.
    2636            0 :         // See https://www.postgresql.org/docs/current/server-shutdown.html for the list of modes and signals.
    2637            0 :         kill(pg_pid, Signal::SIGINT).ok();
    2638            0 :     }
    2639            0 : }
    2640              : 
    2641              : // helper trait to call JoinSet::spawn_blocking(f), but propagates the current
    2642              : // tracing span to the thread.
    2643              : trait JoinSetExt<T> {
    2644              :     fn spawn_blocking_child<F>(&mut self, f: F) -> tokio::task::AbortHandle
    2645              :     where
    2646              :         F: FnOnce() -> T + Send + 'static,
    2647              :         T: Send;
    2648              : }
    2649              : 
    2650              : impl<T: 'static> JoinSetExt<T> for tokio::task::JoinSet<T> {
    2651            0 :     fn spawn_blocking_child<F>(&mut self, f: F) -> tokio::task::AbortHandle
    2652            0 :     where
    2653            0 :         F: FnOnce() -> T + Send + 'static,
    2654            0 :         T: Send,
    2655              :     {
    2656            0 :         let sp = tracing::Span::current();
    2657            0 :         self.spawn_blocking(move || {
    2658            0 :             let _e = sp.enter();
    2659            0 :             f()
    2660            0 :         })
    2661            0 :     }
    2662              : }
    2663              : 
    2664              : #[cfg(test)]
    2665              : mod tests {
    2666              :     use std::fs::File;
    2667              : 
    2668              :     use super::*;
    2669              : 
    2670              :     #[test]
    2671            1 :     fn duplicate_safekeeper_connstring() {
    2672            1 :         let file = File::open("tests/cluster_spec.json").unwrap();
    2673            1 :         let spec: ComputeSpec = serde_json::from_reader(file).unwrap();
    2674              : 
    2675            1 :         match ParsedSpec::try_from(spec.clone()) {
    2676            0 :             Ok(_p) => panic!("Failed to detect duplicate entry"),
    2677            1 :             Err(e) => assert!(e.starts_with("duplicate entry in safekeeper_connstrings:")),
    2678              :         };
    2679            1 :     }
    2680              : }
        

Generated by: LCOV version 2.1-beta