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