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