TLA Line data Source code
1 : //! Code to manage pageservers
2 : //!
3 : //! In the local test environment, the pageserver stores its data directly in
4 : //!
5 : //! .neon/
6 : //!
7 : use std::borrow::Cow;
8 : use std::collections::HashMap;
9 :
10 : use std::io;
11 : use std::io::Write;
12 : use std::num::NonZeroU64;
13 : use std::path::PathBuf;
14 : use std::process::{Child, Command};
15 : use std::time::Duration;
16 :
17 : use anyhow::{bail, Context};
18 : use camino::Utf8PathBuf;
19 : use futures::SinkExt;
20 : use pageserver_api::models::{self, LocationConfig, TenantInfo, TimelineInfo};
21 : use pageserver_api::shard::TenantShardId;
22 : use pageserver_client::mgmt_api;
23 : use postgres_backend::AuthType;
24 : use postgres_connection::{parse_host_port, PgConnectionConfig};
25 : use utils::auth::{Claims, Scope};
26 : use utils::{
27 : id::{TenantId, TimelineId},
28 : lsn::Lsn,
29 : };
30 :
31 : use crate::local_env::PageServerConf;
32 : use crate::{background_process, local_env::LocalEnv};
33 :
34 : /// Directory within .neon which will be used by default for LocalFs remote storage.
35 : pub const PAGESERVER_REMOTE_STORAGE_DIR: &str = "local_fs_remote_storage/pageserver";
36 :
37 : //
38 : // Control routines for pageserver.
39 : //
40 : // Used in CLI and tests.
41 : //
42 UBC 0 : #[derive(Debug)]
43 : pub struct PageServerNode {
44 : pub pg_connection_config: PgConnectionConfig,
45 : pub conf: PageServerConf,
46 : pub env: LocalEnv,
47 : pub http_client: mgmt_api::Client,
48 : }
49 :
50 : impl PageServerNode {
51 CBC 6612 : pub fn from_env(env: &LocalEnv, conf: &PageServerConf) -> PageServerNode {
52 6612 : let (host, port) =
53 6612 : parse_host_port(&conf.listen_pg_addr).expect("Unable to parse listen_pg_addr");
54 6612 : let port = port.unwrap_or(5432);
55 6612 : Self {
56 6612 : pg_connection_config: PgConnectionConfig::new_host_port(host, port),
57 6612 : conf: conf.clone(),
58 6612 : env: env.clone(),
59 6612 : http_client: mgmt_api::Client::new(
60 6612 : format!("http://{}", conf.listen_http_addr),
61 6612 : {
62 6612 : match conf.http_auth_type {
63 6486 : AuthType::Trust => None,
64 126 : AuthType::NeonJWT => Some(
65 126 : env.generate_auth_token(&Claims::new(None, Scope::PageServerApi))
66 126 : .unwrap(),
67 126 : ),
68 : }
69 : }
70 6612 : .as_deref(),
71 6612 : ),
72 6612 : }
73 6612 : }
74 :
75 : /// Merge overrides provided by the user on the command line with our default overides derived from neon_local configuration.
76 : ///
77 : /// These all end up on the command line of the `pageserver` binary.
78 906 : fn neon_local_overrides(&self, cli_overrides: &[&str]) -> Vec<String> {
79 906 : let id = format!("id={}", self.conf.id);
80 906 : // FIXME: the paths should be shell-escaped to handle paths with spaces, quotas etc.
81 906 : let pg_distrib_dir_param = format!(
82 906 : "pg_distrib_dir='{}'",
83 906 : self.env.pg_distrib_dir_raw().display()
84 906 : );
85 906 :
86 906 : let http_auth_type_param = format!("http_auth_type='{}'", self.conf.http_auth_type);
87 906 : let listen_http_addr_param = format!("listen_http_addr='{}'", self.conf.listen_http_addr);
88 906 :
89 906 : let pg_auth_type_param = format!("pg_auth_type='{}'", self.conf.pg_auth_type);
90 906 : let listen_pg_addr_param = format!("listen_pg_addr='{}'", self.conf.listen_pg_addr);
91 906 :
92 906 : let broker_endpoint_param = format!("broker_endpoint='{}'", self.env.broker.client_url());
93 906 :
94 906 : let mut overrides = vec![
95 906 : id,
96 906 : pg_distrib_dir_param,
97 906 : http_auth_type_param,
98 906 : pg_auth_type_param,
99 906 : listen_http_addr_param,
100 906 : listen_pg_addr_param,
101 906 : broker_endpoint_param,
102 906 : ];
103 :
104 906 : if let Some(control_plane_api) = &self.env.control_plane_api {
105 906 : overrides.push(format!(
106 906 : "control_plane_api='{}'",
107 906 : control_plane_api.as_str()
108 906 : ));
109 906 : }
110 :
111 906 : if !cli_overrides
112 906 : .iter()
113 911 : .any(|c| c.starts_with("remote_storage"))
114 4 : {
115 4 : overrides.push(format!(
116 4 : "remote_storage={{local_path='../{PAGESERVER_REMOTE_STORAGE_DIR}'}}"
117 4 : ));
118 902 : }
119 :
120 906 : if self.conf.http_auth_type != AuthType::Trust || self.conf.pg_auth_type != AuthType::Trust
121 22 : {
122 22 : // Keys are generated in the toplevel repo dir, pageservers' workdirs
123 22 : // are one level below that, so refer to keys with ../
124 22 : overrides.push("auth_validation_public_key_path='../auth_public_key.pem'".to_owned());
125 884 : }
126 :
127 : // Apply the user-provided overrides
128 1066 : overrides.extend(cli_overrides.iter().map(|&c| c.to_owned()));
129 906 :
130 906 : overrides
131 906 : }
132 :
133 : /// Initializes a pageserver node by creating its config with the overrides provided.
134 349 : pub fn initialize(&self, config_overrides: &[&str]) -> anyhow::Result<()> {
135 349 : // First, run `pageserver --init` and wait for it to write a config into FS and exit.
136 349 : self.pageserver_init(config_overrides)
137 349 : .with_context(|| format!("Failed to run init for pageserver node {}", self.conf.id))
138 349 : }
139 :
140 2020 : pub fn repo_path(&self) -> PathBuf {
141 2020 : self.env.pageserver_data_dir(self.conf.id)
142 2020 : }
143 :
144 : /// The pid file is created by the pageserver process, with its pid stored inside.
145 : /// Other pageservers cannot lock the same file and overwrite it for as long as the current
146 : /// pageserver runs. (Unless someone removes the file manually; never do that!)
147 1114 : fn pid_file(&self) -> Utf8PathBuf {
148 1114 : Utf8PathBuf::from_path_buf(self.repo_path().join("pageserver.pid"))
149 1114 : .expect("non-Unicode path")
150 1114 : }
151 :
152 557 : pub async fn start(&self, config_overrides: &[&str]) -> anyhow::Result<Child> {
153 3356 : self.start_node(config_overrides, false).await
154 557 : }
155 :
156 349 : fn pageserver_init(&self, config_overrides: &[&str]) -> anyhow::Result<()> {
157 349 : let datadir = self.repo_path();
158 349 : let node_id = self.conf.id;
159 349 : println!(
160 349 : "Initializing pageserver node {} at '{}' in {:?}",
161 349 : node_id,
162 349 : self.pg_connection_config.raw_address(),
163 349 : datadir
164 349 : );
165 349 : io::stdout().flush()?;
166 :
167 349 : if !datadir.exists() {
168 349 : std::fs::create_dir(&datadir)?;
169 UBC 0 : }
170 :
171 CBC 349 : let datadir_path_str = datadir.to_str().with_context(|| {
172 UBC 0 : format!("Cannot start pageserver node {node_id} in path that has no string representation: {datadir:?}")
173 CBC 349 : })?;
174 349 : let mut args = self.pageserver_basic_args(config_overrides, datadir_path_str);
175 349 : args.push(Cow::Borrowed("--init"));
176 :
177 349 : let init_output = Command::new(self.env.pageserver_bin())
178 349 : .args(args.iter().map(Cow::as_ref))
179 349 : .envs(self.pageserver_env_variables()?)
180 349 : .output()
181 349 : .with_context(|| format!("Failed to run pageserver init for node {node_id}"))?;
182 :
183 : anyhow::ensure!(
184 349 : init_output.status.success(),
185 UBC 0 : "Pageserver init for node {} did not finish successfully, stdout: {}, stderr: {}",
186 0 : node_id,
187 0 : String::from_utf8_lossy(&init_output.stdout),
188 0 : String::from_utf8_lossy(&init_output.stderr),
189 : );
190 :
191 CBC 349 : Ok(())
192 349 : }
193 :
194 557 : async fn start_node(
195 557 : &self,
196 557 : config_overrides: &[&str],
197 557 : update_config: bool,
198 557 : ) -> anyhow::Result<Child> {
199 557 : // TODO: using a thread here because start_process() is not async but we need to call check_status()
200 557 : let datadir = self.repo_path();
201 557 : print!(
202 557 : "Starting pageserver node {} at '{}' in {:?}",
203 557 : self.conf.id,
204 557 : self.pg_connection_config.raw_address(),
205 557 : datadir
206 557 : );
207 557 : io::stdout().flush().context("flush stdout")?;
208 :
209 557 : let datadir_path_str = datadir.to_str().with_context(|| {
210 UBC 0 : format!(
211 0 : "Cannot start pageserver node {} in path that has no string representation: {:?}",
212 0 : self.conf.id, datadir,
213 0 : )
214 CBC 557 : })?;
215 557 : let mut args = self.pageserver_basic_args(config_overrides, datadir_path_str);
216 557 : if update_config {
217 UBC 0 : args.push(Cow::Borrowed("--update-config"));
218 CBC 557 : }
219 : background_process::start_process(
220 557 : "pageserver",
221 557 : &datadir,
222 557 : &self.env.pageserver_bin(),
223 557 : args.iter().map(Cow::as_ref),
224 557 : self.pageserver_env_variables()?,
225 557 : background_process::InitialPidFile::Expect(self.pid_file()),
226 1121 : || async {
227 3356 : let st = self.check_status().await;
228 564 : match st {
229 557 : Ok(()) => Ok(true),
230 564 : Err(mgmt_api::Error::ReceiveBody(_)) => Ok(false),
231 UBC 0 : Err(e) => Err(anyhow::anyhow!("Failed to check node status: {e}")),
232 : }
233 CBC 1121 : },
234 : )
235 3356 : .await
236 557 : }
237 :
238 906 : fn pageserver_basic_args<'a>(
239 906 : &self,
240 906 : config_overrides: &'a [&'a str],
241 906 : datadir_path_str: &'a str,
242 906 : ) -> Vec<Cow<'a, str>> {
243 906 : let mut args = vec![Cow::Borrowed("-D"), Cow::Borrowed(datadir_path_str)];
244 906 :
245 906 : let overrides = self.neon_local_overrides(config_overrides);
246 9246 : for config_override in overrides {
247 8340 : args.push(Cow::Borrowed("-c"));
248 8340 : args.push(Cow::Owned(config_override));
249 8340 : }
250 :
251 906 : args
252 906 : }
253 :
254 906 : fn pageserver_env_variables(&self) -> anyhow::Result<Vec<(String, String)>> {
255 906 : // FIXME: why is this tied to pageserver's auth type? Whether or not the safekeeper
256 906 : // needs a token, and how to generate that token, seems independent to whether
257 906 : // the pageserver requires a token in incoming requests.
258 906 : Ok(if self.conf.http_auth_type != AuthType::Trust {
259 : // Generate a token to connect from the pageserver to a safekeeper
260 22 : let token = self
261 22 : .env
262 22 : .generate_auth_token(&Claims::new(None, Scope::SafekeeperData))?;
263 22 : vec![("NEON_AUTH_TOKEN".to_owned(), token)]
264 : } else {
265 884 : Vec::new()
266 : })
267 906 : }
268 :
269 : ///
270 : /// Stop the server.
271 : ///
272 : /// If 'immediate' is true, we use SIGQUIT, killing the process immediately.
273 : /// Otherwise we use SIGTERM, triggering a clean shutdown
274 : ///
275 : /// If the server is not running, returns success
276 : ///
277 557 : pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
278 557 : background_process::stop_process(immediate, "pageserver", &self.pid_file())
279 557 : }
280 :
281 5 : pub async fn page_server_psql_client(
282 5 : &self,
283 5 : ) -> anyhow::Result<(
284 5 : tokio_postgres::Client,
285 5 : tokio_postgres::Connection<tokio_postgres::Socket, tokio_postgres::tls::NoTlsStream>,
286 5 : )> {
287 5 : let mut config = self.pg_connection_config.clone();
288 5 : if self.conf.pg_auth_type == AuthType::NeonJWT {
289 UBC 0 : let token = self
290 0 : .env
291 0 : .generate_auth_token(&Claims::new(None, Scope::PageServerApi))?;
292 0 : config = config.set_password(Some(token));
293 CBC 5 : }
294 15 : Ok(config.connect_no_tls().await?)
295 5 : }
296 :
297 1121 : pub async fn check_status(&self) -> mgmt_api::Result<()> {
298 3356 : self.http_client.status().await
299 1121 : }
300 :
301 6 : pub async fn tenant_list(&self) -> mgmt_api::Result<Vec<TenantInfo>> {
302 24 : self.http_client.list_tenants().await
303 6 : }
304 :
305 409 : pub async fn tenant_create(
306 409 : &self,
307 409 : new_tenant_id: TenantId,
308 409 : generation: Option<u32>,
309 409 : settings: HashMap<&str, &str>,
310 409 : ) -> anyhow::Result<TenantId> {
311 409 : let mut settings = settings.clone();
312 :
313 409 : let config = models::TenantConfig {
314 409 : checkpoint_distance: settings
315 409 : .remove("checkpoint_distance")
316 409 : .map(|x| x.parse::<u64>())
317 409 : .transpose()?,
318 409 : checkpoint_timeout: settings.remove("checkpoint_timeout").map(|x| x.to_string()),
319 409 : compaction_target_size: settings
320 409 : .remove("compaction_target_size")
321 409 : .map(|x| x.parse::<u64>())
322 409 : .transpose()?,
323 409 : compaction_period: settings.remove("compaction_period").map(|x| x.to_string()),
324 409 : compaction_threshold: settings
325 409 : .remove("compaction_threshold")
326 409 : .map(|x| x.parse::<usize>())
327 409 : .transpose()?,
328 409 : gc_horizon: settings
329 409 : .remove("gc_horizon")
330 409 : .map(|x| x.parse::<u64>())
331 409 : .transpose()?,
332 409 : gc_period: settings.remove("gc_period").map(|x| x.to_string()),
333 409 : image_creation_threshold: settings
334 409 : .remove("image_creation_threshold")
335 409 : .map(|x| x.parse::<usize>())
336 409 : .transpose()?,
337 409 : pitr_interval: settings.remove("pitr_interval").map(|x| x.to_string()),
338 409 : walreceiver_connect_timeout: settings
339 409 : .remove("walreceiver_connect_timeout")
340 409 : .map(|x| x.to_string()),
341 409 : lagging_wal_timeout: settings
342 409 : .remove("lagging_wal_timeout")
343 409 : .map(|x| x.to_string()),
344 409 : max_lsn_wal_lag: settings
345 409 : .remove("max_lsn_wal_lag")
346 409 : .map(|x| x.parse::<NonZeroU64>())
347 409 : .transpose()
348 409 : .context("Failed to parse 'max_lsn_wal_lag' as non zero integer")?,
349 409 : trace_read_requests: settings
350 409 : .remove("trace_read_requests")
351 409 : .map(|x| x.parse::<bool>())
352 409 : .transpose()
353 409 : .context("Failed to parse 'trace_read_requests' as bool")?,
354 409 : eviction_policy: settings
355 409 : .remove("eviction_policy")
356 409 : .map(serde_json::from_str)
357 409 : .transpose()
358 409 : .context("Failed to parse 'eviction_policy' json")?,
359 409 : min_resident_size_override: settings
360 409 : .remove("min_resident_size_override")
361 409 : .map(|x| x.parse::<u64>())
362 409 : .transpose()
363 409 : .context("Failed to parse 'min_resident_size_override' as integer")?,
364 409 : evictions_low_residence_duration_metric_threshold: settings
365 409 : .remove("evictions_low_residence_duration_metric_threshold")
366 409 : .map(|x| x.to_string()),
367 409 : gc_feedback: settings
368 409 : .remove("gc_feedback")
369 409 : .map(|x| x.parse::<bool>())
370 409 : .transpose()
371 409 : .context("Failed to parse 'gc_feedback' as bool")?,
372 409 : heatmap_period: settings.remove("heatmap_period").map(|x| x.to_string()),
373 409 : };
374 409 :
375 409 : let request = models::TenantCreateRequest {
376 409 : new_tenant_id: TenantShardId::unsharded(new_tenant_id),
377 409 : generation,
378 409 : config,
379 409 : };
380 409 : if !settings.is_empty() {
381 1 : bail!("Unrecognized tenant settings: {settings:?}")
382 408 : }
383 1632 : Ok(self.http_client.tenant_create(&request).await?)
384 409 : }
385 :
386 14 : pub async fn tenant_config(
387 14 : &self,
388 14 : tenant_id: TenantId,
389 14 : mut settings: HashMap<&str, &str>,
390 14 : ) -> anyhow::Result<()> {
391 14 : let config = {
392 : // Braces to make the diff easier to read
393 : models::TenantConfig {
394 14 : checkpoint_distance: settings
395 14 : .remove("checkpoint_distance")
396 14 : .map(|x| x.parse::<u64>())
397 14 : .transpose()
398 14 : .context("Failed to parse 'checkpoint_distance' as an integer")?,
399 14 : checkpoint_timeout: settings.remove("checkpoint_timeout").map(|x| x.to_string()),
400 14 : compaction_target_size: settings
401 14 : .remove("compaction_target_size")
402 14 : .map(|x| x.parse::<u64>())
403 14 : .transpose()
404 14 : .context("Failed to parse 'compaction_target_size' as an integer")?,
405 14 : compaction_period: settings.remove("compaction_period").map(|x| x.to_string()),
406 14 : compaction_threshold: settings
407 14 : .remove("compaction_threshold")
408 14 : .map(|x| x.parse::<usize>())
409 14 : .transpose()
410 14 : .context("Failed to parse 'compaction_threshold' as an integer")?,
411 14 : gc_horizon: settings
412 14 : .remove("gc_horizon")
413 14 : .map(|x| x.parse::<u64>())
414 14 : .transpose()
415 14 : .context("Failed to parse 'gc_horizon' as an integer")?,
416 14 : gc_period: settings.remove("gc_period").map(|x| x.to_string()),
417 14 : image_creation_threshold: settings
418 14 : .remove("image_creation_threshold")
419 14 : .map(|x| x.parse::<usize>())
420 14 : .transpose()
421 14 : .context("Failed to parse 'image_creation_threshold' as non zero integer")?,
422 14 : pitr_interval: settings.remove("pitr_interval").map(|x| x.to_string()),
423 14 : walreceiver_connect_timeout: settings
424 14 : .remove("walreceiver_connect_timeout")
425 14 : .map(|x| x.to_string()),
426 14 : lagging_wal_timeout: settings
427 14 : .remove("lagging_wal_timeout")
428 14 : .map(|x| x.to_string()),
429 14 : max_lsn_wal_lag: settings
430 14 : .remove("max_lsn_wal_lag")
431 14 : .map(|x| x.parse::<NonZeroU64>())
432 14 : .transpose()
433 14 : .context("Failed to parse 'max_lsn_wal_lag' as non zero integer")?,
434 14 : trace_read_requests: settings
435 14 : .remove("trace_read_requests")
436 14 : .map(|x| x.parse::<bool>())
437 14 : .transpose()
438 14 : .context("Failed to parse 'trace_read_requests' as bool")?,
439 14 : eviction_policy: settings
440 14 : .remove("eviction_policy")
441 14 : .map(serde_json::from_str)
442 14 : .transpose()
443 14 : .context("Failed to parse 'eviction_policy' json")?,
444 14 : min_resident_size_override: settings
445 14 : .remove("min_resident_size_override")
446 14 : .map(|x| x.parse::<u64>())
447 14 : .transpose()
448 14 : .context("Failed to parse 'min_resident_size_override' as an integer")?,
449 14 : evictions_low_residence_duration_metric_threshold: settings
450 14 : .remove("evictions_low_residence_duration_metric_threshold")
451 14 : .map(|x| x.to_string()),
452 14 : gc_feedback: settings
453 14 : .remove("gc_feedback")
454 14 : .map(|x| x.parse::<bool>())
455 14 : .transpose()
456 14 : .context("Failed to parse 'gc_feedback' as bool")?,
457 14 : heatmap_period: settings.remove("heatmap_period").map(|x| x.to_string()),
458 14 : }
459 14 : };
460 14 :
461 14 : if !settings.is_empty() {
462 UBC 0 : bail!("Unrecognized tenant settings: {settings:?}")
463 CBC 14 : }
464 14 :
465 14 : self.http_client
466 14 : .tenant_config(&models::TenantConfigRequest { tenant_id, config })
467 56 : .await?;
468 :
469 14 : Ok(())
470 14 : }
471 :
472 UBC 0 : pub async fn location_config(
473 0 : &self,
474 0 : tenant_id: TenantId,
475 0 : config: LocationConfig,
476 0 : flush_ms: Option<Duration>,
477 0 : ) -> anyhow::Result<()> {
478 0 : Ok(self
479 0 : .http_client
480 0 : .location_config(tenant_id, config, flush_ms)
481 0 : .await?)
482 0 : }
483 :
484 CBC 19 : pub async fn timeline_list(&self, tenant_id: &TenantId) -> anyhow::Result<Vec<TimelineInfo>> {
485 76 : Ok(self.http_client.list_timelines(*tenant_id).await?)
486 19 : }
487 :
488 UBC 0 : pub async fn tenant_secondary_download(&self, tenant_id: &TenantShardId) -> anyhow::Result<()> {
489 0 : Ok(self
490 0 : .http_client
491 0 : .tenant_secondary_download(*tenant_id)
492 0 : .await?)
493 0 : }
494 :
495 CBC 748 : pub async fn timeline_create(
496 748 : &self,
497 748 : tenant_id: TenantId,
498 748 : new_timeline_id: Option<TimelineId>,
499 748 : ancestor_start_lsn: Option<Lsn>,
500 748 : ancestor_timeline_id: Option<TimelineId>,
501 748 : pg_version: Option<u32>,
502 748 : existing_initdb_timeline_id: Option<TimelineId>,
503 748 : ) -> anyhow::Result<TimelineInfo> {
504 748 : // If timeline ID was not specified, generate one
505 748 : let new_timeline_id = new_timeline_id.unwrap_or(TimelineId::generate());
506 748 : let req = models::TimelineCreateRequest {
507 748 : new_timeline_id,
508 748 : ancestor_start_lsn,
509 748 : ancestor_timeline_id,
510 748 : pg_version,
511 748 : existing_initdb_timeline_id,
512 748 : };
513 1771 : Ok(self.http_client.timeline_create(tenant_id, &req).await?)
514 748 : }
515 :
516 : /// Import a basebackup prepared using either:
517 : /// a) `pg_basebackup -F tar`, or
518 : /// b) The `fullbackup` pageserver endpoint
519 : ///
520 : /// # Arguments
521 : /// * `tenant_id` - tenant to import into. Created if not exists
522 : /// * `timeline_id` - id to assign to imported timeline
523 : /// * `base` - (start lsn of basebackup, path to `base.tar` file)
524 : /// * `pg_wal` - if there's any wal to import: (end lsn, path to `pg_wal.tar`)
525 5 : pub async fn timeline_import(
526 5 : &self,
527 5 : tenant_id: TenantId,
528 5 : timeline_id: TimelineId,
529 5 : base: (Lsn, PathBuf),
530 5 : pg_wal: Option<(Lsn, PathBuf)>,
531 5 : pg_version: u32,
532 5 : ) -> anyhow::Result<()> {
533 15 : let (client, conn) = self.page_server_psql_client().await?;
534 : // The connection object performs the actual communication with the database,
535 : // so spawn it off to run on its own.
536 5 : tokio::spawn(async move {
537 52040 : if let Err(e) = conn.await {
538 LBC (1) : eprintln!("connection error: {}", e);
539 (1) : }
540 CBC 5 : });
541 5 : tokio::pin!(client);
542 5 :
543 5 : // Init base reader
544 5 : let (start_lsn, base_tarfile_path) = base;
545 5 : let base_tarfile = tokio::fs::File::open(base_tarfile_path).await?;
546 5 : let base_tarfile = tokio_util::io::ReaderStream::new(base_tarfile);
547 :
548 : // Init wal reader if necessary
549 5 : let (end_lsn, wal_reader) = if let Some((end_lsn, wal_tarfile_path)) = pg_wal {
550 4 : let wal_tarfile = tokio::fs::File::open(wal_tarfile_path).await?;
551 4 : let wal_reader = tokio_util::io::ReaderStream::new(wal_tarfile);
552 4 : (end_lsn, Some(wal_reader))
553 : } else {
554 1 : (start_lsn, None)
555 : };
556 :
557 7 : let copy_in = |reader, cmd| {
558 7 : let client = &client;
559 7 : async move {
560 23 : let writer = client.copy_in(&cmd).await?;
561 7 : let writer = std::pin::pin!(writer);
562 7 : let mut writer = writer.sink_map_err(|e| {
563 UBC 0 : std::io::Error::new(std::io::ErrorKind::Other, format!("{e}"))
564 CBC 7 : });
565 7 : let mut reader = std::pin::pin!(reader);
566 51996 : writer.send_all(&mut reader).await?;
567 7 : writer.into_inner().finish().await?;
568 5 : anyhow::Ok(())
569 7 : }
570 7 : };
571 :
572 : // Import base
573 5 : copy_in(
574 5 : base_tarfile,
575 5 : format!(
576 5 : "import basebackup {tenant_id} {timeline_id} {start_lsn} {end_lsn} {pg_version}"
577 5 : ),
578 5 : )
579 43830 : .await?;
580 : // Import wal if necessary
581 3 : if let Some(wal_reader) = wal_reader {
582 2 : copy_in(
583 2 : wal_reader,
584 2 : format!("import wal {tenant_id} {timeline_id} {start_lsn} {end_lsn}"),
585 2 : )
586 8196 : .await?;
587 1 : }
588 :
589 3 : Ok(())
590 5 : }
591 : }
|
