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