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