LCOV - code coverage report
Current view: top level - compute_tools/src/bin - compute_ctl.rs (source / functions) Coverage Total Hit
Test: 1e20c4f2b28aa592527961bb32170ebbd2c9172f.info Lines: 52.4 % 170 89
Test Date: 2025-07-16 12:29:03 Functions: 38.5 % 13 5

            Line data    Source code
       1              : //!
       2              : //! Postgres wrapper (`compute_ctl`) is intended to be run as a Docker entrypoint or as a `systemd`
       3              : //! `ExecStart` option. It will handle all the `Neon` specifics during compute node
       4              : //! initialization:
       5              : //! - `compute_ctl` accepts cluster (compute node) specification as a JSON file.
       6              : //! - Every start is a fresh start, so the data directory is removed and
       7              : //!   initialized again on each run.
       8              : //! - If remote_extension_config is provided, it will be used to fetch extensions list
       9              : //!   and download `shared_preload_libraries` from the remote storage.
      10              : //! - Next it will put configuration files into the `PGDATA` directory.
      11              : //! - Sync safekeepers and get commit LSN.
      12              : //! - Get `basebackup` from pageserver using the returned on the previous step LSN.
      13              : //! - Try to start `postgres` and wait until it is ready to accept connections.
      14              : //! - Check and alter/drop/create roles and databases.
      15              : //! - Hang waiting on the `postmaster` process to exit.
      16              : //!
      17              : //! Also `compute_ctl` spawns two separate service threads:
      18              : //! - `compute-monitor` checks the last Postgres activity timestamp and saves it
      19              : //!   into the shared `ComputeNode`;
      20              : //! - `http-endpoint` runs a Hyper HTTP API server, which serves readiness and the
      21              : //!   last activity requests.
      22              : //!
      23              : //! If `AUTOSCALING` environment variable is set, `compute_ctl` will start the
      24              : //! `vm-monitor` located in [`neon/libs/vm_monitor`]. For VM compute nodes,
      25              : //! `vm-monitor` communicates with the VM autoscaling system. It coordinates
      26              : //! downscaling and requests immediate upscaling under resource pressure.
      27              : //!
      28              : //! Usage example:
      29              : //! ```sh
      30              : //! compute_ctl -D /var/db/postgres/compute \
      31              : //!             -C 'postgresql://cloud_admin@localhost/postgres' \
      32              : //!             -c /var/db/postgres/configs/config.json \
      33              : //!             -b /usr/local/bin/postgres \
      34              : //!             -r http://pg-ext-s3-gateway \
      35              : //! ```
      36              : use std::ffi::OsString;
      37              : use std::fs::File;
      38              : use std::process::exit;
      39              : use std::sync::Arc;
      40              : use std::sync::atomic::AtomicU64;
      41              : use std::sync::mpsc;
      42              : use std::thread;
      43              : use std::time::Duration;
      44              : 
      45              : use anyhow::{Context, Result, bail};
      46              : use clap::Parser;
      47              : use compute_api::responses::ComputeConfig;
      48              : use compute_tools::compute::{
      49              :     BUILD_TAG, ComputeNode, ComputeNodeParams, forward_termination_signal,
      50              : };
      51              : use compute_tools::extension_server::get_pg_version_string;
      52              : use compute_tools::logger::*;
      53              : use compute_tools::params::*;
      54              : use compute_tools::spec::*;
      55              : use rlimit::{Resource, setrlimit};
      56              : use signal_hook::consts::{SIGINT, SIGQUIT, SIGTERM};
      57              : use signal_hook::iterator::Signals;
      58              : use tracing::{error, info};
      59              : use url::Url;
      60              : use utils::failpoint_support;
      61              : 
      62              : #[derive(Debug, Parser)]
      63              : #[command(rename_all = "kebab-case")]
      64              : struct Cli {
      65              :     #[arg(short = 'b', long, default_value = "postgres", env = "POSTGRES_PATH")]
      66              :     pub pgbin: String,
      67              : 
      68              :     /// The base URL for the remote extension storage proxy gateway.
      69              :     #[arg(short = 'r', long, value_parser = Self::parse_remote_ext_base_url)]
      70              :     pub remote_ext_base_url: Option<Url>,
      71              : 
      72              :     /// The port to bind the external listening HTTP server to. Clients running
      73              :     /// outside the compute will talk to the compute through this port. Keep
      74              :     /// the previous name for this argument around for a smoother release
      75              :     /// with the control plane.
      76              :     #[arg(long, default_value_t = 3080)]
      77              :     pub external_http_port: u16,
      78              : 
      79              :     /// The port to bind the internal listening HTTP server to. Clients include
      80              :     /// the neon extension (for installing remote extensions) and local_proxy.
      81              :     #[arg(long, default_value_t = 3081)]
      82              :     pub internal_http_port: u16,
      83              : 
      84              :     #[arg(short = 'D', long, value_name = "DATADIR")]
      85              :     pub pgdata: String,
      86              : 
      87              :     #[arg(short = 'C', long, value_name = "DATABASE_URL")]
      88              :     pub connstr: String,
      89              : 
      90              :     #[arg(
      91              :         long,
      92              :         default_value = "neon_superuser",
      93              :         value_name = "PRIVILEGED_ROLE_NAME",
      94              :         value_parser = Self::parse_privileged_role_name
      95              :     )]
      96              :     pub privileged_role_name: String,
      97              : 
      98              :     #[cfg(target_os = "linux")]
      99              :     #[arg(long, default_value = "neon-postgres")]
     100              :     pub cgroup: String,
     101              : 
     102              :     #[cfg(target_os = "linux")]
     103              :     #[arg(
     104              :         long,
     105              :         default_value = "host=localhost port=5432 dbname=postgres user=cloud_admin sslmode=disable application_name=vm-monitor"
     106              :     )]
     107              :     pub filecache_connstr: String,
     108              : 
     109              :     #[cfg(target_os = "linux")]
     110              :     #[arg(long, default_value = "0.0.0.0:10301")]
     111              :     pub vm_monitor_addr: String,
     112              : 
     113              :     #[arg(long, action = clap::ArgAction::SetTrue)]
     114              :     pub resize_swap_on_bind: bool,
     115              : 
     116              :     #[arg(long)]
     117              :     pub set_disk_quota_for_fs: Option<String>,
     118              : 
     119              :     #[arg(short = 'c', long)]
     120              :     pub config: Option<OsString>,
     121              : 
     122              :     #[arg(short = 'i', long, group = "compute-id")]
     123              :     pub compute_id: String,
     124              : 
     125              :     #[arg(
     126              :         short = 'p',
     127              :         long,
     128              :         conflicts_with = "config",
     129              :         value_name = "CONTROL_PLANE_API_BASE_URL",
     130              :         requires = "compute-id"
     131              :     )]
     132              :     pub control_plane_uri: Option<String>,
     133              : 
     134              :     /// Interval in seconds for collecting installed extensions statistics
     135              :     #[arg(long, default_value = "3600")]
     136              :     pub installed_extensions_collection_interval: u64,
     137              : 
     138              :     /// Run in development mode, skipping VM-specific operations like process termination
     139              :     #[arg(long, action = clap::ArgAction::SetTrue)]
     140              :     pub dev: bool,
     141              : }
     142              : 
     143              : impl Cli {
     144              :     /// Parse a URL from an argument. By default, this isn't necessary, but we
     145              :     /// want to do some sanity checking.
     146            3 :     fn parse_remote_ext_base_url(value: &str) -> Result<Url> {
     147              :         // Remove extra trailing slashes, and add one. We use Url::join() later
     148              :         // when downloading remote extensions. If the base URL is something like
     149              :         // http://example.com/pg-ext-s3-gateway, and join() is called with
     150              :         // something like "xyz", the resulting URL is http://example.com/xyz.
     151            3 :         let value = value.trim_end_matches('/').to_owned() + "/";
     152            3 :         let url = Url::parse(&value)?;
     153              : 
     154            3 :         if url.query_pairs().count() != 0 {
     155            1 :             bail!("parameters detected in remote extensions base URL")
     156            2 :         }
     157              : 
     158            2 :         Ok(url)
     159            3 :     }
     160              : 
     161              :     /// For simplicity, we do not escape `privileged_role_name` anywhere in the code.
     162              :     /// Since it's a system role, which we fully control, that's fine. Still, let's
     163              :     /// validate it to avoid any surprises.
     164            6 :     fn parse_privileged_role_name(value: &str) -> Result<String> {
     165              :         use regex::Regex;
     166              : 
     167            6 :         let pattern = Regex::new(r"^[a-z_]+$").unwrap();
     168              : 
     169            6 :         if !pattern.is_match(value) {
     170            3 :             bail!("--privileged-role-name can only contain lowercase letters and underscores")
     171            3 :         }
     172              : 
     173            3 :         Ok(value.to_string())
     174            6 :     }
     175              : }
     176              : 
     177            0 : fn main() -> Result<()> {
     178            0 :     let cli = Cli::parse();
     179              : 
     180            0 :     let scenario = failpoint_support::init();
     181              : 
     182              :     // For historical reasons, the main thread that processes the config and launches postgres
     183              :     // is synchronous, but we always have this tokio runtime available and we "enter" it so
     184              :     // that you can use tokio::spawn() and tokio::runtime::Handle::current().block_on(...)
     185              :     // from all parts of compute_ctl.
     186            0 :     let runtime = tokio::runtime::Builder::new_multi_thread()
     187            0 :         .enable_all()
     188            0 :         .build()?;
     189            0 :     let _rt_guard = runtime.enter();
     190              : 
     191            0 :     runtime.block_on(init(cli.dev))?;
     192              : 
     193              :     // enable core dumping for all child processes
     194            0 :     setrlimit(Resource::CORE, rlimit::INFINITY, rlimit::INFINITY)?;
     195              : 
     196            0 :     let connstr = Url::parse(&cli.connstr).context("cannot parse connstr as a URL")?;
     197              : 
     198            0 :     let config = get_config(&cli)?;
     199              : 
     200            0 :     let compute_node = ComputeNode::new(
     201            0 :         ComputeNodeParams {
     202            0 :             compute_id: cli.compute_id,
     203            0 :             connstr,
     204            0 :             privileged_role_name: cli.privileged_role_name.clone(),
     205            0 :             pgdata: cli.pgdata.clone(),
     206            0 :             pgbin: cli.pgbin.clone(),
     207            0 :             pgversion: get_pg_version_string(&cli.pgbin),
     208            0 :             external_http_port: cli.external_http_port,
     209            0 :             internal_http_port: cli.internal_http_port,
     210            0 :             remote_ext_base_url: cli.remote_ext_base_url.clone(),
     211            0 :             resize_swap_on_bind: cli.resize_swap_on_bind,
     212            0 :             set_disk_quota_for_fs: cli.set_disk_quota_for_fs,
     213            0 :             #[cfg(target_os = "linux")]
     214            0 :             filecache_connstr: cli.filecache_connstr,
     215            0 :             #[cfg(target_os = "linux")]
     216            0 :             cgroup: cli.cgroup,
     217            0 :             #[cfg(target_os = "linux")]
     218            0 :             vm_monitor_addr: cli.vm_monitor_addr,
     219            0 :             installed_extensions_collection_interval: Arc::new(AtomicU64::new(
     220            0 :                 cli.installed_extensions_collection_interval,
     221            0 :             )),
     222            0 :         },
     223            0 :         config,
     224            0 :     )?;
     225              : 
     226            0 :     let exit_code = compute_node.run()?;
     227              : 
     228            0 :     scenario.teardown();
     229              : 
     230            0 :     deinit_and_exit(exit_code);
     231            0 : }
     232              : 
     233            0 : async fn init(dev_mode: bool) -> Result<()> {
     234            0 :     init_tracing_and_logging(DEFAULT_LOG_LEVEL).await?;
     235              : 
     236            0 :     let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT])?;
     237            0 :     thread::spawn(move || {
     238            0 :         for sig in signals.forever() {
     239            0 :             handle_exit_signal(sig, dev_mode);
     240            0 :         }
     241            0 :     });
     242              : 
     243            0 :     info!("compute build_tag: {}", &BUILD_TAG.to_string());
     244              : 
     245            0 :     Ok(())
     246            0 : }
     247              : 
     248            0 : fn get_config(cli: &Cli) -> Result<ComputeConfig> {
     249              :     // First, read the config from the path if provided
     250            0 :     if let Some(ref config) = cli.config {
     251            0 :         let file = File::open(config)?;
     252            0 :         return Ok(serde_json::from_reader(&file)?);
     253            0 :     }
     254              : 
     255              :     // If the config wasn't provided in the CLI arguments, then retrieve it from
     256              :     // the control plane
     257            0 :     match get_config_from_control_plane(cli.control_plane_uri.as_ref().unwrap(), &cli.compute_id) {
     258            0 :         Ok(config) => Ok(config),
     259            0 :         Err(e) => {
     260            0 :             error!(
     261            0 :                 "cannot get response from control plane: {}\n\
     262            0 :                 neither spec nor confirmation that compute is in the Empty state was received",
     263              :                 e
     264              :             );
     265            0 :             Err(e)
     266              :         }
     267              :     }
     268            0 : }
     269              : 
     270            0 : fn deinit_and_exit(exit_code: Option<i32>) -> ! {
     271              :     // Shutdown trace pipeline gracefully, so that it has a chance to send any
     272              :     // pending traces before we exit. Shutting down OTEL tracing provider may
     273              :     // hang for quite some time, see, for example:
     274              :     // - https://github.com/open-telemetry/opentelemetry-rust/issues/868
     275              :     // - and our problems with staging https://github.com/neondatabase/cloud/issues/3707#issuecomment-1493983636
     276              :     //
     277              :     // Yet, we want computes to shut down fast enough, as we may need a new one
     278              :     // for the same timeline ASAP. So wait no longer than 2s for the shutdown to
     279              :     // complete, then just error out and exit the main thread.
     280            0 :     info!("shutting down tracing");
     281            0 :     let (sender, receiver) = mpsc::channel();
     282            0 :     let _ = thread::spawn(move || {
     283            0 :         tracing_utils::shutdown_tracing();
     284            0 :         sender.send(()).ok()
     285            0 :     });
     286            0 :     let shutdown_res = receiver.recv_timeout(Duration::from_millis(2000));
     287            0 :     if shutdown_res.is_err() {
     288            0 :         error!("timed out while shutting down tracing, exiting anyway");
     289            0 :     }
     290              : 
     291            0 :     info!("shutting down");
     292            0 :     exit(exit_code.unwrap_or(1))
     293              : }
     294              : 
     295              : /// When compute_ctl is killed, send also termination signal to sync-safekeepers
     296              : /// to prevent leakage. TODO: it is better to convert compute_ctl to async and
     297              : /// wait for termination which would be easy then.
     298            0 : fn handle_exit_signal(sig: i32, dev_mode: bool) {
     299            0 :     info!("received {sig} termination signal");
     300            0 :     forward_termination_signal(dev_mode);
     301            0 :     exit(1);
     302              : }
     303              : 
     304              : #[cfg(test)]
     305              : mod test {
     306              :     use clap::{CommandFactory, Parser};
     307              :     use url::Url;
     308              : 
     309              :     use super::Cli;
     310              : 
     311              :     #[test]
     312            1 :     fn verify_cli() {
     313            1 :         Cli::command().debug_assert()
     314            1 :     }
     315              : 
     316              :     #[test]
     317            1 :     fn verify_remote_ext_base_url() {
     318            1 :         let cli = Cli::parse_from([
     319            1 :             "compute_ctl",
     320            1 :             "--pgdata=test",
     321            1 :             "--connstr=test",
     322            1 :             "--compute-id=test",
     323            1 :             "--remote-ext-base-url",
     324            1 :             "https://example.com/subpath",
     325            1 :         ]);
     326            1 :         assert_eq!(
     327            1 :             cli.remote_ext_base_url.unwrap(),
     328            1 :             Url::parse("https://example.com/subpath/").unwrap()
     329              :         );
     330              : 
     331            1 :         let cli = Cli::parse_from([
     332            1 :             "compute_ctl",
     333            1 :             "--pgdata=test",
     334            1 :             "--connstr=test",
     335            1 :             "--compute-id=test",
     336            1 :             "--remote-ext-base-url",
     337            1 :             "https://example.com//",
     338            1 :         ]);
     339            1 :         assert_eq!(
     340            1 :             cli.remote_ext_base_url.unwrap(),
     341            1 :             Url::parse("https://example.com").unwrap()
     342              :         );
     343              : 
     344            1 :         Cli::try_parse_from([
     345            1 :             "compute_ctl",
     346            1 :             "--pgdata=test",
     347            1 :             "--connstr=test",
     348            1 :             "--compute-id=test",
     349            1 :             "--remote-ext-base-url",
     350            1 :             "https://example.com?hello=world",
     351            1 :         ])
     352            1 :         .expect_err("URL parameters are not allowed");
     353            1 :     }
     354              : 
     355              :     #[test]
     356            1 :     fn verify_privileged_role_name() {
     357              :         // Valid name
     358            1 :         let cli = Cli::parse_from([
     359            1 :             "compute_ctl",
     360            1 :             "--pgdata=test",
     361            1 :             "--connstr=test",
     362            1 :             "--compute-id=test",
     363            1 :             "--privileged-role-name",
     364            1 :             "my_superuser",
     365            1 :         ]);
     366            1 :         assert_eq!(cli.privileged_role_name, "my_superuser");
     367              : 
     368              :         // Invalid names
     369            1 :         Cli::try_parse_from([
     370            1 :             "compute_ctl",
     371            1 :             "--pgdata=test",
     372            1 :             "--connstr=test",
     373            1 :             "--compute-id=test",
     374            1 :             "--privileged-role-name",
     375            1 :             "NeonSuperuser",
     376            1 :         ])
     377            1 :         .expect_err("uppercase letters are not allowed");
     378              : 
     379            1 :         Cli::try_parse_from([
     380            1 :             "compute_ctl",
     381            1 :             "--pgdata=test",
     382            1 :             "--connstr=test",
     383            1 :             "--compute-id=test",
     384            1 :             "--privileged-role-name",
     385            1 :             "$'neon_superuser",
     386            1 :         ])
     387            1 :         .expect_err("special characters are not allowed");
     388              : 
     389            1 :         Cli::try_parse_from([
     390            1 :             "compute_ctl",
     391            1 :             "--pgdata=test",
     392            1 :             "--connstr=test",
     393            1 :             "--compute-id=test",
     394            1 :             "--privileged-role-name",
     395            1 :             "",
     396            1 :         ])
     397            1 :         .expect_err("empty name is not allowed");
     398            1 :     }
     399              : }
        

Generated by: LCOV version 2.1-beta