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 : }
|
