LCOV - code coverage report
Current view: top level - control_plane/src - safekeeper.rs (source / functions) Coverage Total Hit
Test: a1cc1f33dc9899e4da66eb51e44e911a4b3bd648.info Lines: 0.0 % 186 0
Test Date: 2025-07-31 11:35:14 Functions: 0.0 % 21 0

            Line data    Source code
       1              : //! Code to manage safekeepers
       2              : //!
       3              : //! In the local test environment, the data for each safekeeper is stored in
       4              : //!
       5              : //! ```text
       6              : //!   .neon/safekeepers/<safekeeper id>
       7              : //! ```
       8              : use std::error::Error as _;
       9              : use std::io::Write;
      10              : use std::path::PathBuf;
      11              : use std::time::Duration;
      12              : use std::{io, result};
      13              : 
      14              : use anyhow::Context;
      15              : use camino::Utf8PathBuf;
      16              : use postgres_backend::AuthType;
      17              : use postgres_connection::PgConnectionConfig;
      18              : use safekeeper_api::models::TimelineCreateRequest;
      19              : use safekeeper_client::mgmt_api;
      20              : use thiserror::Error;
      21              : use utils::auth::{Claims, Scope};
      22              : use utils::id::NodeId;
      23              : 
      24              : use crate::background_process;
      25              : use crate::local_env::{LocalEnv, SafekeeperConf};
      26              : 
      27              : #[derive(Error, Debug)]
      28              : pub enum SafekeeperHttpError {
      29            0 :     #[error("request error: {0}{}", .0.source().map(|e| format!(": {e}")).unwrap_or_default())]
      30              :     Transport(#[from] reqwest::Error),
      31              : 
      32              :     #[error("Error: {0}")]
      33              :     Response(String),
      34              : }
      35              : 
      36              : type Result<T> = result::Result<T, SafekeeperHttpError>;
      37              : 
      38            0 : fn err_from_client_err(err: mgmt_api::Error) -> SafekeeperHttpError {
      39              :     use mgmt_api::Error::*;
      40            0 :     match err {
      41            0 :         ApiError(_, str) => SafekeeperHttpError::Response(str),
      42            0 :         Cancelled => SafekeeperHttpError::Response("Cancelled".to_owned()),
      43            0 :         ReceiveBody(err) => SafekeeperHttpError::Transport(err),
      44            0 :         ReceiveErrorBody(err) => SafekeeperHttpError::Response(err),
      45            0 :         Timeout(str) => SafekeeperHttpError::Response(format!("timeout: {str}")),
      46              :     }
      47            0 : }
      48              : 
      49              : //
      50              : // Control routines for safekeeper.
      51              : //
      52              : // Used in CLI and tests.
      53              : //
      54              : #[derive(Debug)]
      55              : pub struct SafekeeperNode {
      56              :     pub id: NodeId,
      57              : 
      58              :     pub conf: SafekeeperConf,
      59              : 
      60              :     pub pg_connection_config: PgConnectionConfig,
      61              :     pub env: LocalEnv,
      62              :     pub http_client: mgmt_api::Client,
      63              :     pub listen_addr: String,
      64              : }
      65              : 
      66              : impl SafekeeperNode {
      67            0 :     pub fn from_env(env: &LocalEnv, conf: &SafekeeperConf) -> SafekeeperNode {
      68            0 :         let listen_addr = if let Some(ref listen_addr) = conf.listen_addr {
      69            0 :             listen_addr.clone()
      70              :         } else {
      71            0 :             "127.0.0.1".to_string()
      72              :         };
      73            0 :         let jwt = None;
      74            0 :         let http_base_url = format!("http://{}:{}", listen_addr, conf.http_port);
      75            0 :         SafekeeperNode {
      76            0 :             id: conf.id,
      77            0 :             conf: conf.clone(),
      78            0 :             pg_connection_config: Self::safekeeper_connection_config(&listen_addr, conf.pg_port),
      79            0 :             env: env.clone(),
      80            0 :             http_client: mgmt_api::Client::new(env.create_http_client(), http_base_url, jwt),
      81            0 :             listen_addr,
      82            0 :         }
      83            0 :     }
      84              : 
      85              :     /// Construct libpq connection string for connecting to this safekeeper.
      86            0 :     fn safekeeper_connection_config(addr: &str, port: u16) -> PgConnectionConfig {
      87            0 :         PgConnectionConfig::new_host_port(url::Host::parse(addr).unwrap(), port)
      88            0 :     }
      89              : 
      90            0 :     pub fn datadir_path_by_id(env: &LocalEnv, sk_id: NodeId) -> PathBuf {
      91            0 :         env.safekeeper_data_dir(&format!("sk{sk_id}"))
      92            0 :     }
      93              : 
      94            0 :     pub fn datadir_path(&self) -> PathBuf {
      95            0 :         SafekeeperNode::datadir_path_by_id(&self.env, self.id)
      96            0 :     }
      97              : 
      98            0 :     pub fn pid_file(&self) -> Utf8PathBuf {
      99            0 :         Utf8PathBuf::from_path_buf(self.datadir_path().join("safekeeper.pid"))
     100            0 :             .expect("non-Unicode path")
     101            0 :     }
     102              : 
     103              :     /// Initializes a safekeeper node by creating all necessary files,
     104              :     /// e.g. SSL certificates and JWT token file.
     105            0 :     pub fn initialize(&self) -> anyhow::Result<()> {
     106            0 :         if self.env.generate_local_ssl_certs {
     107            0 :             self.env.generate_ssl_cert(
     108            0 :                 &self.datadir_path().join("server.crt"),
     109            0 :                 &self.datadir_path().join("server.key"),
     110            0 :             )?;
     111            0 :         }
     112              : 
     113              :         // Generate a token file for authentication with other safekeepers
     114            0 :         if self.conf.auth_type != AuthType::Trust {
     115            0 :             let token = self
     116            0 :                 .env
     117            0 :                 .generate_auth_token(&Claims::new(None, Scope::SafekeeperData))?;
     118              : 
     119            0 :             let token_path = self.datadir_path().join("peer_jwt_token");
     120            0 :             std::fs::write(token_path, token)?;
     121            0 :         }
     122              : 
     123            0 :         Ok(())
     124            0 :     }
     125              : 
     126            0 :     pub async fn start(
     127            0 :         &self,
     128            0 :         extra_opts: &[String],
     129            0 :         retry_timeout: &Duration,
     130            0 :     ) -> anyhow::Result<()> {
     131            0 :         println!(
     132            0 :             "Starting safekeeper at '{}' in '{}', retrying for {:?}",
     133            0 :             self.pg_connection_config.raw_address(),
     134            0 :             self.datadir_path().display(),
     135              :             retry_timeout,
     136              :         );
     137            0 :         io::stdout().flush().unwrap();
     138              : 
     139            0 :         let listen_pg = format!("{}:{}", self.listen_addr, self.conf.pg_port);
     140            0 :         let listen_http = format!("{}:{}", self.listen_addr, self.conf.http_port);
     141            0 :         let id = self.id;
     142            0 :         let datadir = self.datadir_path();
     143              : 
     144            0 :         let id_string = id.to_string();
     145              :         // TODO: add availability_zone to the config.
     146              :         // Right now we just specify any value here and use it to check metrics in tests.
     147            0 :         let availability_zone = format!("sk-{id_string}");
     148              : 
     149            0 :         let mut args = vec![
     150            0 :             "-D".to_owned(),
     151            0 :             datadir
     152            0 :                 .to_str()
     153            0 :                 .with_context(|| {
     154            0 :                     format!("Datadir path {datadir:?} cannot be represented as a unicode string")
     155            0 :                 })?
     156            0 :                 .to_owned(),
     157            0 :             "--id".to_owned(),
     158            0 :             id_string,
     159            0 :             "--listen-pg".to_owned(),
     160            0 :             listen_pg.clone(),
     161            0 :             "--listen-http".to_owned(),
     162            0 :             listen_http,
     163            0 :             "--availability-zone".to_owned(),
     164            0 :             availability_zone,
     165              :         ];
     166            0 :         if let Some(pg_tenant_only_port) = self.conf.pg_tenant_only_port {
     167            0 :             let listen_pg_tenant_only = format!("{}:{}", self.listen_addr, pg_tenant_only_port);
     168            0 :             args.extend(["--listen-pg-tenant-only".to_owned(), listen_pg_tenant_only]);
     169            0 :         }
     170            0 :         if !self.conf.sync {
     171            0 :             args.push("--no-sync".to_owned());
     172            0 :         }
     173              : 
     174            0 :         let broker_endpoint = format!("{}", self.env.broker.client_url());
     175            0 :         args.extend(["--broker-endpoint".to_owned(), broker_endpoint]);
     176              : 
     177            0 :         let mut backup_threads = String::new();
     178            0 :         if let Some(threads) = self.conf.backup_threads {
     179            0 :             backup_threads = threads.to_string();
     180            0 :             args.extend(["--backup-threads".to_owned(), backup_threads]);
     181            0 :         } else {
     182            0 :             drop(backup_threads);
     183            0 :         }
     184              : 
     185            0 :         if let Some(ref remote_storage) = self.conf.remote_storage {
     186            0 :             args.extend(["--remote-storage".to_owned(), remote_storage.clone()]);
     187            0 :         }
     188              : 
     189            0 :         let key_path = self.env.base_data_dir.join("auth_public_key.pem");
     190            0 :         if self.conf.auth_type != AuthType::Trust {
     191            0 :             args.extend([
     192            0 :                 "--token-auth-type".to_owned(),
     193            0 :                 self.conf.auth_type.to_string(),
     194            0 :             ]);
     195            0 :             let key_path_string = key_path
     196            0 :                 .to_str()
     197            0 :                 .with_context(|| {
     198            0 :                     format!("Key path {key_path:?} cannot be represented as a unicode string")
     199            0 :                 })?
     200            0 :                 .to_owned();
     201            0 :             args.extend([
     202            0 :                 "--pg-auth-public-key-path".to_owned(),
     203            0 :                 key_path_string.clone(),
     204            0 :             ]);
     205            0 :             args.extend([
     206            0 :                 "--pg-tenant-only-auth-public-key-path".to_owned(),
     207            0 :                 key_path_string.clone(),
     208            0 :             ]);
     209            0 :             args.extend([
     210            0 :                 "--http-auth-public-key-path".to_owned(),
     211            0 :                 key_path_string.clone(),
     212            0 :             ]);
     213              : 
     214            0 :             let token_path = self.datadir_path().join("peer_jwt_token");
     215            0 :             let token_path_str = token_path
     216            0 :                 .to_str()
     217            0 :                 .with_context(|| {
     218            0 :                     format!("Token path {token_path:?} cannot be represented as a unicode string")
     219            0 :                 })?
     220            0 :                 .to_owned();
     221            0 :             args.extend(["--auth-token-path".to_owned(), token_path_str]);
     222            0 :         }
     223              : 
     224            0 :         if let Some(https_port) = self.conf.https_port {
     225            0 :             args.extend([
     226            0 :                 "--listen-https".to_owned(),
     227            0 :                 format!("{}:{}", self.listen_addr, https_port),
     228            0 :             ]);
     229            0 :         }
     230            0 :         if let Some(ssl_ca_file) = self.env.ssl_ca_cert_path() {
     231            0 :             args.push(format!("--ssl-ca-file={}", ssl_ca_file.to_str().unwrap()));
     232            0 :         }
     233              : 
     234            0 :         args.extend_from_slice(extra_opts);
     235              : 
     236            0 :         background_process::start_process(
     237            0 :             &format!("safekeeper-{id}"),
     238            0 :             &datadir,
     239            0 :             &self.env.safekeeper_bin(),
     240            0 :             &args,
     241            0 :             self.safekeeper_env_variables()?,
     242            0 :             background_process::InitialPidFile::Expect(self.pid_file()),
     243            0 :             retry_timeout,
     244            0 :             || async {
     245            0 :                 match self.check_status().await {
     246            0 :                     Ok(()) => Ok(true),
     247            0 :                     Err(SafekeeperHttpError::Transport(_)) => Ok(false),
     248            0 :                     Err(e) => Err(anyhow::anyhow!("Failed to check node status: {e}")),
     249              :                 }
     250            0 :             },
     251              :         )
     252            0 :         .await
     253            0 :     }
     254              : 
     255            0 :     fn safekeeper_env_variables(&self) -> anyhow::Result<Vec<(String, String)>> {
     256              :         // TODO: remove me
     257            0 :         Ok(vec![])
     258            0 :     }
     259              : 
     260              :     ///
     261              :     /// Stop the server.
     262              :     ///
     263              :     /// If 'immediate' is true, we use SIGQUIT, killing the process immediately.
     264              :     /// Otherwise we use SIGTERM, triggering a clean shutdown
     265              :     ///
     266              :     /// If the server is not running, returns success
     267              :     ///
     268            0 :     pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
     269            0 :         background_process::stop_process(
     270            0 :             immediate,
     271            0 :             &format!("safekeeper {}", self.id),
     272            0 :             &self.pid_file(),
     273              :         )
     274            0 :     }
     275              : 
     276            0 :     pub async fn check_status(&self) -> Result<()> {
     277            0 :         self.http_client
     278            0 :             .status()
     279            0 :             .await
     280            0 :             .map_err(err_from_client_err)?;
     281            0 :         Ok(())
     282            0 :     }
     283              : 
     284            0 :     pub async fn create_timeline(&self, req: &TimelineCreateRequest) -> Result<()> {
     285            0 :         self.http_client
     286            0 :             .create_timeline(req)
     287            0 :             .await
     288            0 :             .map_err(err_from_client_err)?;
     289            0 :         Ok(())
     290            0 :     }
     291              : }
        

Generated by: LCOV version 2.1-beta