LCOV - code coverage report
Current view: top level - compute_tools/src - compute.rs (source / functions) Coverage Total Hit
Test: fc67f8dc6087a0b4f4f0bcd74f6e1dc25fab8cf3.info Lines: 0.0 % 535 0
Test Date: 2024-09-24 13:57:57 Functions: 0.0 % 57 0

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::env;
       3              : use std::fs;
       4              : use std::io::BufRead;
       5              : use std::os::unix::fs::{symlink, PermissionsExt};
       6              : use std::path::Path;
       7              : use std::process::{Command, Stdio};
       8              : use std::str::FromStr;
       9              : use std::sync::atomic::AtomicU32;
      10              : use std::sync::atomic::Ordering;
      11              : use std::sync::{Condvar, Mutex, RwLock};
      12              : use std::thread;
      13              : use std::time::Instant;
      14              : 
      15              : use anyhow::{Context, Result};
      16              : use chrono::{DateTime, Utc};
      17              : use futures::future::join_all;
      18              : use futures::stream::FuturesUnordered;
      19              : use futures::StreamExt;
      20              : use nix::unistd::Pid;
      21              : use postgres::error::SqlState;
      22              : use postgres::{Client, NoTls};
      23              : use tracing::{debug, error, info, instrument, warn};
      24              : use utils::id::{TenantId, TimelineId};
      25              : use utils::lsn::Lsn;
      26              : 
      27              : use compute_api::responses::{ComputeMetrics, ComputeStatus};
      28              : use compute_api::spec::{ComputeFeature, ComputeMode, ComputeSpec};
      29              : use utils::measured_stream::MeasuredReader;
      30              : 
      31              : use nix::sys::signal::{kill, Signal};
      32              : 
      33              : use remote_storage::{DownloadError, RemotePath};
      34              : 
      35              : use crate::checker::create_availability_check_data;
      36              : use crate::logger::inlinify;
      37              : use crate::pg_helpers::*;
      38              : use crate::spec::*;
      39              : use crate::sync_sk::{check_if_synced, ping_safekeeper};
      40              : use crate::{config, extension_server};
      41              : 
      42              : pub static SYNC_SAFEKEEPERS_PID: AtomicU32 = AtomicU32::new(0);
      43              : pub static PG_PID: AtomicU32 = AtomicU32::new(0);
      44              : 
      45              : /// Compute node info shared across several `compute_ctl` threads.
      46              : pub struct ComputeNode {
      47              :     // Url type maintains proper escaping
      48              :     pub connstr: url::Url,
      49              :     pub pgdata: String,
      50              :     pub pgbin: String,
      51              :     pub pgversion: String,
      52              :     /// We should only allow live re- / configuration of the compute node if
      53              :     /// it uses 'pull model', i.e. it can go to control-plane and fetch
      54              :     /// the latest configuration. Otherwise, there could be a case:
      55              :     /// - we start compute with some spec provided as argument
      56              :     /// - we push new spec and it does reconfiguration
      57              :     /// - but then something happens and compute pod / VM is destroyed,
      58              :     ///   so k8s controller starts it again with the **old** spec
      59              :     ///
      60              :     /// and the same for empty computes:
      61              :     /// - we started compute without any spec
      62              :     /// - we push spec and it does configuration
      63              :     /// - but then it is restarted without any spec again
      64              :     pub live_config_allowed: bool,
      65              :     /// Volatile part of the `ComputeNode`, which should be used under `Mutex`.
      66              :     /// To allow HTTP API server to serving status requests, while configuration
      67              :     /// is in progress, lock should be held only for short periods of time to do
      68              :     /// read/write, not the whole configuration process.
      69              :     pub state: Mutex<ComputeState>,
      70              :     /// `Condvar` to allow notifying waiters about state changes.
      71              :     pub state_changed: Condvar,
      72              :     /// the address of extension storage proxy gateway
      73              :     pub ext_remote_storage: Option<String>,
      74              :     // key: ext_archive_name, value: started download time, download_completed?
      75              :     pub ext_download_progress: RwLock<HashMap<String, (DateTime<Utc>, bool)>>,
      76              :     pub build_tag: String,
      77              : }
      78              : 
      79              : // store some metrics about download size that might impact startup time
      80              : #[derive(Clone, Debug)]
      81              : pub struct RemoteExtensionMetrics {
      82              :     num_ext_downloaded: u64,
      83              :     largest_ext_size: u64,
      84              :     total_ext_download_size: u64,
      85              : }
      86              : 
      87              : #[derive(Clone, Debug)]
      88              : pub struct ComputeState {
      89              :     pub start_time: DateTime<Utc>,
      90              :     pub status: ComputeStatus,
      91              :     /// Timestamp of the last Postgres activity. It could be `None` if
      92              :     /// compute wasn't used since start.
      93              :     pub last_active: Option<DateTime<Utc>>,
      94              :     pub error: Option<String>,
      95              :     pub pspec: Option<ParsedSpec>,
      96              :     pub metrics: ComputeMetrics,
      97              : }
      98              : 
      99              : impl ComputeState {
     100            0 :     pub fn new() -> Self {
     101            0 :         Self {
     102            0 :             start_time: Utc::now(),
     103            0 :             status: ComputeStatus::Empty,
     104            0 :             last_active: None,
     105            0 :             error: None,
     106            0 :             pspec: None,
     107            0 :             metrics: ComputeMetrics::default(),
     108            0 :         }
     109            0 :     }
     110              : }
     111              : 
     112              : impl Default for ComputeState {
     113            0 :     fn default() -> Self {
     114            0 :         Self::new()
     115            0 :     }
     116              : }
     117              : 
     118              : #[derive(Clone, Debug)]
     119              : pub struct ParsedSpec {
     120              :     pub spec: ComputeSpec,
     121              :     pub tenant_id: TenantId,
     122              :     pub timeline_id: TimelineId,
     123              :     pub pageserver_connstr: String,
     124              :     pub safekeeper_connstrings: Vec<String>,
     125              :     pub storage_auth_token: Option<String>,
     126              : }
     127              : 
     128              : impl TryFrom<ComputeSpec> for ParsedSpec {
     129              :     type Error = String;
     130            0 :     fn try_from(spec: ComputeSpec) -> Result<Self, String> {
     131              :         // Extract the options from the spec file that are needed to connect to
     132              :         // the storage system.
     133              :         //
     134              :         // For backwards-compatibility, the top-level fields in the spec file
     135              :         // may be empty. In that case, we need to dig them from the GUCs in the
     136              :         // cluster.settings field.
     137            0 :         let pageserver_connstr = spec
     138            0 :             .pageserver_connstring
     139            0 :             .clone()
     140            0 :             .or_else(|| spec.cluster.settings.find("neon.pageserver_connstring"))
     141            0 :             .ok_or("pageserver connstr should be provided")?;
     142            0 :         let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
     143            0 :             if matches!(spec.mode, ComputeMode::Primary) {
     144            0 :                 spec.cluster
     145            0 :                     .settings
     146            0 :                     .find("neon.safekeepers")
     147            0 :                     .ok_or("safekeeper connstrings should be provided")?
     148            0 :                     .split(',')
     149            0 :                     .map(|str| str.to_string())
     150            0 :                     .collect()
     151              :             } else {
     152            0 :                 vec![]
     153              :             }
     154              :         } else {
     155            0 :             spec.safekeeper_connstrings.clone()
     156              :         };
     157            0 :         let storage_auth_token = spec.storage_auth_token.clone();
     158            0 :         let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
     159            0 :             tenant_id
     160              :         } else {
     161            0 :             spec.cluster
     162            0 :                 .settings
     163            0 :                 .find("neon.tenant_id")
     164            0 :                 .ok_or("tenant id should be provided")
     165            0 :                 .map(|s| TenantId::from_str(&s))?
     166            0 :                 .or(Err("invalid tenant id"))?
     167              :         };
     168            0 :         let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
     169            0 :             timeline_id
     170              :         } else {
     171            0 :             spec.cluster
     172            0 :                 .settings
     173            0 :                 .find("neon.timeline_id")
     174            0 :                 .ok_or("timeline id should be provided")
     175            0 :                 .map(|s| TimelineId::from_str(&s))?
     176            0 :                 .or(Err("invalid timeline id"))?
     177              :         };
     178              : 
     179            0 :         Ok(ParsedSpec {
     180            0 :             spec,
     181            0 :             pageserver_connstr,
     182            0 :             safekeeper_connstrings,
     183            0 :             storage_auth_token,
     184            0 :             tenant_id,
     185            0 :             timeline_id,
     186            0 :         })
     187            0 :     }
     188              : }
     189              : 
     190              : /// If we are a VM, returns a [`Command`] that will run in the `neon-postgres`
     191              : /// cgroup. Otherwise returns the default `Command::new(cmd)`
     192              : ///
     193              : /// This function should be used to start postgres, as it will start it in the
     194              : /// neon-postgres cgroup if we are a VM. This allows autoscaling to control
     195              : /// postgres' resource usage. The cgroup will exist in VMs because vm-builder
     196              : /// creates it during the sysinit phase of its inittab.
     197            0 : fn maybe_cgexec(cmd: &str) -> Command {
     198            0 :     // The cplane sets this env var for autoscaling computes.
     199            0 :     // use `var_os` so we don't have to worry about the variable being valid
     200            0 :     // unicode. Should never be an concern . . . but just in case
     201            0 :     if env::var_os("AUTOSCALING").is_some() {
     202            0 :         let mut command = Command::new("cgexec");
     203            0 :         command.args(["-g", "memory:neon-postgres"]);
     204            0 :         command.arg(cmd);
     205            0 :         command
     206              :     } else {
     207            0 :         Command::new(cmd)
     208              :     }
     209            0 : }
     210              : 
     211              : /// Create special neon_superuser role, that's a slightly nerfed version of a real superuser
     212              : /// that we give to customers
     213            0 : #[instrument(skip_all)]
     214              : fn create_neon_superuser(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     215              :     let roles = spec
     216              :         .cluster
     217              :         .roles
     218              :         .iter()
     219            0 :         .map(|r| escape_literal(&r.name))
     220              :         .collect::<Vec<_>>();
     221              : 
     222              :     let dbs = spec
     223              :         .cluster
     224              :         .databases
     225              :         .iter()
     226            0 :         .map(|db| escape_literal(&db.name))
     227              :         .collect::<Vec<_>>();
     228              : 
     229              :     let roles_decl = if roles.is_empty() {
     230              :         String::from("roles text[] := NULL;")
     231              :     } else {
     232              :         format!(
     233              :             r#"
     234              :                roles text[] := ARRAY(SELECT rolname
     235              :                                      FROM pg_catalog.pg_roles
     236              :                                      WHERE rolname IN ({}));"#,
     237              :             roles.join(", ")
     238              :         )
     239              :     };
     240              : 
     241              :     let database_decl = if dbs.is_empty() {
     242              :         String::from("dbs text[] := NULL;")
     243              :     } else {
     244              :         format!(
     245              :             r#"
     246              :                dbs text[] := ARRAY(SELECT datname
     247              :                                    FROM pg_catalog.pg_database
     248              :                                    WHERE datname IN ({}));"#,
     249              :             dbs.join(", ")
     250              :         )
     251              :     };
     252              : 
     253              :     // ALL PRIVILEGES grants CREATE, CONNECT, and TEMPORARY on all databases
     254              :     // (see https://www.postgresql.org/docs/current/ddl-priv.html)
     255              :     let query = format!(
     256              :         r#"
     257              :             DO $$
     258              :                 DECLARE
     259              :                     r text;
     260              :                     {}
     261              :                     {}
     262              :                 BEGIN
     263              :                     IF NOT EXISTS (
     264              :                         SELECT FROM pg_catalog.pg_roles WHERE rolname = 'neon_superuser')
     265              :                     THEN
     266              :                         CREATE ROLE neon_superuser CREATEDB CREATEROLE NOLOGIN REPLICATION BYPASSRLS IN ROLE pg_read_all_data, pg_write_all_data;
     267              :                         IF array_length(roles, 1) IS NOT NULL THEN
     268              :                             EXECUTE format('GRANT neon_superuser TO %s',
     269              :                                            array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(roles) as x), ', '));
     270              :                             FOREACH r IN ARRAY roles LOOP
     271              :                                 EXECUTE format('ALTER ROLE %s CREATEROLE CREATEDB', quote_ident(r));
     272              :                             END LOOP;
     273              :                         END IF;
     274              :                         IF array_length(dbs, 1) IS NOT NULL THEN
     275              :                             EXECUTE format('GRANT ALL PRIVILEGES ON DATABASE %s TO neon_superuser',
     276              :                                            array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(dbs) as x), ', '));
     277              :                         END IF;
     278              :                     END IF;
     279              :                 END
     280              :             $$;"#,
     281              :         roles_decl, database_decl,
     282              :     );
     283              :     info!("Neon superuser created: {}", inlinify(&query));
     284              :     client
     285              :         .simple_query(&query)
     286            0 :         .map_err(|e| anyhow::anyhow!(e).context(query))?;
     287              :     Ok(())
     288              : }
     289              : 
     290              : impl ComputeNode {
     291              :     /// Check that compute node has corresponding feature enabled.
     292            0 :     pub fn has_feature(&self, feature: ComputeFeature) -> bool {
     293            0 :         let state = self.state.lock().unwrap();
     294              : 
     295            0 :         if let Some(s) = state.pspec.as_ref() {
     296            0 :             s.spec.features.contains(&feature)
     297              :         } else {
     298            0 :             false
     299              :         }
     300            0 :     }
     301              : 
     302            0 :     pub fn set_status(&self, status: ComputeStatus) {
     303            0 :         let mut state = self.state.lock().unwrap();
     304            0 :         state.status = status;
     305            0 :         self.state_changed.notify_all();
     306            0 :     }
     307              : 
     308            0 :     pub fn get_status(&self) -> ComputeStatus {
     309            0 :         self.state.lock().unwrap().status
     310            0 :     }
     311              : 
     312              :     // Remove `pgdata` directory and create it again with right permissions.
     313            0 :     fn create_pgdata(&self) -> Result<()> {
     314            0 :         // Ignore removal error, likely it is a 'No such file or directory (os error 2)'.
     315            0 :         // If it is something different then create_dir() will error out anyway.
     316            0 :         let _ok = fs::remove_dir_all(&self.pgdata);
     317            0 :         fs::create_dir(&self.pgdata)?;
     318            0 :         fs::set_permissions(&self.pgdata, fs::Permissions::from_mode(0o700))?;
     319              : 
     320            0 :         Ok(())
     321            0 :     }
     322              : 
     323              :     // Get basebackup from the libpq connection to pageserver using `connstr` and
     324              :     // unarchive it to `pgdata` directory overriding all its previous content.
     325            0 :     #[instrument(skip_all, fields(%lsn))]
     326              :     fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
     327              :         let spec = compute_state.pspec.as_ref().expect("spec must be set");
     328              :         let start_time = Instant::now();
     329              : 
     330              :         let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
     331              :         let mut config = postgres::Config::from_str(shard0_connstr)?;
     332              : 
     333              :         // Use the storage auth token from the config file, if given.
     334              :         // Note: this overrides any password set in the connection string.
     335              :         if let Some(storage_auth_token) = &spec.storage_auth_token {
     336              :             info!("Got storage auth token from spec file");
     337              :             config.password(storage_auth_token);
     338              :         } else {
     339              :             info!("Storage auth token not set");
     340              :         }
     341              : 
     342              :         // Connect to pageserver
     343              :         let mut client = config.connect(NoTls)?;
     344              :         let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;
     345              : 
     346              :         let basebackup_cmd = match lsn {
     347              :             // HACK We don't use compression on first start (Lsn(0)) because there's no API for it
     348              :             Lsn(0) => format!("basebackup {} {}", spec.tenant_id, spec.timeline_id),
     349              :             _ => format!(
     350              :                 "basebackup {} {} {} --gzip",
     351              :                 spec.tenant_id, spec.timeline_id, lsn
     352              :             ),
     353              :         };
     354              : 
     355              :         let copyreader = client.copy_out(basebackup_cmd.as_str())?;
     356              :         let mut measured_reader = MeasuredReader::new(copyreader);
     357              : 
     358              :         // Check the magic number to see if it's a gzip or not. Even though
     359              :         // we might explicitly ask for gzip, an old pageserver with no implementation
     360              :         // of gzip compression might send us uncompressed data. After some time
     361              :         // passes we can assume all pageservers know how to compress and we can
     362              :         // delete this check.
     363              :         //
     364              :         // If the data is not gzip, it will be tar. It will not be mistakenly
     365              :         // recognized as gzip because tar starts with an ascii encoding of a filename,
     366              :         // and 0x1f and 0x8b are unlikely first characters for any filename. Moreover,
     367              :         // we send the "global" directory first from the pageserver, so it definitely
     368              :         // won't be recognized as gzip.
     369              :         let mut bufreader = std::io::BufReader::new(&mut measured_reader);
     370              :         let gzip = {
     371              :             let peek = bufreader.fill_buf().unwrap();
     372              :             peek[0] == 0x1f && peek[1] == 0x8b
     373              :         };
     374              : 
     375              :         // Read the archive directly from the `CopyOutReader`
     376              :         //
     377              :         // Set `ignore_zeros` so that unpack() reads all the Copy data and
     378              :         // doesn't stop at the end-of-archive marker. Otherwise, if the server
     379              :         // sends an Error after finishing the tarball, we will not notice it.
     380              :         if gzip {
     381              :             let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
     382              :             ar.set_ignore_zeros(true);
     383              :             ar.unpack(&self.pgdata)?;
     384              :         } else {
     385              :             let mut ar = tar::Archive::new(&mut bufreader);
     386              :             ar.set_ignore_zeros(true);
     387              :             ar.unpack(&self.pgdata)?;
     388              :         };
     389              : 
     390              :         // Report metrics
     391              :         let mut state = self.state.lock().unwrap();
     392              :         state.metrics.pageserver_connect_micros = pageserver_connect_micros;
     393              :         state.metrics.basebackup_bytes = measured_reader.get_byte_count() as u64;
     394              :         state.metrics.basebackup_ms = start_time.elapsed().as_millis() as u64;
     395              :         Ok(())
     396              :     }
     397              : 
     398              :     // Gets the basebackup in a retry loop
     399            0 :     #[instrument(skip_all, fields(%lsn))]
     400              :     pub fn get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
     401              :         let mut retry_period_ms = 500.0;
     402              :         let mut attempts = 0;
     403              :         const DEFAULT_ATTEMPTS: u16 = 10;
     404              :         #[cfg(feature = "testing")]
     405              :         let max_attempts = if let Ok(v) = env::var("NEON_COMPUTE_TESTING_BASEBACKUP_RETRIES") {
     406              :             u16::from_str(&v).unwrap()
     407              :         } else {
     408              :             DEFAULT_ATTEMPTS
     409              :         };
     410              :         #[cfg(not(feature = "testing"))]
     411              :         let max_attempts = DEFAULT_ATTEMPTS;
     412              :         loop {
     413              :             let result = self.try_get_basebackup(compute_state, lsn);
     414              :             match result {
     415              :                 Ok(_) => {
     416              :                     return result;
     417              :                 }
     418              :                 Err(ref e) if attempts < max_attempts => {
     419              :                     warn!(
     420              :                         "Failed to get basebackup: {} (attempt {}/{})",
     421              :                         e, attempts, max_attempts
     422              :                     );
     423              :                     std::thread::sleep(std::time::Duration::from_millis(retry_period_ms as u64));
     424              :                     retry_period_ms *= 1.5;
     425              :                 }
     426              :                 Err(_) => {
     427              :                     return result;
     428              :                 }
     429              :             }
     430              :             attempts += 1;
     431              :         }
     432              :     }
     433              : 
     434            0 :     pub async fn check_safekeepers_synced_async(
     435            0 :         &self,
     436            0 :         compute_state: &ComputeState,
     437            0 :     ) -> Result<Option<Lsn>> {
     438            0 :         // Construct a connection config for each safekeeper
     439            0 :         let pspec: ParsedSpec = compute_state
     440            0 :             .pspec
     441            0 :             .as_ref()
     442            0 :             .expect("spec must be set")
     443            0 :             .clone();
     444            0 :         let sk_connstrs: Vec<String> = pspec.safekeeper_connstrings.clone();
     445            0 :         let sk_configs = sk_connstrs.into_iter().map(|connstr| {
     446            0 :             // Format connstr
     447            0 :             let id = connstr.clone();
     448            0 :             let connstr = format!("postgresql://no_user@{}", connstr);
     449            0 :             let options = format!(
     450            0 :                 "-c timeline_id={} tenant_id={}",
     451            0 :                 pspec.timeline_id, pspec.tenant_id
     452            0 :             );
     453            0 : 
     454            0 :             // Construct client
     455            0 :             let mut config = tokio_postgres::Config::from_str(&connstr).unwrap();
     456            0 :             config.options(&options);
     457            0 :             if let Some(storage_auth_token) = pspec.storage_auth_token.clone() {
     458            0 :                 config.password(storage_auth_token);
     459            0 :             }
     460              : 
     461            0 :             (id, config)
     462            0 :         });
     463            0 : 
     464            0 :         // Create task set to query all safekeepers
     465            0 :         let mut tasks = FuturesUnordered::new();
     466            0 :         let quorum = sk_configs.len() / 2 + 1;
     467            0 :         for (id, config) in sk_configs {
     468            0 :             let timeout = tokio::time::Duration::from_millis(100);
     469            0 :             let task = tokio::time::timeout(timeout, ping_safekeeper(id, config));
     470            0 :             tasks.push(tokio::spawn(task));
     471            0 :         }
     472              : 
     473              :         // Get a quorum of responses or errors
     474            0 :         let mut responses = Vec::new();
     475            0 :         let mut join_errors = Vec::new();
     476            0 :         let mut task_errors = Vec::new();
     477            0 :         let mut timeout_errors = Vec::new();
     478            0 :         while let Some(response) = tasks.next().await {
     479            0 :             match response {
     480            0 :                 Ok(Ok(Ok(r))) => responses.push(r),
     481            0 :                 Ok(Ok(Err(e))) => task_errors.push(e),
     482            0 :                 Ok(Err(e)) => timeout_errors.push(e),
     483            0 :                 Err(e) => join_errors.push(e),
     484              :             };
     485            0 :             if responses.len() >= quorum {
     486            0 :                 break;
     487            0 :             }
     488            0 :             if join_errors.len() + task_errors.len() + timeout_errors.len() >= quorum {
     489            0 :                 break;
     490            0 :             }
     491              :         }
     492              : 
     493              :         // In case of error, log and fail the check, but don't crash.
     494              :         // We're playing it safe because these errors could be transient
     495              :         // and we don't yet retry. Also being careful here allows us to
     496              :         // be backwards compatible with safekeepers that don't have the
     497              :         // TIMELINE_STATUS API yet.
     498            0 :         if responses.len() < quorum {
     499            0 :             error!(
     500            0 :                 "failed sync safekeepers check {:?} {:?} {:?}",
     501              :                 join_errors, task_errors, timeout_errors
     502              :             );
     503            0 :             return Ok(None);
     504            0 :         }
     505            0 : 
     506            0 :         Ok(check_if_synced(responses))
     507            0 :     }
     508              : 
     509              :     // Fast path for sync_safekeepers. If they're already synced we get the lsn
     510              :     // in one roundtrip. If not, we should do a full sync_safekeepers.
     511            0 :     pub fn check_safekeepers_synced(&self, compute_state: &ComputeState) -> Result<Option<Lsn>> {
     512            0 :         let start_time = Utc::now();
     513            0 : 
     514            0 :         // Run actual work with new tokio runtime
     515            0 :         let rt = tokio::runtime::Builder::new_current_thread()
     516            0 :             .enable_all()
     517            0 :             .build()
     518            0 :             .expect("failed to create rt");
     519            0 :         let result = rt.block_on(self.check_safekeepers_synced_async(compute_state));
     520            0 : 
     521            0 :         // Record runtime
     522            0 :         self.state.lock().unwrap().metrics.sync_sk_check_ms = Utc::now()
     523            0 :             .signed_duration_since(start_time)
     524            0 :             .to_std()
     525            0 :             .unwrap()
     526            0 :             .as_millis() as u64;
     527            0 :         result
     528            0 :     }
     529              : 
     530              :     // Run `postgres` in a special mode with `--sync-safekeepers` argument
     531              :     // and return the reported LSN back to the caller.
     532            0 :     #[instrument(skip_all)]
     533              :     pub fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
     534              :         let start_time = Utc::now();
     535              : 
     536              :         let mut sync_handle = maybe_cgexec(&self.pgbin)
     537              :             .args(["--sync-safekeepers"])
     538              :             .env("PGDATA", &self.pgdata) // we cannot use -D in this mode
     539              :             .envs(if let Some(storage_auth_token) = &storage_auth_token {
     540              :                 vec![("NEON_AUTH_TOKEN", storage_auth_token)]
     541              :             } else {
     542              :                 vec![]
     543              :             })
     544              :             .stdout(Stdio::piped())
     545              :             .stderr(Stdio::piped())
     546              :             .spawn()
     547              :             .expect("postgres --sync-safekeepers failed to start");
     548              :         SYNC_SAFEKEEPERS_PID.store(sync_handle.id(), Ordering::SeqCst);
     549              : 
     550              :         // `postgres --sync-safekeepers` will print all log output to stderr and
     551              :         // final LSN to stdout. So we leave stdout to collect LSN, while stderr logs
     552              :         // will be collected in a child thread.
     553              :         let stderr = sync_handle
     554              :             .stderr
     555              :             .take()
     556              :             .expect("stderr should be captured");
     557              :         let logs_handle = handle_postgres_logs(stderr);
     558              : 
     559              :         let sync_output = sync_handle
     560              :             .wait_with_output()
     561              :             .expect("postgres --sync-safekeepers failed");
     562              :         SYNC_SAFEKEEPERS_PID.store(0, Ordering::SeqCst);
     563              : 
     564              :         // Process has exited, so we can join the logs thread.
     565              :         let _ = logs_handle
     566              :             .join()
     567            0 :             .map_err(|e| tracing::error!("log thread panicked: {:?}", e));
     568              : 
     569              :         if !sync_output.status.success() {
     570              :             anyhow::bail!(
     571              :                 "postgres --sync-safekeepers exited with non-zero status: {}. stdout: {}",
     572              :                 sync_output.status,
     573              :                 String::from_utf8(sync_output.stdout)
     574              :                     .expect("postgres --sync-safekeepers exited, and stdout is not utf-8"),
     575              :             );
     576              :         }
     577              : 
     578              :         self.state.lock().unwrap().metrics.sync_safekeepers_ms = Utc::now()
     579              :             .signed_duration_since(start_time)
     580              :             .to_std()
     581              :             .unwrap()
     582              :             .as_millis() as u64;
     583              : 
     584              :         let lsn = Lsn::from_str(String::from_utf8(sync_output.stdout)?.trim())?;
     585              : 
     586              :         Ok(lsn)
     587              :     }
     588              : 
     589              :     /// Do all the preparations like PGDATA directory creation, configuration,
     590              :     /// safekeepers sync, basebackup, etc.
     591            0 :     #[instrument(skip_all)]
     592              :     pub fn prepare_pgdata(
     593              :         &self,
     594              :         compute_state: &ComputeState,
     595              :         extension_server_port: u16,
     596              :     ) -> Result<()> {
     597              :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
     598              :         let spec = &pspec.spec;
     599              :         let pgdata_path = Path::new(&self.pgdata);
     600              : 
     601              :         // Remove/create an empty pgdata directory and put configuration there.
     602              :         self.create_pgdata()?;
     603              :         config::write_postgres_conf(
     604              :             &pgdata_path.join("postgresql.conf"),
     605              :             &pspec.spec,
     606              :             Some(extension_server_port),
     607              :         )?;
     608              : 
     609              :         // Syncing safekeepers is only safe with primary nodes: if a primary
     610              :         // is already connected it will be kicked out, so a secondary (standby)
     611              :         // cannot sync safekeepers.
     612              :         let lsn = match spec.mode {
     613              :             ComputeMode::Primary => {
     614              :                 info!("checking if safekeepers are synced");
     615              :                 let lsn = if let Ok(Some(lsn)) = self.check_safekeepers_synced(compute_state) {
     616              :                     lsn
     617              :                 } else {
     618              :                     info!("starting safekeepers syncing");
     619              :                     self.sync_safekeepers(pspec.storage_auth_token.clone())
     620            0 :                         .with_context(|| "failed to sync safekeepers")?
     621              :                 };
     622              :                 info!("safekeepers synced at LSN {}", lsn);
     623              :                 lsn
     624              :             }
     625              :             ComputeMode::Static(lsn) => {
     626              :                 info!("Starting read-only node at static LSN {}", lsn);
     627              :                 lsn
     628              :             }
     629              :             ComputeMode::Replica => {
     630              :                 info!("Initializing standby from latest Pageserver LSN");
     631              :                 Lsn(0)
     632              :             }
     633              :         };
     634              : 
     635              :         info!(
     636              :             "getting basebackup@{} from pageserver {}",
     637              :             lsn, &pspec.pageserver_connstr
     638              :         );
     639            0 :         self.get_basebackup(compute_state, lsn).with_context(|| {
     640            0 :             format!(
     641            0 :                 "failed to get basebackup@{} from pageserver {}",
     642            0 :                 lsn, &pspec.pageserver_connstr
     643            0 :             )
     644            0 :         })?;
     645              : 
     646              :         // Update pg_hba.conf received with basebackup.
     647              :         update_pg_hba(pgdata_path)?;
     648              : 
     649              :         // Place pg_dynshmem under /dev/shm. This allows us to use
     650              :         // 'dynamic_shared_memory_type = mmap' so that the files are placed in
     651              :         // /dev/shm, similar to how 'dynamic_shared_memory_type = posix' works.
     652              :         //
     653              :         // Why on earth don't we just stick to the 'posix' default, you might
     654              :         // ask.  It turns out that making large allocations with 'posix' doesn't
     655              :         // work very well with autoscaling. The behavior we want is that:
     656              :         //
     657              :         // 1. You can make large DSM allocations, larger than the current RAM
     658              :         //    size of the VM, without errors
     659              :         //
     660              :         // 2. If the allocated memory is really used, the VM is scaled up
     661              :         //    automatically to accommodate that
     662              :         //
     663              :         // We try to make that possible by having swap in the VM. But with the
     664              :         // default 'posix' DSM implementation, we fail step 1, even when there's
     665              :         // plenty of swap available. PostgreSQL uses posix_fallocate() to create
     666              :         // the shmem segment, which is really just a file in /dev/shm in Linux,
     667              :         // but posix_fallocate() on tmpfs returns ENOMEM if the size is larger
     668              :         // than available RAM.
     669              :         //
     670              :         // Using 'dynamic_shared_memory_type = mmap' works around that, because
     671              :         // the Postgres 'mmap' DSM implementation doesn't use
     672              :         // posix_fallocate(). Instead, it uses repeated calls to write(2) to
     673              :         // fill the file with zeros. It's weird that that differs between
     674              :         // 'posix' and 'mmap', but we take advantage of it. When the file is
     675              :         // filled slowly with write(2), the kernel allows it to grow larger, as
     676              :         // long as there's swap available.
     677              :         //
     678              :         // In short, using 'dynamic_shared_memory_type = mmap' allows us one DSM
     679              :         // segment to be larger than currently available RAM. But because we
     680              :         // don't want to store it on a real file, which the kernel would try to
     681              :         // flush to disk, so symlink pg_dynshm to /dev/shm.
     682              :         //
     683              :         // We don't set 'dynamic_shared_memory_type = mmap' here, we let the
     684              :         // control plane control that option. If 'mmap' is not used, this
     685              :         // symlink doesn't affect anything.
     686              :         //
     687              :         // See https://github.com/neondatabase/autoscaling/issues/800
     688              :         std::fs::remove_dir(pgdata_path.join("pg_dynshmem"))?;
     689              :         symlink("/dev/shm/", pgdata_path.join("pg_dynshmem"))?;
     690              : 
     691              :         match spec.mode {
     692              :             ComputeMode::Primary => {}
     693              :             ComputeMode::Replica | ComputeMode::Static(..) => {
     694              :                 add_standby_signal(pgdata_path)?;
     695              :             }
     696              :         }
     697              : 
     698              :         Ok(())
     699              :     }
     700              : 
     701              :     /// Start and stop a postgres process to warm up the VM for startup.
     702            0 :     pub fn prewarm_postgres(&self) -> Result<()> {
     703            0 :         info!("prewarming");
     704              : 
     705              :         // Create pgdata
     706            0 :         let pgdata = &format!("{}.warmup", self.pgdata);
     707            0 :         create_pgdata(pgdata)?;
     708              : 
     709              :         // Run initdb to completion
     710            0 :         info!("running initdb");
     711            0 :         let initdb_bin = Path::new(&self.pgbin).parent().unwrap().join("initdb");
     712            0 :         Command::new(initdb_bin)
     713            0 :             .args(["-D", pgdata])
     714            0 :             .output()
     715            0 :             .expect("cannot start initdb process");
     716              : 
     717              :         // Write conf
     718              :         use std::io::Write;
     719            0 :         let conf_path = Path::new(pgdata).join("postgresql.conf");
     720            0 :         let mut file = std::fs::File::create(conf_path)?;
     721            0 :         writeln!(file, "shared_buffers=65536")?;
     722            0 :         writeln!(file, "port=51055")?; // Nobody should be connecting
     723            0 :         writeln!(file, "shared_preload_libraries = 'neon'")?;
     724              : 
     725              :         // Start postgres
     726            0 :         info!("starting postgres");
     727            0 :         let mut pg = maybe_cgexec(&self.pgbin)
     728            0 :             .args(["-D", pgdata])
     729            0 :             .spawn()
     730            0 :             .expect("cannot start postgres process");
     731            0 : 
     732            0 :         // Stop it when it's ready
     733            0 :         info!("waiting for postgres");
     734            0 :         wait_for_postgres(&mut pg, Path::new(pgdata))?;
     735              :         // SIGQUIT orders postgres to exit immediately. We don't want to SIGKILL
     736              :         // it to avoid orphaned processes prowling around while datadir is
     737              :         // wiped.
     738            0 :         let pm_pid = Pid::from_raw(pg.id() as i32);
     739            0 :         kill(pm_pid, Signal::SIGQUIT)?;
     740            0 :         info!("sent SIGQUIT signal");
     741            0 :         pg.wait()?;
     742            0 :         info!("done prewarming");
     743              : 
     744              :         // clean up
     745            0 :         let _ok = fs::remove_dir_all(pgdata);
     746            0 :         Ok(())
     747            0 :     }
     748              : 
     749              :     /// Start Postgres as a child process and manage DBs/roles.
     750              :     /// After that this will hang waiting on the postmaster process to exit.
     751              :     /// Returns a handle to the child process and a handle to the logs thread.
     752            0 :     #[instrument(skip_all)]
     753              :     pub fn start_postgres(
     754              :         &self,
     755              :         storage_auth_token: Option<String>,
     756              :     ) -> Result<(std::process::Child, std::thread::JoinHandle<()>)> {
     757              :         let pgdata_path = Path::new(&self.pgdata);
     758              : 
     759              :         // Run postgres as a child process.
     760              :         let mut pg = maybe_cgexec(&self.pgbin)
     761              :             .args(["-D", &self.pgdata])
     762              :             .envs(if let Some(storage_auth_token) = &storage_auth_token {
     763              :                 vec![("NEON_AUTH_TOKEN", storage_auth_token)]
     764              :             } else {
     765              :                 vec![]
     766              :             })
     767              :             .stderr(Stdio::piped())
     768              :             .spawn()
     769              :             .expect("cannot start postgres process");
     770              :         PG_PID.store(pg.id(), Ordering::SeqCst);
     771              : 
     772              :         // Start a thread to collect logs from stderr.
     773              :         let stderr = pg.stderr.take().expect("stderr should be captured");
     774              :         let logs_handle = handle_postgres_logs(stderr);
     775              : 
     776              :         wait_for_postgres(&mut pg, pgdata_path)?;
     777              : 
     778              :         Ok((pg, logs_handle))
     779              :     }
     780              : 
     781              :     /// Do post configuration of the already started Postgres. This function spawns a background thread to
     782              :     /// configure the database after applying the compute spec. Currently, it upgrades the neon extension
     783              :     /// version. In the future, it may upgrade all 3rd-party extensions.
     784            0 :     #[instrument(skip_all)]
     785              :     pub fn post_apply_config(&self) -> Result<()> {
     786              :         let connstr = self.connstr.clone();
     787            0 :         thread::spawn(move || {
     788            0 :             let func = || {
     789            0 :                 let mut client = Client::connect(connstr.as_str(), NoTls)?;
     790            0 :                 handle_neon_extension_upgrade(&mut client)
     791            0 :                     .context("handle_neon_extension_upgrade")?;
     792            0 :                 Ok::<_, anyhow::Error>(())
     793            0 :             };
     794            0 :             if let Err(err) = func() {
     795            0 :                 error!("error while post_apply_config: {err:#}");
     796            0 :             }
     797            0 :         });
     798              :         Ok(())
     799              :     }
     800              : 
     801              :     /// Do initial configuration of the already started Postgres.
     802            0 :     #[instrument(skip_all)]
     803              :     pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
     804              :         // If connection fails,
     805              :         // it may be the old node with `zenith_admin` superuser.
     806              :         //
     807              :         // In this case we need to connect with old `zenith_admin` name
     808              :         // and create new user. We cannot simply rename connected user,
     809              :         // but we can create a new one and grant it all privileges.
     810              :         let mut connstr = self.connstr.clone();
     811              :         connstr
     812              :             .query_pairs_mut()
     813              :             .append_pair("application_name", "apply_config");
     814              : 
     815              :         let mut client = match Client::connect(connstr.as_str(), NoTls) {
     816              :             Err(e) => match e.code() {
     817              :                 Some(&SqlState::INVALID_PASSWORD)
     818              :                 | Some(&SqlState::INVALID_AUTHORIZATION_SPECIFICATION) => {
     819              :                     // connect with zenith_admin if cloud_admin could not authenticate
     820              :                     info!(
     821              :                         "cannot connect to postgres: {}, retrying with `zenith_admin` username",
     822              :                         e
     823              :                     );
     824              :                     let mut zenith_admin_connstr = connstr.clone();
     825              : 
     826              :                     zenith_admin_connstr
     827              :                         .set_username("zenith_admin")
     828            0 :                         .map_err(|_| anyhow::anyhow!("invalid connstr"))?;
     829              : 
     830              :                     let mut client =
     831              :                         Client::connect(zenith_admin_connstr.as_str(), NoTls)
     832              :                             .context("broken cloud_admin credential: tried connecting with cloud_admin but could not authenticate, and zenith_admin does not work either")?;
     833              :                     // Disable forwarding so that users don't get a cloud_admin role
     834              : 
     835            0 :                     let mut func = || {
     836            0 :                         client.simple_query("SET neon.forward_ddl = false")?;
     837            0 :                         client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?;
     838            0 :                         client.simple_query("GRANT zenith_admin TO cloud_admin")?;
     839            0 :                         Ok::<_, anyhow::Error>(())
     840            0 :                     };
     841              :                     func().context("apply_config setup cloud_admin")?;
     842              : 
     843              :                     drop(client);
     844              : 
     845              :                     // reconnect with connstring with expected name
     846              :                     Client::connect(connstr.as_str(), NoTls)?
     847              :                 }
     848              :                 _ => return Err(e.into()),
     849              :             },
     850              :             Ok(client) => client,
     851              :         };
     852              : 
     853              :         // Disable DDL forwarding because control plane already knows about these roles/databases.
     854              :         client
     855              :             .simple_query("SET neon.forward_ddl = false")
     856              :             .context("apply_config SET neon.forward_ddl = false")?;
     857              : 
     858              :         // Proceed with post-startup configuration. Note, that order of operations is important.
     859              :         let spec = &compute_state.pspec.as_ref().expect("spec must be set").spec;
     860              :         create_neon_superuser(spec, &mut client).context("apply_config create_neon_superuser")?;
     861              :         cleanup_instance(&mut client).context("apply_config cleanup_instance")?;
     862              :         handle_roles(spec, &mut client).context("apply_config handle_roles")?;
     863              :         handle_databases(spec, &mut client).context("apply_config handle_databases")?;
     864              :         handle_role_deletions(spec, connstr.as_str(), &mut client)
     865              :             .context("apply_config handle_role_deletions")?;
     866              :         handle_grants(
     867              :             spec,
     868              :             &mut client,
     869              :             connstr.as_str(),
     870              :             self.has_feature(ComputeFeature::AnonExtension),
     871              :         )
     872              :         .context("apply_config handle_grants")?;
     873              :         handle_extensions(spec, &mut client).context("apply_config handle_extensions")?;
     874              :         handle_extension_neon(&mut client).context("apply_config handle_extension_neon")?;
     875              :         create_availability_check_data(&mut client)
     876              :             .context("apply_config create_availability_check_data")?;
     877              : 
     878              :         // 'Close' connection
     879              :         drop(client);
     880              : 
     881              :         // Run migrations separately to not hold up cold starts
     882            0 :         thread::spawn(move || {
     883            0 :             let mut connstr = connstr.clone();
     884            0 :             connstr
     885            0 :                 .query_pairs_mut()
     886            0 :                 .append_pair("application_name", "migrations");
     887              : 
     888            0 :             let mut client = Client::connect(connstr.as_str(), NoTls)?;
     889            0 :             handle_migrations(&mut client).context("apply_config handle_migrations")
     890            0 :         });
     891              :         Ok(())
     892              :     }
     893              : 
     894              :     // Wrapped this around `pg_ctl reload`, but right now we don't use
     895              :     // `pg_ctl` for start / stop.
     896            0 :     #[instrument(skip_all)]
     897              :     fn pg_reload_conf(&self) -> Result<()> {
     898              :         let pgctl_bin = Path::new(&self.pgbin).parent().unwrap().join("pg_ctl");
     899              :         Command::new(pgctl_bin)
     900              :             .args(["reload", "-D", &self.pgdata])
     901              :             .output()
     902              :             .expect("cannot run pg_ctl process");
     903              :         Ok(())
     904              :     }
     905              : 
     906              :     /// Similar to `apply_config()`, but does a bit different sequence of operations,
     907              :     /// as it's used to reconfigure a previously started and configured Postgres node.
     908            0 :     #[instrument(skip_all)]
     909              :     pub fn reconfigure(&self) -> Result<()> {
     910              :         let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
     911              : 
     912              :         if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
     913              :             info!("tuning pgbouncer");
     914              : 
     915              :             let rt = tokio::runtime::Builder::new_current_thread()
     916              :                 .enable_all()
     917              :                 .build()
     918              :                 .expect("failed to create rt");
     919              : 
     920              :             // Spawn a thread to do the tuning,
     921              :             // so that we don't block the main thread that starts Postgres.
     922              :             let pgbouncer_settings = pgbouncer_settings.clone();
     923            0 :             let _handle = thread::spawn(move || {
     924            0 :                 let res = rt.block_on(tune_pgbouncer(pgbouncer_settings));
     925            0 :                 if let Err(err) = res {
     926            0 :                     error!("error while tuning pgbouncer: {err:?}");
     927            0 :                 }
     928            0 :             });
     929              :         }
     930              : 
     931              :         // Write new config
     932              :         let pgdata_path = Path::new(&self.pgdata);
     933              :         let postgresql_conf_path = pgdata_path.join("postgresql.conf");
     934              :         config::write_postgres_conf(&postgresql_conf_path, &spec, None)?;
     935              :         // temporarily reset max_cluster_size in config
     936              :         // to avoid the possibility of hitting the limit, while we are reconfiguring:
     937              :         // creating new extensions, roles, etc...
     938            0 :         config::with_compute_ctl_tmp_override(pgdata_path, "neon.max_cluster_size=-1", || {
     939            0 :             self.pg_reload_conf()?;
     940              : 
     941            0 :             let mut client = Client::connect(self.connstr.as_str(), NoTls)?;
     942              : 
     943              :             // Proceed with post-startup configuration. Note, that order of operations is important.
     944              :             // Disable DDL forwarding because control plane already knows about these roles/databases.
     945            0 :             if spec.mode == ComputeMode::Primary {
     946            0 :                 client.simple_query("SET neon.forward_ddl = false")?;
     947            0 :                 cleanup_instance(&mut client)?;
     948            0 :                 handle_roles(&spec, &mut client)?;
     949            0 :                 handle_databases(&spec, &mut client)?;
     950            0 :                 handle_role_deletions(&spec, self.connstr.as_str(), &mut client)?;
     951            0 :                 handle_grants(
     952            0 :                     &spec,
     953            0 :                     &mut client,
     954            0 :                     self.connstr.as_str(),
     955            0 :                     self.has_feature(ComputeFeature::AnonExtension),
     956            0 :                 )?;
     957            0 :                 handle_extensions(&spec, &mut client)?;
     958            0 :                 handle_extension_neon(&mut client)?;
     959              :                 // We can skip handle_migrations here because a new migration can only appear
     960              :                 // if we have a new version of the compute_ctl binary, which can only happen
     961              :                 // if compute got restarted, in which case we'll end up inside of apply_config
     962              :                 // instead of reconfigure.
     963            0 :             }
     964              : 
     965              :             // 'Close' connection
     966            0 :             drop(client);
     967            0 : 
     968            0 :             Ok(())
     969            0 :         })?;
     970              : 
     971              :         self.pg_reload_conf()?;
     972              : 
     973              :         let unknown_op = "unknown".to_string();
     974              :         let op_id = spec.operation_uuid.as_ref().unwrap_or(&unknown_op);
     975              :         info!(
     976              :             "finished reconfiguration of compute node for operation {}",
     977              :             op_id
     978              :         );
     979              : 
     980              :         Ok(())
     981              :     }
     982              : 
     983            0 :     #[instrument(skip_all)]
     984              :     pub fn start_compute(
     985              :         &self,
     986              :         extension_server_port: u16,
     987              :     ) -> Result<(std::process::Child, std::thread::JoinHandle<()>)> {
     988              :         let compute_state = self.state.lock().unwrap().clone();
     989              :         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
     990              :         info!(
     991              :             "starting compute for project {}, operation {}, tenant {}, timeline {}",
     992              :             pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
     993              :             pspec.spec.operation_uuid.as_deref().unwrap_or("None"),
     994              :             pspec.tenant_id,
     995              :             pspec.timeline_id,
     996              :         );
     997              : 
     998              :         // tune pgbouncer
     999              :         if let Some(pgbouncer_settings) = &pspec.spec.pgbouncer_settings {
    1000              :             info!("tuning pgbouncer");
    1001              : 
    1002              :             let rt = tokio::runtime::Builder::new_current_thread()
    1003              :                 .enable_all()
    1004              :                 .build()
    1005              :                 .expect("failed to create rt");
    1006              : 
    1007              :             // Spawn a thread to do the tuning,
    1008              :             // so that we don't block the main thread that starts Postgres.
    1009              :             let pgbouncer_settings = pgbouncer_settings.clone();
    1010            0 :             let _handle = thread::spawn(move || {
    1011            0 :                 let res = rt.block_on(tune_pgbouncer(pgbouncer_settings));
    1012            0 :                 if let Err(err) = res {
    1013            0 :                     error!("error while tuning pgbouncer: {err:?}");
    1014            0 :                 }
    1015            0 :             });
    1016              :         }
    1017              : 
    1018              :         info!(
    1019              :             "start_compute spec.remote_extensions {:?}",
    1020              :             pspec.spec.remote_extensions
    1021              :         );
    1022              : 
    1023              :         // This part is sync, because we need to download
    1024              :         // remote shared_preload_libraries before postgres start (if any)
    1025              :         if let Some(remote_extensions) = &pspec.spec.remote_extensions {
    1026              :             // First, create control files for all availale extensions
    1027              :             extension_server::create_control_files(remote_extensions, &self.pgbin);
    1028              : 
    1029              :             let library_load_start_time = Utc::now();
    1030              :             let remote_ext_metrics = self.prepare_preload_libraries(&pspec.spec)?;
    1031              : 
    1032              :             let library_load_time = Utc::now()
    1033              :                 .signed_duration_since(library_load_start_time)
    1034              :                 .to_std()
    1035              :                 .unwrap()
    1036              :                 .as_millis() as u64;
    1037              :             let mut state = self.state.lock().unwrap();
    1038              :             state.metrics.load_ext_ms = library_load_time;
    1039              :             state.metrics.num_ext_downloaded = remote_ext_metrics.num_ext_downloaded;
    1040              :             state.metrics.largest_ext_size = remote_ext_metrics.largest_ext_size;
    1041              :             state.metrics.total_ext_download_size = remote_ext_metrics.total_ext_download_size;
    1042              :             info!(
    1043              :                 "Loading shared_preload_libraries took {:?}ms",
    1044              :                 library_load_time
    1045              :             );
    1046              :             info!("{:?}", remote_ext_metrics);
    1047              :         }
    1048              : 
    1049              :         self.prepare_pgdata(&compute_state, extension_server_port)?;
    1050              : 
    1051              :         let start_time = Utc::now();
    1052              :         let pg_process = self.start_postgres(pspec.storage_auth_token.clone())?;
    1053              : 
    1054              :         let config_time = Utc::now();
    1055              :         if pspec.spec.mode == ComputeMode::Primary && !pspec.spec.skip_pg_catalog_updates {
    1056              :             let pgdata_path = Path::new(&self.pgdata);
    1057              :             // temporarily reset max_cluster_size in config
    1058              :             // to avoid the possibility of hitting the limit, while we are applying config:
    1059              :             // creating new extensions, roles, etc...
    1060            0 :             config::with_compute_ctl_tmp_override(pgdata_path, "neon.max_cluster_size=-1", || {
    1061            0 :                 self.pg_reload_conf()?;
    1062              : 
    1063            0 :                 self.apply_config(&compute_state)?;
    1064              : 
    1065            0 :                 Ok(())
    1066            0 :             })?;
    1067              :             self.pg_reload_conf()?;
    1068              :         }
    1069              : 
    1070              :         let startup_end_time = Utc::now();
    1071              :         {
    1072              :             let mut state = self.state.lock().unwrap();
    1073              :             state.metrics.start_postgres_ms = config_time
    1074              :                 .signed_duration_since(start_time)
    1075              :                 .to_std()
    1076              :                 .unwrap()
    1077              :                 .as_millis() as u64;
    1078              :             state.metrics.config_ms = startup_end_time
    1079              :                 .signed_duration_since(config_time)
    1080              :                 .to_std()
    1081              :                 .unwrap()
    1082              :                 .as_millis() as u64;
    1083              :             state.metrics.total_startup_ms = startup_end_time
    1084              :                 .signed_duration_since(compute_state.start_time)
    1085              :                 .to_std()
    1086              :                 .unwrap()
    1087              :                 .as_millis() as u64;
    1088              :         }
    1089              :         self.set_status(ComputeStatus::Running);
    1090              : 
    1091              :         info!(
    1092              :             "finished configuration of compute for project {}",
    1093              :             pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None")
    1094              :         );
    1095              : 
    1096              :         // Log metrics so that we can search for slow operations in logs
    1097              :         let metrics = {
    1098              :             let state = self.state.lock().unwrap();
    1099              :             state.metrics.clone()
    1100              :         };
    1101              :         info!(?metrics, "compute start finished");
    1102              : 
    1103              :         Ok(pg_process)
    1104              :     }
    1105              : 
    1106              :     /// Update the `last_active` in the shared state, but ensure that it's a more recent one.
    1107            0 :     pub fn update_last_active(&self, last_active: Option<DateTime<Utc>>) {
    1108            0 :         let mut state = self.state.lock().unwrap();
    1109            0 :         // NB: `Some(<DateTime>)` is always greater than `None`.
    1110            0 :         if last_active > state.last_active {
    1111            0 :             state.last_active = last_active;
    1112            0 :             debug!("set the last compute activity time to: {:?}", last_active);
    1113            0 :         }
    1114            0 :     }
    1115              : 
    1116              :     // Look for core dumps and collect backtraces.
    1117              :     //
    1118              :     // EKS worker nodes have following core dump settings:
    1119              :     //   /proc/sys/kernel/core_pattern -> core
    1120              :     //   /proc/sys/kernel/core_uses_pid -> 1
    1121              :     //   ulimit -c -> unlimited
    1122              :     // which results in core dumps being written to postgres data directory as core.<pid>.
    1123              :     //
    1124              :     // Use that as a default location and pattern, except macos where core dumps are written
    1125              :     // to /cores/ directory by default.
    1126            0 :     pub fn check_for_core_dumps(&self) -> Result<()> {
    1127            0 :         let core_dump_dir = match std::env::consts::OS {
    1128            0 :             "macos" => Path::new("/cores/"),
    1129            0 :             _ => Path::new(&self.pgdata),
    1130              :         };
    1131              : 
    1132              :         // Collect core dump paths if any
    1133            0 :         info!("checking for core dumps in {}", core_dump_dir.display());
    1134            0 :         let files = fs::read_dir(core_dump_dir)?;
    1135            0 :         let cores = files.filter_map(|entry| {
    1136            0 :             let entry = entry.ok()?;
    1137            0 :             let _ = entry.file_name().to_str()?.strip_prefix("core.")?;
    1138            0 :             Some(entry.path())
    1139            0 :         });
    1140              : 
    1141              :         // Print backtrace for each core dump
    1142            0 :         for core_path in cores {
    1143            0 :             warn!(
    1144            0 :                 "core dump found: {}, collecting backtrace",
    1145            0 :                 core_path.display()
    1146              :             );
    1147              : 
    1148              :             // Try first with gdb
    1149            0 :             let backtrace = Command::new("gdb")
    1150            0 :                 .args(["--batch", "-q", "-ex", "bt", &self.pgbin])
    1151            0 :                 .arg(&core_path)
    1152            0 :                 .output();
    1153              : 
    1154              :             // Try lldb if no gdb is found -- that is handy for local testing on macOS
    1155            0 :             let backtrace = match backtrace {
    1156            0 :                 Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
    1157            0 :                     warn!("cannot find gdb, trying lldb");
    1158            0 :                     Command::new("lldb")
    1159            0 :                         .arg("-c")
    1160            0 :                         .arg(&core_path)
    1161            0 :                         .args(["--batch", "-o", "bt all", "-o", "quit"])
    1162            0 :                         .output()
    1163              :                 }
    1164            0 :                 _ => backtrace,
    1165            0 :             }?;
    1166              : 
    1167            0 :             warn!(
    1168            0 :                 "core dump backtrace: {}",
    1169            0 :                 String::from_utf8_lossy(&backtrace.stdout)
    1170              :             );
    1171            0 :             warn!(
    1172            0 :                 "debugger stderr: {}",
    1173            0 :                 String::from_utf8_lossy(&backtrace.stderr)
    1174              :             );
    1175              :         }
    1176              : 
    1177            0 :         Ok(())
    1178            0 :     }
    1179              : 
    1180              :     /// Select `pg_stat_statements` data and return it as a stringified JSON
    1181            0 :     pub async fn collect_insights(&self) -> String {
    1182            0 :         let mut result_rows: Vec<String> = Vec::new();
    1183            0 :         let connect_result = tokio_postgres::connect(self.connstr.as_str(), NoTls).await;
    1184            0 :         let (client, connection) = connect_result.unwrap();
    1185            0 :         tokio::spawn(async move {
    1186            0 :             if let Err(e) = connection.await {
    1187            0 :                 eprintln!("connection error: {}", e);
    1188            0 :             }
    1189            0 :         });
    1190            0 :         let result = client
    1191            0 :             .simple_query(
    1192            0 :                 "SELECT
    1193            0 :     row_to_json(pg_stat_statements)
    1194            0 : FROM
    1195            0 :     pg_stat_statements
    1196            0 : WHERE
    1197            0 :     userid != 'cloud_admin'::regrole::oid
    1198            0 : ORDER BY
    1199            0 :     (mean_exec_time + mean_plan_time) DESC
    1200            0 : LIMIT 100",
    1201            0 :             )
    1202            0 :             .await;
    1203              : 
    1204            0 :         if let Ok(raw_rows) = result {
    1205            0 :             for message in raw_rows.iter() {
    1206            0 :                 if let postgres::SimpleQueryMessage::Row(row) = message {
    1207            0 :                     if let Some(json) = row.get(0) {
    1208            0 :                         result_rows.push(json.to_string());
    1209            0 :                     }
    1210            0 :                 }
    1211              :             }
    1212              : 
    1213            0 :             format!("{{\"pg_stat_statements\": [{}]}}", result_rows.join(","))
    1214              :         } else {
    1215            0 :             "{{\"pg_stat_statements\": []}}".to_string()
    1216              :         }
    1217            0 :     }
    1218              : 
    1219              :     // download an archive, unzip and place files in correct locations
    1220            0 :     pub async fn download_extension(
    1221            0 :         &self,
    1222            0 :         real_ext_name: String,
    1223            0 :         ext_path: RemotePath,
    1224            0 :     ) -> Result<u64, DownloadError> {
    1225            0 :         let ext_remote_storage =
    1226            0 :             self.ext_remote_storage
    1227            0 :                 .as_ref()
    1228            0 :                 .ok_or(DownloadError::BadInput(anyhow::anyhow!(
    1229            0 :                     "Remote extensions storage is not configured",
    1230            0 :                 )))?;
    1231              : 
    1232            0 :         let ext_archive_name = ext_path.object_name().expect("bad path");
    1233            0 : 
    1234            0 :         let mut first_try = false;
    1235            0 :         if !self
    1236            0 :             .ext_download_progress
    1237            0 :             .read()
    1238            0 :             .expect("lock err")
    1239            0 :             .contains_key(ext_archive_name)
    1240            0 :         {
    1241            0 :             self.ext_download_progress
    1242            0 :                 .write()
    1243            0 :                 .expect("lock err")
    1244            0 :                 .insert(ext_archive_name.to_string(), (Utc::now(), false));
    1245            0 :             first_try = true;
    1246            0 :         }
    1247            0 :         let (download_start, download_completed) =
    1248            0 :             self.ext_download_progress.read().expect("lock err")[ext_archive_name];
    1249            0 :         let start_time_delta = Utc::now()
    1250            0 :             .signed_duration_since(download_start)
    1251            0 :             .to_std()
    1252            0 :             .unwrap()
    1253            0 :             .as_millis() as u64;
    1254              : 
    1255              :         // how long to wait for extension download if it was started by another process
    1256              :         const HANG_TIMEOUT: u64 = 3000; // milliseconds
    1257              : 
    1258            0 :         if download_completed {
    1259            0 :             info!("extension already downloaded, skipping re-download");
    1260            0 :             return Ok(0);
    1261            0 :         } else if start_time_delta < HANG_TIMEOUT && !first_try {
    1262            0 :             info!("download {ext_archive_name} already started by another process, hanging untill completion or timeout");
    1263            0 :             let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500));
    1264              :             loop {
    1265            0 :                 info!("waiting for download");
    1266            0 :                 interval.tick().await;
    1267            0 :                 let (_, download_completed_now) =
    1268            0 :                     self.ext_download_progress.read().expect("lock")[ext_archive_name];
    1269            0 :                 if download_completed_now {
    1270            0 :                     info!("download finished by whoever else downloaded it");
    1271            0 :                     return Ok(0);
    1272            0 :                 }
    1273              :             }
    1274              :             // NOTE: the above loop will get terminated
    1275              :             // based on the timeout of the download function
    1276            0 :         }
    1277            0 : 
    1278            0 :         // if extension hasn't been downloaded before or the previous
    1279            0 :         // attempt to download was at least HANG_TIMEOUT ms ago
    1280            0 :         // then we try to download it here
    1281            0 :         info!("downloading new extension {ext_archive_name}");
    1282              : 
    1283            0 :         let download_size = extension_server::download_extension(
    1284            0 :             &real_ext_name,
    1285            0 :             &ext_path,
    1286            0 :             ext_remote_storage,
    1287            0 :             &self.pgbin,
    1288            0 :         )
    1289            0 :         .await
    1290            0 :         .map_err(DownloadError::Other);
    1291            0 : 
    1292            0 :         if download_size.is_ok() {
    1293            0 :             self.ext_download_progress
    1294            0 :                 .write()
    1295            0 :                 .expect("bad lock")
    1296            0 :                 .insert(ext_archive_name.to_string(), (download_start, true));
    1297            0 :         }
    1298              : 
    1299            0 :         download_size
    1300            0 :     }
    1301              : 
    1302              :     #[tokio::main]
    1303            0 :     pub async fn prepare_preload_libraries(
    1304            0 :         &self,
    1305            0 :         spec: &ComputeSpec,
    1306            0 :     ) -> Result<RemoteExtensionMetrics> {
    1307            0 :         if self.ext_remote_storage.is_none() {
    1308            0 :             return Ok(RemoteExtensionMetrics {
    1309            0 :                 num_ext_downloaded: 0,
    1310            0 :                 largest_ext_size: 0,
    1311            0 :                 total_ext_download_size: 0,
    1312            0 :             });
    1313            0 :         }
    1314            0 :         let remote_extensions = spec
    1315            0 :             .remote_extensions
    1316            0 :             .as_ref()
    1317            0 :             .ok_or(anyhow::anyhow!("Remote extensions are not configured"))?;
    1318            0 : 
    1319            0 :         info!("parse shared_preload_libraries from spec.cluster.settings");
    1320            0 :         let mut libs_vec = Vec::new();
    1321            0 :         if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
    1322            0 :             libs_vec = libs
    1323            0 :                 .split(&[',', '\'', ' '])
    1324            0 :                 .filter(|s| *s != "neon" && !s.is_empty())
    1325            0 :                 .map(str::to_string)
    1326            0 :                 .collect();
    1327            0 :         }
    1328            0 :         info!("parse shared_preload_libraries from provided postgresql.conf");
    1329            0 : 
    1330            0 :         // that is used in neon_local and python tests
    1331            0 :         if let Some(conf) = &spec.cluster.postgresql_conf {
    1332            0 :             let conf_lines = conf.split('\n').collect::<Vec<&str>>();
    1333            0 :             let mut shared_preload_libraries_line = "";
    1334            0 :             for line in conf_lines {
    1335            0 :                 if line.starts_with("shared_preload_libraries") {
    1336            0 :                     shared_preload_libraries_line = line;
    1337            0 :                 }
    1338            0 :             }
    1339            0 :             let mut preload_libs_vec = Vec::new();
    1340            0 :             if let Some(libs) = shared_preload_libraries_line.split("='").nth(1) {
    1341            0 :                 preload_libs_vec = libs
    1342            0 :                     .split(&[',', '\'', ' '])
    1343            0 :                     .filter(|s| *s != "neon" && !s.is_empty())
    1344            0 :                     .map(str::to_string)
    1345            0 :                     .collect();
    1346            0 :             }
    1347            0 :             libs_vec.extend(preload_libs_vec);
    1348            0 :         }
    1349            0 : 
    1350            0 :         // Don't try to download libraries that are not in the index.
    1351            0 :         // Assume that they are already present locally.
    1352            0 :         libs_vec.retain(|lib| remote_extensions.library_index.contains_key(lib));
    1353            0 : 
    1354            0 :         info!("Downloading to shared preload libraries: {:?}", &libs_vec);
    1355            0 : 
    1356            0 :         let mut download_tasks = Vec::new();
    1357            0 :         for library in &libs_vec {
    1358            0 :             let (ext_name, ext_path) =
    1359            0 :                 remote_extensions.get_ext(library, true, &self.build_tag, &self.pgversion)?;
    1360            0 :             download_tasks.push(self.download_extension(ext_name, ext_path));
    1361            0 :         }
    1362            0 :         let results = join_all(download_tasks).await;
    1363            0 : 
    1364            0 :         let mut remote_ext_metrics = RemoteExtensionMetrics {
    1365            0 :             num_ext_downloaded: 0,
    1366            0 :             largest_ext_size: 0,
    1367            0 :             total_ext_download_size: 0,
    1368            0 :         };
    1369            0 :         for result in results {
    1370            0 :             let download_size = match result {
    1371            0 :                 Ok(res) => {
    1372            0 :                     remote_ext_metrics.num_ext_downloaded += 1;
    1373            0 :                     res
    1374            0 :                 }
    1375            0 :                 Err(err) => {
    1376            0 :                     // if we failed to download an extension, we don't want to fail the whole
    1377            0 :                     // process, but we do want to log the error
    1378            0 :                     error!("Failed to download extension: {}", err);
    1379            0 :                     0
    1380            0 :                 }
    1381            0 :             };
    1382            0 : 
    1383            0 :             remote_ext_metrics.largest_ext_size =
    1384            0 :                 std::cmp::max(remote_ext_metrics.largest_ext_size, download_size);
    1385            0 :             remote_ext_metrics.total_ext_download_size += download_size;
    1386            0 :         }
    1387            0 :         Ok(remote_ext_metrics)
    1388            0 :     }
    1389              : }
    1390              : 
    1391            0 : pub fn forward_termination_signal() {
    1392            0 :     let ss_pid = SYNC_SAFEKEEPERS_PID.load(Ordering::SeqCst);
    1393            0 :     if ss_pid != 0 {
    1394            0 :         let ss_pid = nix::unistd::Pid::from_raw(ss_pid as i32);
    1395            0 :         kill(ss_pid, Signal::SIGTERM).ok();
    1396            0 :     }
    1397            0 :     let pg_pid = PG_PID.load(Ordering::SeqCst);
    1398            0 :     if pg_pid != 0 {
    1399            0 :         let pg_pid = nix::unistd::Pid::from_raw(pg_pid as i32);
    1400            0 :         // Use 'fast' shutdown (SIGINT) because it also creates a shutdown checkpoint, which is important for
    1401            0 :         // ROs to get a list of running xacts faster instead of going through the CLOG.
    1402            0 :         // See https://www.postgresql.org/docs/current/server-shutdown.html for the list of modes and signals.
    1403            0 :         kill(pg_pid, Signal::SIGINT).ok();
    1404            0 :     }
    1405            0 : }
        

Generated by: LCOV version 2.1-beta