LCOV - code coverage report
Current view: top level - control_plane/src - safekeeper.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 170 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 20 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::future::Future;
      10              : use std::io::Write;
      11              : use std::path::PathBuf;
      12              : use std::time::Duration;
      13              : use std::{io, result};
      14              : 
      15              : use anyhow::Context;
      16              : use camino::Utf8PathBuf;
      17              : use postgres_connection::PgConnectionConfig;
      18              : use reqwest::{IntoUrl, Method};
      19              : use thiserror::Error;
      20              : 
      21              : use http_utils::error::HttpErrorBody;
      22              : use utils::auth::{Claims, Scope};
      23              : use utils::id::NodeId;
      24              : 
      25              : use crate::{
      26              :     background_process,
      27              :     local_env::{LocalEnv, SafekeeperConf},
      28              : };
      29              : 
      30              : #[derive(Error, Debug)]
      31              : pub enum SafekeeperHttpError {
      32            0 :     #[error("request error: {0}{}", .0.source().map(|e| format!(": {e}")).unwrap_or_default())]
      33              :     Transport(#[from] reqwest::Error),
      34              : 
      35              :     #[error("Error: {0}")]
      36              :     Response(String),
      37              : }
      38              : 
      39              : type Result<T> = result::Result<T, SafekeeperHttpError>;
      40              : 
      41              : pub(crate) trait ResponseErrorMessageExt: Sized {
      42              :     fn error_from_body(self) -> impl Future<Output = Result<Self>> + Send;
      43              : }
      44              : 
      45              : impl ResponseErrorMessageExt for reqwest::Response {
      46            0 :     async fn error_from_body(self) -> Result<Self> {
      47            0 :         let status = self.status();
      48            0 :         if !(status.is_client_error() || status.is_server_error()) {
      49            0 :             return Ok(self);
      50            0 :         }
      51            0 : 
      52            0 :         // reqwest does not export its error construction utility functions, so let's craft the message ourselves
      53            0 :         let url = self.url().to_owned();
      54            0 :         Err(SafekeeperHttpError::Response(
      55            0 :             match self.json::<HttpErrorBody>().await {
      56            0 :                 Ok(err_body) => format!("Error: {}", err_body.msg),
      57            0 :                 Err(_) => format!("Http error ({}) at {}.", status.as_u16(), url),
      58              :             },
      59              :         ))
      60            0 :     }
      61              : }
      62              : 
      63              : //
      64              : // Control routines for safekeeper.
      65              : //
      66              : // Used in CLI and tests.
      67              : //
      68              : #[derive(Debug)]
      69              : pub struct SafekeeperNode {
      70              :     pub id: NodeId,
      71              : 
      72              :     pub conf: SafekeeperConf,
      73              : 
      74              :     pub pg_connection_config: PgConnectionConfig,
      75              :     pub env: LocalEnv,
      76              :     pub http_client: reqwest::Client,
      77              :     pub listen_addr: String,
      78              :     pub http_base_url: String,
      79              : }
      80              : 
      81              : impl SafekeeperNode {
      82            0 :     pub fn from_env(env: &LocalEnv, conf: &SafekeeperConf) -> SafekeeperNode {
      83            0 :         let listen_addr = if let Some(ref listen_addr) = conf.listen_addr {
      84            0 :             listen_addr.clone()
      85              :         } else {
      86            0 :             "127.0.0.1".to_string()
      87              :         };
      88            0 :         SafekeeperNode {
      89            0 :             id: conf.id,
      90            0 :             conf: conf.clone(),
      91            0 :             pg_connection_config: Self::safekeeper_connection_config(&listen_addr, conf.pg_port),
      92            0 :             env: env.clone(),
      93            0 :             http_client: reqwest::Client::new(),
      94            0 :             http_base_url: format!("http://{}:{}/v1", listen_addr, conf.http_port),
      95            0 :             listen_addr,
      96            0 :         }
      97            0 :     }
      98              : 
      99              :     /// Construct libpq connection string for connecting to this safekeeper.
     100            0 :     fn safekeeper_connection_config(addr: &str, port: u16) -> PgConnectionConfig {
     101            0 :         PgConnectionConfig::new_host_port(url::Host::parse(addr).unwrap(), port)
     102            0 :     }
     103              : 
     104            0 :     pub fn datadir_path_by_id(env: &LocalEnv, sk_id: NodeId) -> PathBuf {
     105            0 :         env.safekeeper_data_dir(&format!("sk{sk_id}"))
     106            0 :     }
     107              : 
     108            0 :     pub fn datadir_path(&self) -> PathBuf {
     109            0 :         SafekeeperNode::datadir_path_by_id(&self.env, self.id)
     110            0 :     }
     111              : 
     112            0 :     pub fn pid_file(&self) -> Utf8PathBuf {
     113            0 :         Utf8PathBuf::from_path_buf(self.datadir_path().join("safekeeper.pid"))
     114            0 :             .expect("non-Unicode path")
     115            0 :     }
     116              : 
     117            0 :     pub async fn start(
     118            0 :         &self,
     119            0 :         extra_opts: &[String],
     120            0 :         retry_timeout: &Duration,
     121            0 :     ) -> anyhow::Result<()> {
     122            0 :         print!(
     123            0 :             "Starting safekeeper at '{}' in '{}', retrying for {:?}",
     124            0 :             self.pg_connection_config.raw_address(),
     125            0 :             self.datadir_path().display(),
     126            0 :             retry_timeout,
     127            0 :         );
     128            0 :         io::stdout().flush().unwrap();
     129            0 : 
     130            0 :         let listen_pg = format!("{}:{}", self.listen_addr, self.conf.pg_port);
     131            0 :         let listen_http = format!("{}:{}", self.listen_addr, self.conf.http_port);
     132            0 :         let id = self.id;
     133            0 :         let datadir = self.datadir_path();
     134            0 : 
     135            0 :         let id_string = id.to_string();
     136            0 :         // TODO: add availability_zone to the config.
     137            0 :         // Right now we just specify any value here and use it to check metrics in tests.
     138            0 :         let availability_zone = format!("sk-{}", id_string);
     139              : 
     140            0 :         let mut args = vec![
     141            0 :             "-D".to_owned(),
     142            0 :             datadir
     143            0 :                 .to_str()
     144            0 :                 .with_context(|| {
     145            0 :                     format!("Datadir path {datadir:?} cannot be represented as a unicode string")
     146            0 :                 })?
     147            0 :                 .to_owned(),
     148            0 :             "--id".to_owned(),
     149            0 :             id_string,
     150            0 :             "--listen-pg".to_owned(),
     151            0 :             listen_pg,
     152            0 :             "--listen-http".to_owned(),
     153            0 :             listen_http,
     154            0 :             "--availability-zone".to_owned(),
     155            0 :             availability_zone,
     156              :         ];
     157            0 :         if let Some(pg_tenant_only_port) = self.conf.pg_tenant_only_port {
     158            0 :             let listen_pg_tenant_only = format!("{}:{}", self.listen_addr, pg_tenant_only_port);
     159            0 :             args.extend(["--listen-pg-tenant-only".to_owned(), listen_pg_tenant_only]);
     160            0 :         }
     161            0 :         if !self.conf.sync {
     162            0 :             args.push("--no-sync".to_owned());
     163            0 :         }
     164              : 
     165            0 :         let broker_endpoint = format!("{}", self.env.broker.client_url());
     166            0 :         args.extend(["--broker-endpoint".to_owned(), broker_endpoint]);
     167            0 : 
     168            0 :         let mut backup_threads = String::new();
     169            0 :         if let Some(threads) = self.conf.backup_threads {
     170            0 :             backup_threads = threads.to_string();
     171            0 :             args.extend(["--backup-threads".to_owned(), backup_threads]);
     172            0 :         } else {
     173            0 :             drop(backup_threads);
     174            0 :         }
     175              : 
     176            0 :         if let Some(ref remote_storage) = self.conf.remote_storage {
     177            0 :             args.extend(["--remote-storage".to_owned(), remote_storage.clone()]);
     178            0 :         }
     179              : 
     180            0 :         let key_path = self.env.base_data_dir.join("auth_public_key.pem");
     181            0 :         if self.conf.auth_enabled {
     182            0 :             let key_path_string = key_path
     183            0 :                 .to_str()
     184            0 :                 .with_context(|| {
     185            0 :                     format!("Key path {key_path:?} cannot be represented as a unicode string")
     186            0 :                 })?
     187            0 :                 .to_owned();
     188            0 :             args.extend([
     189            0 :                 "--pg-auth-public-key-path".to_owned(),
     190            0 :                 key_path_string.clone(),
     191            0 :             ]);
     192            0 :             args.extend([
     193            0 :                 "--pg-tenant-only-auth-public-key-path".to_owned(),
     194            0 :                 key_path_string.clone(),
     195            0 :             ]);
     196            0 :             args.extend([
     197            0 :                 "--http-auth-public-key-path".to_owned(),
     198            0 :                 key_path_string.clone(),
     199            0 :             ]);
     200            0 :         }
     201              : 
     202            0 :         args.extend_from_slice(extra_opts);
     203            0 : 
     204            0 :         background_process::start_process(
     205            0 :             &format!("safekeeper-{id}"),
     206            0 :             &datadir,
     207            0 :             &self.env.safekeeper_bin(),
     208            0 :             &args,
     209            0 :             self.safekeeper_env_variables()?,
     210            0 :             background_process::InitialPidFile::Expect(self.pid_file()),
     211            0 :             retry_timeout,
     212            0 :             || async {
     213            0 :                 match self.check_status().await {
     214            0 :                     Ok(()) => Ok(true),
     215            0 :                     Err(SafekeeperHttpError::Transport(_)) => Ok(false),
     216            0 :                     Err(e) => Err(anyhow::anyhow!("Failed to check node status: {e}")),
     217              :                 }
     218            0 :             },
     219            0 :         )
     220            0 :         .await
     221            0 :     }
     222              : 
     223            0 :     fn safekeeper_env_variables(&self) -> anyhow::Result<Vec<(String, String)>> {
     224            0 :         // Generate a token to connect from safekeeper to peers
     225            0 :         if self.conf.auth_enabled {
     226            0 :             let token = self
     227            0 :                 .env
     228            0 :                 .generate_auth_token(&Claims::new(None, Scope::SafekeeperData))?;
     229            0 :             Ok(vec![("SAFEKEEPER_AUTH_TOKEN".to_owned(), token)])
     230              :         } else {
     231            0 :             Ok(Vec::new())
     232              :         }
     233            0 :     }
     234              : 
     235              :     ///
     236              :     /// Stop the server.
     237              :     ///
     238              :     /// If 'immediate' is true, we use SIGQUIT, killing the process immediately.
     239              :     /// Otherwise we use SIGTERM, triggering a clean shutdown
     240              :     ///
     241              :     /// If the server is not running, returns success
     242              :     ///
     243            0 :     pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
     244            0 :         background_process::stop_process(
     245            0 :             immediate,
     246            0 :             &format!("safekeeper {}", self.id),
     247            0 :             &self.pid_file(),
     248            0 :         )
     249            0 :     }
     250              : 
     251            0 :     fn http_request<U: IntoUrl>(&self, method: Method, url: U) -> reqwest::RequestBuilder {
     252            0 :         // TODO: authentication
     253            0 :         //if self.env.auth_type == AuthType::NeonJWT {
     254            0 :         //    builder = builder.bearer_auth(&self.env.safekeeper_auth_token)
     255            0 :         //}
     256            0 :         self.http_client.request(method, url)
     257            0 :     }
     258              : 
     259            0 :     pub async fn check_status(&self) -> Result<()> {
     260            0 :         self.http_request(Method::GET, format!("{}/{}", self.http_base_url, "status"))
     261            0 :             .send()
     262            0 :             .await?
     263            0 :             .error_from_body()
     264            0 :             .await?;
     265            0 :         Ok(())
     266            0 :     }
     267              : }
        

Generated by: LCOV version 2.1-beta