LCOV - code coverage report
Current view: top level - control_plane/src - safekeeper.rs (source / functions) Coverage Total Hit
Test: c639aa5f7ab62b43d647b10f40d15a15686ce8a9.info Lines: 89.5 % 152 136
Test Date: 2024-02-12 20:26:03 Functions: 69.6 % 23 16

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

Generated by: LCOV version 2.1-beta