Line data Source code
1 : use std::collections::HashMap;
2 : use std::env;
3 : use std::fs;
4 : use std::io::BufRead;
5 : use std::os::unix::fs::PermissionsExt;
6 : use std::path::Path;
7 : use std::process::{Command, Stdio};
8 : use std::str::FromStr;
9 : use std::sync::atomic::AtomicU32;
10 : use std::sync::atomic::Ordering;
11 : use std::sync::{Condvar, Mutex, RwLock};
12 : use std::thread;
13 : use std::time::Instant;
14 :
15 : use anyhow::{Context, Result};
16 : use chrono::{DateTime, Utc};
17 : use futures::future::join_all;
18 : use futures::stream::FuturesUnordered;
19 : use futures::StreamExt;
20 : use postgres::{Client, NoTls};
21 : use tokio;
22 : use tokio_postgres;
23 : use tracing::{debug, error, info, instrument, warn};
24 : use utils::id::{TenantId, TimelineId};
25 : use utils::lsn::Lsn;
26 :
27 : use compute_api::responses::{ComputeMetrics, ComputeStatus};
28 : use compute_api::spec::{ComputeFeature, ComputeMode, ComputeSpec};
29 : use utils::measured_stream::MeasuredReader;
30 :
31 : use remote_storage::{DownloadError, RemotePath};
32 :
33 : use crate::checker::create_availability_check_data;
34 : use crate::logger::inlinify;
35 : use crate::pg_helpers::*;
36 : use crate::spec::*;
37 : use crate::sync_sk::{check_if_synced, ping_safekeeper};
38 : use crate::{config, extension_server};
39 :
40 : pub static SYNC_SAFEKEEPERS_PID: AtomicU32 = AtomicU32::new(0);
41 : pub static PG_PID: AtomicU32 = AtomicU32::new(0);
42 :
43 : /// Compute node info shared across several `compute_ctl` threads.
44 : pub struct ComputeNode {
45 : // Url type maintains proper escaping
46 : pub connstr: url::Url,
47 : pub pgdata: String,
48 : pub pgbin: String,
49 : pub pgversion: String,
50 : /// We should only allow live re- / configuration of the compute node if
51 : /// it uses 'pull model', i.e. it can go to control-plane and fetch
52 : /// the latest configuration. Otherwise, there could be a case:
53 : /// - we start compute with some spec provided as argument
54 : /// - we push new spec and it does reconfiguration
55 : /// - but then something happens and compute pod / VM is destroyed,
56 : /// so k8s controller starts it again with the **old** spec
57 : /// and the same for empty computes:
58 : /// - we started compute without any spec
59 : /// - we push spec and it does configuration
60 : /// - but then it is restarted without any spec again
61 : pub live_config_allowed: bool,
62 : /// Volatile part of the `ComputeNode`, which should be used under `Mutex`.
63 : /// To allow HTTP API server to serving status requests, while configuration
64 : /// is in progress, lock should be held only for short periods of time to do
65 : /// read/write, not the whole configuration process.
66 : pub state: Mutex<ComputeState>,
67 : /// `Condvar` to allow notifying waiters about state changes.
68 : pub state_changed: Condvar,
69 : /// the address of extension storage proxy gateway
70 : pub ext_remote_storage: Option<String>,
71 : // key: ext_archive_name, value: started download time, download_completed?
72 : pub ext_download_progress: RwLock<HashMap<String, (DateTime<Utc>, bool)>>,
73 : pub build_tag: String,
74 : }
75 :
76 : // store some metrics about download size that might impact startup time
77 2 : #[derive(Clone, Debug)]
78 : pub struct RemoteExtensionMetrics {
79 : num_ext_downloaded: u64,
80 : largest_ext_size: u64,
81 : total_ext_download_size: u64,
82 : }
83 :
84 1375 : #[derive(Clone, Debug)]
85 : pub struct ComputeState {
86 : pub start_time: DateTime<Utc>,
87 : pub status: ComputeStatus,
88 : /// Timestamp of the last Postgres activity. It could be `None` if
89 : /// compute wasn't used since start.
90 : pub last_active: Option<DateTime<Utc>>,
91 : pub error: Option<String>,
92 : pub pspec: Option<ParsedSpec>,
93 : pub metrics: ComputeMetrics,
94 : }
95 :
96 : impl ComputeState {
97 572 : pub fn new() -> Self {
98 572 : Self {
99 572 : start_time: Utc::now(),
100 572 : status: ComputeStatus::Empty,
101 572 : last_active: None,
102 572 : error: None,
103 572 : pspec: None,
104 572 : metrics: ComputeMetrics::default(),
105 572 : }
106 572 : }
107 : }
108 :
109 : impl Default for ComputeState {
110 0 : fn default() -> Self {
111 0 : Self::new()
112 0 : }
113 : }
114 :
115 2132 : #[derive(Clone, Debug)]
116 : pub struct ParsedSpec {
117 : pub spec: ComputeSpec,
118 : pub tenant_id: TenantId,
119 : pub timeline_id: TimelineId,
120 : pub pageserver_connstr: String,
121 : pub safekeeper_connstrings: Vec<String>,
122 : pub storage_auth_token: Option<String>,
123 : }
124 :
125 : impl TryFrom<ComputeSpec> for ParsedSpec {
126 : type Error = String;
127 803 : fn try_from(spec: ComputeSpec) -> Result<Self, String> {
128 : // Extract the options from the spec file that are needed to connect to
129 : // the storage system.
130 : //
131 : // For backwards-compatibility, the top-level fields in the spec file
132 : // may be empty. In that case, we need to dig them from the GUCs in the
133 : // cluster.settings field.
134 803 : let pageserver_connstr = spec
135 803 : .pageserver_connstring
136 803 : .clone()
137 803 : .or_else(|| spec.cluster.settings.find("neon.pageserver_connstring"))
138 803 : .ok_or("pageserver connstr should be provided")?;
139 803 : let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
140 46 : if matches!(spec.mode, ComputeMode::Primary) {
141 0 : spec.cluster
142 0 : .settings
143 0 : .find("neon.safekeepers")
144 0 : .ok_or("safekeeper connstrings should be provided")?
145 0 : .split(',')
146 0 : .map(|str| str.to_string())
147 0 : .collect()
148 : } else {
149 46 : vec![]
150 : }
151 : } else {
152 757 : spec.safekeeper_connstrings.clone()
153 : };
154 803 : let storage_auth_token = spec.storage_auth_token.clone();
155 803 : let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
156 803 : tenant_id
157 : } else {
158 0 : spec.cluster
159 0 : .settings
160 0 : .find("neon.tenant_id")
161 0 : .ok_or("tenant id should be provided")
162 0 : .map(|s| TenantId::from_str(&s))?
163 0 : .or(Err("invalid tenant id"))?
164 : };
165 803 : let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
166 803 : timeline_id
167 : } else {
168 0 : spec.cluster
169 0 : .settings
170 0 : .find("neon.timeline_id")
171 0 : .ok_or("timeline id should be provided")
172 0 : .map(|s| TimelineId::from_str(&s))?
173 0 : .or(Err("invalid timeline id"))?
174 : };
175 :
176 803 : Ok(ParsedSpec {
177 803 : spec,
178 803 : pageserver_connstr,
179 803 : safekeeper_connstrings,
180 803 : storage_auth_token,
181 803 : tenant_id,
182 803 : timeline_id,
183 803 : })
184 803 : }
185 : }
186 :
187 : /// If we are a VM, returns a [`Command`] that will run in the `neon-postgres`
188 : /// cgroup. Otherwise returns the default `Command::new(cmd)`
189 : ///
190 : /// This function should be used to start postgres, as it will start it in the
191 : /// neon-postgres cgroup if we are a VM. This allows autoscaling to control
192 : /// postgres' resource usage. The cgroup will exist in VMs because vm-builder
193 : /// creates it during the sysinit phase of its inittab.
194 1502 : fn maybe_cgexec(cmd: &str) -> Command {
195 1502 : // The cplane sets this env var for autoscaling computes.
196 1502 : // use `var_os` so we don't have to worry about the variable being valid
197 1502 : // unicode. Should never be an concern . . . but just in case
198 1502 : if env::var_os("AUTOSCALING").is_some() {
199 0 : let mut command = Command::new("cgexec");
200 0 : command.args(["-g", "memory:neon-postgres"]);
201 0 : command.arg(cmd);
202 0 : command
203 : } else {
204 1502 : Command::new(cmd)
205 : }
206 1502 : }
207 :
208 : /// Create special neon_superuser role, that's a slightly nerfed version of a real superuser
209 : /// that we give to customers
210 8 : #[instrument(skip_all)]
211 : fn create_neon_superuser(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
212 : let roles = spec
213 : .cluster
214 : .roles
215 : .iter()
216 0 : .map(|r| escape_literal(&r.name))
217 : .collect::<Vec<_>>();
218 :
219 : let dbs = spec
220 : .cluster
221 : .databases
222 : .iter()
223 0 : .map(|db| escape_literal(&db.name))
224 : .collect::<Vec<_>>();
225 :
226 : let roles_decl = if roles.is_empty() {
227 : String::from("roles text[] := NULL;")
228 : } else {
229 : format!(
230 : r#"
231 : roles text[] := ARRAY(SELECT rolname
232 : FROM pg_catalog.pg_roles
233 : WHERE rolname IN ({}));"#,
234 : roles.join(", ")
235 : )
236 : };
237 :
238 : let database_decl = if dbs.is_empty() {
239 : String::from("dbs text[] := NULL;")
240 : } else {
241 : format!(
242 : r#"
243 : dbs text[] := ARRAY(SELECT datname
244 : FROM pg_catalog.pg_database
245 : WHERE datname IN ({}));"#,
246 : dbs.join(", ")
247 : )
248 : };
249 :
250 : // ALL PRIVILEGES grants CREATE, CONNECT, and TEMPORARY on all databases
251 : // (see https://www.postgresql.org/docs/current/ddl-priv.html)
252 : let query = format!(
253 : r#"
254 : DO $$
255 : DECLARE
256 : r text;
257 : {}
258 : {}
259 : BEGIN
260 : IF NOT EXISTS (
261 : SELECT FROM pg_catalog.pg_roles WHERE rolname = 'neon_superuser')
262 : THEN
263 : CREATE ROLE neon_superuser CREATEDB CREATEROLE NOLOGIN REPLICATION BYPASSRLS IN ROLE pg_read_all_data, pg_write_all_data;
264 : IF array_length(roles, 1) IS NOT NULL THEN
265 : EXECUTE format('GRANT neon_superuser TO %s',
266 : array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(roles) as x), ', '));
267 : FOREACH r IN ARRAY roles LOOP
268 : EXECUTE format('ALTER ROLE %s CREATEROLE CREATEDB', quote_ident(r));
269 : END LOOP;
270 : END IF;
271 : IF array_length(dbs, 1) IS NOT NULL THEN
272 : EXECUTE format('GRANT ALL PRIVILEGES ON DATABASE %s TO neon_superuser',
273 : array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(dbs) as x), ', '));
274 : END IF;
275 : END IF;
276 : END
277 : $$;"#,
278 : roles_decl, database_decl,
279 : );
280 8 : info!("Neon superuser created: {}", inlinify(&query));
281 : client
282 : .simple_query(&query)
283 0 : .map_err(|e| anyhow::anyhow!(e).context(query))?;
284 : Ok(())
285 : }
286 :
287 : impl ComputeNode {
288 : /// Check that compute node has corresponding feature enabled.
289 12085 : pub fn has_feature(&self, feature: ComputeFeature) -> bool {
290 12085 : let state = self.state.lock().unwrap();
291 :
292 12085 : if let Some(s) = state.pspec.as_ref() {
293 12085 : s.spec.features.contains(&feature)
294 : } else {
295 0 : false
296 : }
297 12085 : }
298 :
299 803 : pub fn set_status(&self, status: ComputeStatus) {
300 803 : let mut state = self.state.lock().unwrap();
301 803 : state.status = status;
302 803 : self.state_changed.notify_all();
303 803 : }
304 :
305 0 : pub fn get_status(&self) -> ComputeStatus {
306 0 : self.state.lock().unwrap().status
307 0 : }
308 :
309 : // Remove `pgdata` directory and create it again with right permissions.
310 572 : fn create_pgdata(&self) -> Result<()> {
311 572 : // Ignore removal error, likely it is a 'No such file or directory (os error 2)'.
312 572 : // If it is something different then create_dir() will error out anyway.
313 572 : let _ok = fs::remove_dir_all(&self.pgdata);
314 572 : fs::create_dir(&self.pgdata)?;
315 572 : fs::set_permissions(&self.pgdata, fs::Permissions::from_mode(0o700))?;
316 :
317 572 : Ok(())
318 572 : }
319 :
320 : // Get basebackup from the libpq connection to pageserver using `connstr` and
321 : // unarchive it to `pgdata` directory overriding all its previous content.
322 572 : #[instrument(skip_all, fields(%lsn))]
323 : fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
324 : let spec = compute_state.pspec.as_ref().expect("spec must be set");
325 : let start_time = Instant::now();
326 :
327 : let mut config = postgres::Config::from_str(&spec.pageserver_connstr)?;
328 :
329 : // Use the storage auth token from the config file, if given.
330 : // Note: this overrides any password set in the connection string.
331 : if let Some(storage_auth_token) = &spec.storage_auth_token {
332 15 : info!("Got storage auth token from spec file");
333 : config.password(storage_auth_token);
334 : } else {
335 557 : info!("Storage auth token not set");
336 : }
337 :
338 : // Connect to pageserver
339 : let mut client = config.connect(NoTls)?;
340 : let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;
341 :
342 : let basebackup_cmd = match lsn {
343 : // HACK We don't use compression on first start (Lsn(0)) because there's no API for it
344 : Lsn(0) => format!("basebackup {} {}", spec.tenant_id, spec.timeline_id),
345 : _ => format!(
346 : "basebackup {} {} {} --gzip",
347 : spec.tenant_id, spec.timeline_id, lsn
348 : ),
349 : };
350 :
351 : let copyreader = client.copy_out(basebackup_cmd.as_str())?;
352 : let mut measured_reader = MeasuredReader::new(copyreader);
353 :
354 : // Check the magic number to see if it's a gzip or not. Even though
355 : // we might explicitly ask for gzip, an old pageserver with no implementation
356 : // of gzip compression might send us uncompressed data. After some time
357 : // passes we can assume all pageservers know how to compress and we can
358 : // delete this check.
359 : //
360 : // If the data is not gzip, it will be tar. It will not be mistakenly
361 : // recognized as gzip because tar starts with an ascii encoding of a filename,
362 : // and 0x1f and 0x8b are unlikely first characters for any filename. Moreover,
363 : // we send the "global" directory first from the pageserver, so it definitely
364 : // won't be recognized as gzip.
365 : let mut bufreader = std::io::BufReader::new(&mut measured_reader);
366 : let gzip = {
367 : let peek = bufreader.fill_buf().unwrap();
368 : peek[0] == 0x1f && peek[1] == 0x8b
369 : };
370 :
371 : // Read the archive directly from the `CopyOutReader`
372 : //
373 : // Set `ignore_zeros` so that unpack() reads all the Copy data and
374 : // doesn't stop at the end-of-archive marker. Otherwise, if the server
375 : // sends an Error after finishing the tarball, we will not notice it.
376 : if gzip {
377 : let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
378 : ar.set_ignore_zeros(true);
379 : ar.unpack(&self.pgdata)?;
380 : } else {
381 : let mut ar = tar::Archive::new(&mut bufreader);
382 : ar.set_ignore_zeros(true);
383 : ar.unpack(&self.pgdata)?;
384 : };
385 :
386 : // Report metrics
387 : let mut state = self.state.lock().unwrap();
388 : state.metrics.pageserver_connect_micros = pageserver_connect_micros;
389 : state.metrics.basebackup_bytes = measured_reader.get_byte_count() as u64;
390 : state.metrics.basebackup_ms = start_time.elapsed().as_millis() as u64;
391 : Ok(())
392 : }
393 :
394 : // Gets the basebackup in a retry loop
395 572 : #[instrument(skip_all, fields(%lsn))]
396 : pub fn get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
397 : let mut retry_period_ms = 500;
398 : let mut attempts = 0;
399 : let max_attempts = 5;
400 : loop {
401 : let result = self.try_get_basebackup(compute_state, lsn);
402 : match result {
403 : Ok(_) => {
404 : return result;
405 : }
406 : Err(ref e) if attempts < max_attempts => {
407 0 : warn!(
408 0 : "Failed to get basebackup: {} (attempt {}/{})",
409 0 : e, attempts, max_attempts
410 0 : );
411 : std::thread::sleep(std::time::Duration::from_millis(retry_period_ms));
412 : retry_period_ms *= 2;
413 : }
414 : Err(_) => {
415 : return result;
416 : }
417 : }
418 : attempts += 1;
419 : }
420 : }
421 :
422 526 : pub async fn check_safekeepers_synced_async(
423 526 : &self,
424 526 : compute_state: &ComputeState,
425 526 : ) -> Result<Option<Lsn>> {
426 526 : // Construct a connection config for each safekeeper
427 526 : let pspec: ParsedSpec = compute_state
428 526 : .pspec
429 526 : .as_ref()
430 526 : .expect("spec must be set")
431 526 : .clone();
432 526 : let sk_connstrs: Vec<String> = pspec.safekeeper_connstrings.clone();
433 669 : let sk_configs = sk_connstrs.into_iter().map(|connstr| {
434 669 : // Format connstr
435 669 : let id = connstr.clone();
436 669 : let connstr = format!("postgresql://no_user@{}", connstr);
437 669 : let options = format!(
438 669 : "-c timeline_id={} tenant_id={}",
439 669 : pspec.timeline_id, pspec.tenant_id
440 669 : );
441 669 :
442 669 : // Construct client
443 669 : let mut config = tokio_postgres::Config::from_str(&connstr).unwrap();
444 669 : config.options(&options);
445 669 : if let Some(storage_auth_token) = pspec.storage_auth_token.clone() {
446 26 : config.password(storage_auth_token);
447 643 : }
448 :
449 669 : (id, config)
450 669 : });
451 526 :
452 526 : // Create task set to query all safekeepers
453 526 : let mut tasks = FuturesUnordered::new();
454 526 : let quorum = sk_configs.len() / 2 + 1;
455 1195 : for (id, config) in sk_configs {
456 669 : let timeout = tokio::time::Duration::from_millis(100);
457 669 : let task = tokio::time::timeout(timeout, ping_safekeeper(id, config));
458 669 : tasks.push(tokio::spawn(task));
459 669 : }
460 :
461 : // Get a quorum of responses or errors
462 526 : let mut responses = Vec::new();
463 526 : let mut join_errors = Vec::new();
464 526 : let mut task_errors = Vec::new();
465 526 : let mut timeout_errors = Vec::new();
466 1119 : while let Some(response) = tasks.next().await {
467 610 : match response {
468 599 : Ok(Ok(Ok(r))) => responses.push(r),
469 11 : Ok(Ok(Err(e))) => task_errors.push(e),
470 0 : Ok(Err(e)) => timeout_errors.push(e),
471 0 : Err(e) => join_errors.push(e),
472 : };
473 610 : if responses.len() >= quorum {
474 526 : break;
475 84 : }
476 84 : if join_errors.len() + task_errors.len() + timeout_errors.len() >= quorum {
477 0 : break;
478 84 : }
479 : }
480 :
481 : // In case of error, log and fail the check, but don't crash.
482 : // We're playing it safe because these errors could be transient
483 : // and we don't yet retry. Also being careful here allows us to
484 : // be backwards compatible with safekeepers that don't have the
485 : // TIMELINE_STATUS API yet.
486 526 : if responses.len() < quorum {
487 0 : error!(
488 0 : "failed sync safekeepers check {:?} {:?} {:?}",
489 0 : join_errors, task_errors, timeout_errors
490 0 : );
491 0 : return Ok(None);
492 526 : }
493 526 :
494 526 : Ok(check_if_synced(responses))
495 526 : }
496 :
497 : // Fast path for sync_safekeepers. If they're already synced we get the lsn
498 : // in one roundtrip. If not, we should do a full sync_safekeepers.
499 526 : pub fn check_safekeepers_synced(&self, compute_state: &ComputeState) -> Result<Option<Lsn>> {
500 526 : let start_time = Utc::now();
501 526 :
502 526 : // Run actual work with new tokio runtime
503 526 : let rt = tokio::runtime::Builder::new_current_thread()
504 526 : .enable_all()
505 526 : .build()
506 526 : .expect("failed to create rt");
507 526 : let result = rt.block_on(self.check_safekeepers_synced_async(compute_state));
508 526 :
509 526 : // Record runtime
510 526 : self.state.lock().unwrap().metrics.sync_sk_check_ms = Utc::now()
511 526 : .signed_duration_since(start_time)
512 526 : .to_std()
513 526 : .unwrap()
514 526 : .as_millis() as u64;
515 526 : result
516 526 : }
517 :
518 : // Run `postgres` in a special mode with `--sync-safekeepers` argument
519 : // and return the reported LSN back to the caller.
520 930 : #[instrument(skip_all)]
521 : pub fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
522 : let start_time = Utc::now();
523 :
524 : let mut sync_handle = maybe_cgexec(&self.pgbin)
525 : .args(["--sync-safekeepers"])
526 : .env("PGDATA", &self.pgdata) // we cannot use -D in this mode
527 : .envs(if let Some(storage_auth_token) = &storage_auth_token {
528 : vec![("NEON_AUTH_TOKEN", storage_auth_token)]
529 : } else {
530 : vec![]
531 : })
532 : .stdout(Stdio::piped())
533 : .stderr(Stdio::piped())
534 : .spawn()
535 : .expect("postgres --sync-safekeepers failed to start");
536 : SYNC_SAFEKEEPERS_PID.store(sync_handle.id(), Ordering::SeqCst);
537 :
538 : // `postgres --sync-safekeepers` will print all log output to stderr and
539 : // final LSN to stdout. So we leave stdout to collect LSN, while stderr logs
540 : // will be collected in a child thread.
541 : let stderr = sync_handle
542 : .stderr
543 : .take()
544 : .expect("stderr should be captured");
545 : let logs_handle = handle_postgres_logs(stderr);
546 :
547 : let sync_output = sync_handle
548 : .wait_with_output()
549 : .expect("postgres --sync-safekeepers failed");
550 : SYNC_SAFEKEEPERS_PID.store(0, Ordering::SeqCst);
551 :
552 : // Process has exited, so we can join the logs thread.
553 : let _ = logs_handle
554 : .join()
555 0 : .map_err(|e| tracing::error!("log thread panicked: {:?}", e));
556 :
557 : if !sync_output.status.success() {
558 : anyhow::bail!(
559 : "postgres --sync-safekeepers exited with non-zero status: {}. stdout: {}",
560 : sync_output.status,
561 : String::from_utf8(sync_output.stdout)
562 : .expect("postgres --sync-safekeepers exited, and stdout is not utf-8"),
563 : );
564 : }
565 :
566 : self.state.lock().unwrap().metrics.sync_safekeepers_ms = Utc::now()
567 : .signed_duration_since(start_time)
568 : .to_std()
569 : .unwrap()
570 : .as_millis() as u64;
571 :
572 : let lsn = Lsn::from_str(String::from_utf8(sync_output.stdout)?.trim())?;
573 :
574 : Ok(lsn)
575 : }
576 :
577 : /// Do all the preparations like PGDATA directory creation, configuration,
578 : /// safekeepers sync, basebackup, etc.
579 572 : #[instrument(skip_all)]
580 : pub fn prepare_pgdata(
581 : &self,
582 : compute_state: &ComputeState,
583 : extension_server_port: u16,
584 : ) -> Result<()> {
585 : let pspec = compute_state.pspec.as_ref().expect("spec must be set");
586 : let spec = &pspec.spec;
587 : let pgdata_path = Path::new(&self.pgdata);
588 :
589 : // Remove/create an empty pgdata directory and put configuration there.
590 : self.create_pgdata()?;
591 : config::write_postgres_conf(
592 : &pgdata_path.join("postgresql.conf"),
593 : &pspec.spec,
594 : Some(extension_server_port),
595 : )?;
596 :
597 : // Syncing safekeepers is only safe with primary nodes: if a primary
598 : // is already connected it will be kicked out, so a secondary (standby)
599 : // cannot sync safekeepers.
600 : let lsn = match spec.mode {
601 : ComputeMode::Primary => {
602 526 : info!("checking if safekeepers are synced");
603 : let lsn = if let Ok(Some(lsn)) = self.check_safekeepers_synced(compute_state) {
604 : lsn
605 : } else {
606 404 : info!("starting safekeepers syncing");
607 : self.sync_safekeepers(pspec.storage_auth_token.clone())
608 0 : .with_context(|| "failed to sync safekeepers")?
609 : };
610 526 : info!("safekeepers synced at LSN {}", lsn);
611 : lsn
612 : }
613 : ComputeMode::Static(lsn) => {
614 44 : info!("Starting read-only node at static LSN {}", lsn);
615 : lsn
616 : }
617 : ComputeMode::Replica => {
618 2 : info!("Initializing standby from latest Pageserver LSN");
619 : Lsn(0)
620 : }
621 : };
622 :
623 572 : info!(
624 572 : "getting basebackup@{} from pageserver {}",
625 572 : lsn, &pspec.pageserver_connstr
626 572 : );
627 0 : self.get_basebackup(compute_state, lsn).with_context(|| {
628 0 : format!(
629 0 : "failed to get basebackup@{} from pageserver {}",
630 0 : lsn, &pspec.pageserver_connstr
631 0 : )
632 0 : })?;
633 :
634 : // Update pg_hba.conf received with basebackup.
635 : update_pg_hba(pgdata_path)?;
636 :
637 : match spec.mode {
638 : ComputeMode::Primary => {}
639 : ComputeMode::Replica | ComputeMode::Static(..) => {
640 : add_standby_signal(pgdata_path)?;
641 : }
642 : }
643 :
644 : Ok(())
645 : }
646 :
647 : /// Start and stop a postgres process to warm up the VM for startup.
648 0 : pub fn prewarm_postgres(&self) -> Result<()> {
649 0 : info!("prewarming");
650 :
651 : // Create pgdata
652 0 : let pgdata = &format!("{}.warmup", self.pgdata);
653 0 : create_pgdata(pgdata)?;
654 :
655 : // Run initdb to completion
656 0 : info!("running initdb");
657 0 : let initdb_bin = Path::new(&self.pgbin).parent().unwrap().join("initdb");
658 0 : Command::new(initdb_bin)
659 0 : .args(["-D", pgdata])
660 0 : .output()
661 0 : .expect("cannot start initdb process");
662 0 :
663 0 : // Write conf
664 0 : use std::io::Write;
665 0 : let conf_path = Path::new(pgdata).join("postgresql.conf");
666 0 : let mut file = std::fs::File::create(conf_path)?;
667 0 : writeln!(file, "shared_buffers=65536")?;
668 0 : writeln!(file, "port=51055")?; // Nobody should be connecting
669 0 : writeln!(file, "shared_preload_libraries = 'neon'")?;
670 :
671 : // Start postgres
672 0 : info!("starting postgres");
673 0 : let mut pg = maybe_cgexec(&self.pgbin)
674 0 : .args(["-D", pgdata])
675 0 : .spawn()
676 0 : .expect("cannot start postgres process");
677 0 :
678 0 : // Stop it when it's ready
679 0 : info!("waiting for postgres");
680 0 : wait_for_postgres(&mut pg, Path::new(pgdata))?;
681 0 : pg.kill()?;
682 0 : info!("sent kill signal");
683 0 : pg.wait()?;
684 0 : info!("done prewarming");
685 :
686 : // clean up
687 0 : let _ok = fs::remove_dir_all(pgdata);
688 0 : Ok(())
689 0 : }
690 :
691 : /// Start Postgres as a child process and manage DBs/roles.
692 : /// After that this will hang waiting on the postmaster process to exit.
693 : /// Returns a handle to the child process and a handle to the logs thread.
694 572 : #[instrument(skip_all)]
695 : pub fn start_postgres(
696 : &self,
697 : storage_auth_token: Option<String>,
698 : ) -> Result<(std::process::Child, std::thread::JoinHandle<()>)> {
699 : let pgdata_path = Path::new(&self.pgdata);
700 :
701 : // Run postgres as a child process.
702 : let mut pg = maybe_cgexec(&self.pgbin)
703 : .args(["-D", &self.pgdata])
704 : .envs(if let Some(storage_auth_token) = &storage_auth_token {
705 : vec![("NEON_AUTH_TOKEN", storage_auth_token)]
706 : } else {
707 : vec![]
708 : })
709 : .stderr(Stdio::piped())
710 : .spawn()
711 : .expect("cannot start postgres process");
712 : PG_PID.store(pg.id(), Ordering::SeqCst);
713 :
714 : // Start a thread to collect logs from stderr.
715 : let stderr = pg.stderr.take().expect("stderr should be captured");
716 : let logs_handle = handle_postgres_logs(stderr);
717 :
718 : wait_for_postgres(&mut pg, pgdata_path)?;
719 :
720 : Ok((pg, logs_handle))
721 : }
722 :
723 : /// Do initial configuration of the already started Postgres.
724 8 : #[instrument(skip_all)]
725 : pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
726 : // If connection fails,
727 : // it may be the old node with `zenith_admin` superuser.
728 : //
729 : // In this case we need to connect with old `zenith_admin` name
730 : // and create new user. We cannot simply rename connected user,
731 : // but we can create a new one and grant it all privileges.
732 : let connstr = self.connstr.clone();
733 : let mut client = match Client::connect(connstr.as_str(), NoTls) {
734 : Err(e) => {
735 0 : info!(
736 0 : "cannot connect to postgres: {}, retrying with `zenith_admin` username",
737 0 : e
738 0 : );
739 : let mut zenith_admin_connstr = connstr.clone();
740 :
741 : zenith_admin_connstr
742 : .set_username("zenith_admin")
743 0 : .map_err(|_| anyhow::anyhow!("invalid connstr"))?;
744 :
745 : let mut client = Client::connect(zenith_admin_connstr.as_str(), NoTls)?;
746 : // Disable forwarding so that users don't get a cloud_admin role
747 : client.simple_query("SET neon.forward_ddl = false")?;
748 : client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?;
749 : client.simple_query("GRANT zenith_admin TO cloud_admin")?;
750 : drop(client);
751 :
752 : // reconnect with connstring with expected name
753 : Client::connect(connstr.as_str(), NoTls)?
754 : }
755 : Ok(client) => client,
756 : };
757 :
758 : // Disable DDL forwarding because control plane already knows about these roles/databases.
759 : client.simple_query("SET neon.forward_ddl = false")?;
760 :
761 : // Proceed with post-startup configuration. Note, that order of operations is important.
762 : let spec = &compute_state.pspec.as_ref().expect("spec must be set").spec;
763 : create_neon_superuser(spec, &mut client)?;
764 : cleanup_instance(&mut client)?;
765 : handle_roles(spec, &mut client)?;
766 : handle_databases(spec, &mut client)?;
767 : handle_role_deletions(spec, connstr.as_str(), &mut client)?;
768 : handle_grants(
769 : spec,
770 : &mut client,
771 : connstr.as_str(),
772 : self.has_feature(ComputeFeature::AnonExtension),
773 : )?;
774 : handle_extensions(spec, &mut client)?;
775 : handle_extension_neon(&mut client)?;
776 : create_availability_check_data(&mut client)?;
777 :
778 : // 'Close' connection
779 : drop(client);
780 :
781 : // Run migrations separately to not hold up cold starts
782 8 : thread::spawn(move || {
783 8 : let mut client = Client::connect(connstr.as_str(), NoTls)?;
784 8 : handle_migrations(&mut client)
785 8 : });
786 : Ok(())
787 : }
788 :
789 : // We could've wrapped this around `pg_ctl reload`, but right now we don't use
790 : // `pg_ctl` for start / stop, so this just seems much easier to do as we already
791 : // have opened connection to Postgres and superuser access.
792 478 : #[instrument(skip_all)]
793 : fn pg_reload_conf(&self) -> Result<()> {
794 : let pgctl_bin = Path::new(&self.pgbin).parent().unwrap().join("pg_ctl");
795 : Command::new(pgctl_bin)
796 : .args(["reload", "-D", &self.pgdata])
797 : .output()
798 : .expect("cannot run pg_ctl process");
799 : Ok(())
800 : }
801 :
802 : /// Similar to `apply_config()`, but does a bit different sequence of operations,
803 : /// as it's used to reconfigure a previously started and configured Postgres node.
804 231 : #[instrument(skip_all)]
805 : pub fn reconfigure(&self) -> Result<()> {
806 : let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
807 :
808 : if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
809 0 : info!("tuning pgbouncer");
810 :
811 : let rt = tokio::runtime::Builder::new_current_thread()
812 : .enable_all()
813 : .build()
814 : .expect("failed to create rt");
815 :
816 : // Spawn a thread to do the tuning,
817 : // so that we don't block the main thread that starts Postgres.
818 : let pgbouncer_settings = pgbouncer_settings.clone();
819 0 : let _handle = thread::spawn(move || {
820 0 : let res = rt.block_on(tune_pgbouncer(pgbouncer_settings));
821 0 : if let Err(err) = res {
822 0 : error!("error while tuning pgbouncer: {err:?}");
823 0 : }
824 0 : });
825 : }
826 :
827 : // Write new config
828 : let pgdata_path = Path::new(&self.pgdata);
829 : let postgresql_conf_path = pgdata_path.join("postgresql.conf");
830 : config::write_postgres_conf(&postgresql_conf_path, &spec, None)?;
831 : // temporarily reset max_cluster_size in config
832 : // to avoid the possibility of hitting the limit, while we are reconfiguring:
833 : // creating new extensions, roles, etc...
834 : config::compute_ctl_temp_override_create(pgdata_path, "neon.max_cluster_size=-1")?;
835 : self.pg_reload_conf()?;
836 :
837 : let mut client = Client::connect(self.connstr.as_str(), NoTls)?;
838 :
839 : // Proceed with post-startup configuration. Note, that order of operations is important.
840 : // Disable DDL forwarding because control plane already knows about these roles/databases.
841 : if spec.mode == ComputeMode::Primary {
842 : client.simple_query("SET neon.forward_ddl = false")?;
843 : cleanup_instance(&mut client)?;
844 : handle_roles(&spec, &mut client)?;
845 : handle_databases(&spec, &mut client)?;
846 : handle_role_deletions(&spec, self.connstr.as_str(), &mut client)?;
847 : handle_grants(
848 : &spec,
849 : &mut client,
850 : self.connstr.as_str(),
851 : self.has_feature(ComputeFeature::AnonExtension),
852 : )?;
853 : handle_extensions(&spec, &mut client)?;
854 : handle_extension_neon(&mut client)?;
855 : // We can skip handle_migrations here because a new migration can only appear
856 : // if we have a new version of the compute_ctl binary, which can only happen
857 : // if compute got restarted, in which case we'll end up inside of apply_config
858 : // instead of reconfigure.
859 : }
860 :
861 : // 'Close' connection
862 : drop(client);
863 :
864 : // reset max_cluster_size in config back to original value and reload config
865 : config::compute_ctl_temp_override_remove(pgdata_path)?;
866 : self.pg_reload_conf()?;
867 :
868 : let unknown_op = "unknown".to_string();
869 : let op_id = spec.operation_uuid.as_ref().unwrap_or(&unknown_op);
870 231 : info!(
871 231 : "finished reconfiguration of compute node for operation {}",
872 231 : op_id
873 231 : );
874 :
875 : Ok(())
876 : }
877 :
878 572 : #[instrument(skip_all)]
879 : pub fn start_compute(
880 : &self,
881 : extension_server_port: u16,
882 : ) -> Result<(std::process::Child, std::thread::JoinHandle<()>)> {
883 : let compute_state = self.state.lock().unwrap().clone();
884 : let pspec = compute_state.pspec.as_ref().expect("spec must be set");
885 572 : info!(
886 572 : "starting compute for project {}, operation {}, tenant {}, timeline {}",
887 572 : pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
888 572 : pspec.spec.operation_uuid.as_deref().unwrap_or("None"),
889 572 : pspec.tenant_id,
890 572 : pspec.timeline_id,
891 572 : );
892 :
893 : // tune pgbouncer
894 : if let Some(pgbouncer_settings) = &pspec.spec.pgbouncer_settings {
895 0 : info!("tuning pgbouncer");
896 :
897 : let rt = tokio::runtime::Builder::new_current_thread()
898 : .enable_all()
899 : .build()
900 : .expect("failed to create rt");
901 :
902 : // Spawn a thread to do the tuning,
903 : // so that we don't block the main thread that starts Postgres.
904 : let pgbouncer_settings = pgbouncer_settings.clone();
905 0 : let _handle = thread::spawn(move || {
906 0 : let res = rt.block_on(tune_pgbouncer(pgbouncer_settings));
907 0 : if let Err(err) = res {
908 0 : error!("error while tuning pgbouncer: {err:?}");
909 0 : }
910 0 : });
911 : }
912 :
913 572 : info!(
914 572 : "start_compute spec.remote_extensions {:?}",
915 572 : pspec.spec.remote_extensions
916 572 : );
917 :
918 : // This part is sync, because we need to download
919 : // remote shared_preload_libraries before postgres start (if any)
920 : if let Some(remote_extensions) = &pspec.spec.remote_extensions {
921 : // First, create control files for all availale extensions
922 : extension_server::create_control_files(remote_extensions, &self.pgbin);
923 :
924 : let library_load_start_time = Utc::now();
925 : let remote_ext_metrics = self.prepare_preload_libraries(&pspec.spec)?;
926 :
927 : let library_load_time = Utc::now()
928 : .signed_duration_since(library_load_start_time)
929 : .to_std()
930 : .unwrap()
931 : .as_millis() as u64;
932 : let mut state = self.state.lock().unwrap();
933 : state.metrics.load_ext_ms = library_load_time;
934 : state.metrics.num_ext_downloaded = remote_ext_metrics.num_ext_downloaded;
935 : state.metrics.largest_ext_size = remote_ext_metrics.largest_ext_size;
936 : state.metrics.total_ext_download_size = remote_ext_metrics.total_ext_download_size;
937 1 : info!(
938 1 : "Loading shared_preload_libraries took {:?}ms",
939 1 : library_load_time
940 1 : );
941 1 : info!("{:?}", remote_ext_metrics);
942 : }
943 :
944 : self.prepare_pgdata(&compute_state, extension_server_port)?;
945 :
946 : let start_time = Utc::now();
947 : let pg_process = self.start_postgres(pspec.storage_auth_token.clone())?;
948 :
949 : let config_time = Utc::now();
950 : if pspec.spec.mode == ComputeMode::Primary && !pspec.spec.skip_pg_catalog_updates {
951 : let pgdata_path = Path::new(&self.pgdata);
952 : // temporarily reset max_cluster_size in config
953 : // to avoid the possibility of hitting the limit, while we are applying config:
954 : // creating new extensions, roles, etc...
955 : config::compute_ctl_temp_override_create(pgdata_path, "neon.max_cluster_size=-1")?;
956 : self.pg_reload_conf()?;
957 :
958 : self.apply_config(&compute_state)?;
959 :
960 : config::compute_ctl_temp_override_remove(pgdata_path)?;
961 : self.pg_reload_conf()?;
962 : }
963 :
964 : let startup_end_time = Utc::now();
965 : {
966 : let mut state = self.state.lock().unwrap();
967 : state.metrics.start_postgres_ms = config_time
968 : .signed_duration_since(start_time)
969 : .to_std()
970 : .unwrap()
971 : .as_millis() as u64;
972 : state.metrics.config_ms = startup_end_time
973 : .signed_duration_since(config_time)
974 : .to_std()
975 : .unwrap()
976 : .as_millis() as u64;
977 : state.metrics.total_startup_ms = startup_end_time
978 : .signed_duration_since(compute_state.start_time)
979 : .to_std()
980 : .unwrap()
981 : .as_millis() as u64;
982 : }
983 : self.set_status(ComputeStatus::Running);
984 :
985 572 : info!(
986 572 : "finished configuration of compute for project {}",
987 572 : pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None")
988 572 : );
989 :
990 : // Log metrics so that we can search for slow operations in logs
991 : let metrics = {
992 : let state = self.state.lock().unwrap();
993 : state.metrics.clone()
994 : };
995 572 : info!(?metrics, "compute start finished");
996 :
997 : Ok(pg_process)
998 : }
999 :
1000 : /// Update the `last_active` in the shared state, but ensure that it's a more recent one.
1001 11502 : pub fn update_last_active(&self, last_active: Option<DateTime<Utc>>) {
1002 11502 : let mut state = self.state.lock().unwrap();
1003 11502 : // NB: `Some(<DateTime>)` is always greater than `None`.
1004 11502 : if last_active > state.last_active {
1005 278 : state.last_active = last_active;
1006 278 : debug!("set the last compute activity time to: {:?}", last_active);
1007 11224 : }
1008 11502 : }
1009 :
1010 : // Look for core dumps and collect backtraces.
1011 : //
1012 : // EKS worker nodes have following core dump settings:
1013 : // /proc/sys/kernel/core_pattern -> core
1014 : // /proc/sys/kernel/core_uses_pid -> 1
1015 : // ulimint -c -> unlimited
1016 : // which results in core dumps being written to postgres data directory as core.<pid>.
1017 : //
1018 : // Use that as a default location and pattern, except macos where core dumps are written
1019 : // to /cores/ directory by default.
1020 571 : pub fn check_for_core_dumps(&self) -> Result<()> {
1021 571 : let core_dump_dir = match std::env::consts::OS {
1022 571 : "macos" => Path::new("/cores/"),
1023 571 : _ => Path::new(&self.pgdata),
1024 : };
1025 :
1026 : // Collect core dump paths if any
1027 571 : info!("checking for core dumps in {}", core_dump_dir.display());
1028 571 : let files = fs::read_dir(core_dump_dir)?;
1029 13755 : let cores = files.filter_map(|entry| {
1030 13755 : let entry = entry.ok()?;
1031 13755 : let _ = entry.file_name().to_str()?.strip_prefix("core.")?;
1032 0 : Some(entry.path())
1033 13755 : });
1034 :
1035 : // Print backtrace for each core dump
1036 571 : for core_path in cores {
1037 0 : warn!(
1038 0 : "core dump found: {}, collecting backtrace",
1039 0 : core_path.display()
1040 0 : );
1041 :
1042 : // Try first with gdb
1043 0 : let backtrace = Command::new("gdb")
1044 0 : .args(["--batch", "-q", "-ex", "bt", &self.pgbin])
1045 0 : .arg(&core_path)
1046 0 : .output();
1047 :
1048 : // Try lldb if no gdb is found -- that is handy for local testing on macOS
1049 0 : let backtrace = match backtrace {
1050 0 : Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
1051 0 : warn!("cannot find gdb, trying lldb");
1052 0 : Command::new("lldb")
1053 0 : .arg("-c")
1054 0 : .arg(&core_path)
1055 0 : .args(["--batch", "-o", "bt all", "-o", "quit"])
1056 0 : .output()
1057 : }
1058 0 : _ => backtrace,
1059 0 : }?;
1060 :
1061 0 : warn!(
1062 0 : "core dump backtrace: {}",
1063 0 : String::from_utf8_lossy(&backtrace.stdout)
1064 0 : );
1065 0 : warn!(
1066 0 : "debugger stderr: {}",
1067 0 : String::from_utf8_lossy(&backtrace.stderr)
1068 0 : );
1069 : }
1070 :
1071 571 : Ok(())
1072 571 : }
1073 :
1074 : /// Select `pg_stat_statements` data and return it as a stringified JSON
1075 0 : pub async fn collect_insights(&self) -> String {
1076 0 : let mut result_rows: Vec<String> = Vec::new();
1077 0 : let connect_result = tokio_postgres::connect(self.connstr.as_str(), NoTls).await;
1078 0 : let (client, connection) = connect_result.unwrap();
1079 0 : tokio::spawn(async move {
1080 0 : if let Err(e) = connection.await {
1081 0 : eprintln!("connection error: {}", e);
1082 0 : }
1083 0 : });
1084 0 : let result = client
1085 0 : .simple_query(
1086 0 : "SELECT
1087 0 : row_to_json(pg_stat_statements)
1088 0 : FROM
1089 0 : pg_stat_statements
1090 0 : WHERE
1091 0 : userid != 'cloud_admin'::regrole::oid
1092 0 : ORDER BY
1093 0 : (mean_exec_time + mean_plan_time) DESC
1094 0 : LIMIT 100",
1095 0 : )
1096 0 : .await;
1097 :
1098 0 : if let Ok(raw_rows) = result {
1099 0 : for message in raw_rows.iter() {
1100 0 : if let postgres::SimpleQueryMessage::Row(row) = message {
1101 0 : if let Some(json) = row.get(0) {
1102 0 : result_rows.push(json.to_string());
1103 0 : }
1104 0 : }
1105 : }
1106 :
1107 0 : format!("{{\"pg_stat_statements\": [{}]}}", result_rows.join(","))
1108 : } else {
1109 0 : "{{\"pg_stat_statements\": []}}".to_string()
1110 : }
1111 0 : }
1112 :
1113 : // download an archive, unzip and place files in correct locations
1114 2 : pub async fn download_extension(
1115 2 : &self,
1116 2 : real_ext_name: String,
1117 2 : ext_path: RemotePath,
1118 2 : ) -> Result<u64, DownloadError> {
1119 2 : let ext_remote_storage =
1120 2 : self.ext_remote_storage
1121 2 : .as_ref()
1122 2 : .ok_or(DownloadError::BadInput(anyhow::anyhow!(
1123 2 : "Remote extensions storage is not configured",
1124 2 : )))?;
1125 :
1126 2 : let ext_archive_name = ext_path.object_name().expect("bad path");
1127 2 :
1128 2 : let mut first_try = false;
1129 2 : if !self
1130 2 : .ext_download_progress
1131 2 : .read()
1132 2 : .expect("lock err")
1133 2 : .contains_key(ext_archive_name)
1134 1 : {
1135 1 : self.ext_download_progress
1136 1 : .write()
1137 1 : .expect("lock err")
1138 1 : .insert(ext_archive_name.to_string(), (Utc::now(), false));
1139 1 : first_try = true;
1140 1 : }
1141 2 : let (download_start, download_completed) =
1142 2 : self.ext_download_progress.read().expect("lock err")[ext_archive_name];
1143 2 : let start_time_delta = Utc::now()
1144 2 : .signed_duration_since(download_start)
1145 2 : .to_std()
1146 2 : .unwrap()
1147 2 : .as_millis() as u64;
1148 2 :
1149 2 : // how long to wait for extension download if it was started by another process
1150 2 : const HANG_TIMEOUT: u64 = 3000; // milliseconds
1151 2 :
1152 2 : if download_completed {
1153 1 : info!("extension already downloaded, skipping re-download");
1154 1 : return Ok(0);
1155 1 : } else if start_time_delta < HANG_TIMEOUT && !first_try {
1156 0 : info!("download {ext_archive_name} already started by another process, hanging untill completion or timeout");
1157 0 : let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500));
1158 : loop {
1159 0 : info!("waiting for download");
1160 0 : interval.tick().await;
1161 0 : let (_, download_completed_now) =
1162 0 : self.ext_download_progress.read().expect("lock")[ext_archive_name];
1163 0 : if download_completed_now {
1164 0 : info!("download finished by whoever else downloaded it");
1165 0 : return Ok(0);
1166 0 : }
1167 : }
1168 : // NOTE: the above loop will get terminated
1169 : // based on the timeout of the download function
1170 1 : }
1171 :
1172 : // if extension hasn't been downloaded before or the previous
1173 : // attempt to download was at least HANG_TIMEOUT ms ago
1174 : // then we try to download it here
1175 1 : info!("downloading new extension {ext_archive_name}");
1176 :
1177 1 : let download_size = extension_server::download_extension(
1178 1 : &real_ext_name,
1179 1 : &ext_path,
1180 1 : ext_remote_storage,
1181 1 : &self.pgbin,
1182 1 : )
1183 32 : .await
1184 1 : .map_err(DownloadError::Other);
1185 1 :
1186 1 : self.ext_download_progress
1187 1 : .write()
1188 1 : .expect("bad lock")
1189 1 : .insert(ext_archive_name.to_string(), (download_start, true));
1190 1 :
1191 1 : download_size
1192 2 : }
1193 :
1194 : #[tokio::main]
1195 1 : pub async fn prepare_preload_libraries(
1196 1 : &self,
1197 1 : spec: &ComputeSpec,
1198 1 : ) -> Result<RemoteExtensionMetrics> {
1199 1 : if self.ext_remote_storage.is_none() {
1200 1 : return Ok(RemoteExtensionMetrics {
1201 0 : num_ext_downloaded: 0,
1202 0 : largest_ext_size: 0,
1203 0 : total_ext_download_size: 0,
1204 0 : });
1205 1 : }
1206 1 : let remote_extensions = spec
1207 1 : .remote_extensions
1208 1 : .as_ref()
1209 1 : .ok_or(anyhow::anyhow!("Remote extensions are not configured"))?;
1210 1 :
1211 1 : info!("parse shared_preload_libraries from spec.cluster.settings");
1212 1 : let mut libs_vec = Vec::new();
1213 1 : if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
1214 0 : libs_vec = libs
1215 0 : .split(&[',', '\'', ' '])
1216 0 : .filter(|s| *s != "neon" && !s.is_empty())
1217 0 : .map(str::to_string)
1218 0 : .collect();
1219 1 : }
1220 1 : info!("parse shared_preload_libraries from provided postgresql.conf");
1221 1 :
1222 1 : // that is used in neon_local and python tests
1223 1 : if let Some(conf) = &spec.cluster.postgresql_conf {
1224 1 : let conf_lines = conf.split('\n').collect::<Vec<&str>>();
1225 1 : let mut shared_preload_libraries_line = "";
1226 22 : for line in conf_lines {
1227 21 : if line.starts_with("shared_preload_libraries") {
1228 1 : shared_preload_libraries_line = line;
1229 20 : }
1230 1 : }
1231 1 : let mut preload_libs_vec = Vec::new();
1232 1 : if let Some(libs) = shared_preload_libraries_line.split("='").nth(1) {
1233 0 : preload_libs_vec = libs
1234 0 : .split(&[',', '\'', ' '])
1235 0 : .filter(|s| *s != "neon" && !s.is_empty())
1236 0 : .map(str::to_string)
1237 0 : .collect();
1238 1 : }
1239 1 : libs_vec.extend(preload_libs_vec);
1240 1 : }
1241 1 :
1242 1 : // Don't try to download libraries that are not in the index.
1243 1 : // Assume that they are already present locally.
1244 1 : libs_vec.retain(|lib| remote_extensions.library_index.contains_key(lib));
1245 1 :
1246 1 : info!("Downloading to shared preload libraries: {:?}", &libs_vec);
1247 1 :
1248 1 : let mut download_tasks = Vec::new();
1249 1 : for library in &libs_vec {
1250 0 : let (ext_name, ext_path) =
1251 1 : remote_extensions.get_ext(library, true, &self.build_tag, &self.pgversion)?;
1252 1 : download_tasks.push(self.download_extension(ext_name, ext_path));
1253 1 : }
1254 1 : let results = join_all(download_tasks).await;
1255 1 :
1256 1 : let mut remote_ext_metrics = RemoteExtensionMetrics {
1257 1 : num_ext_downloaded: 0,
1258 1 : largest_ext_size: 0,
1259 1 : total_ext_download_size: 0,
1260 1 : };
1261 1 : for result in results {
1262 0 : let download_size = match result {
1263 1 : Ok(res) => {
1264 0 : remote_ext_metrics.num_ext_downloaded += 1;
1265 0 : res
1266 1 : }
1267 1 : Err(err) => {
1268 0 : // if we failed to download an extension, we don't want to fail the whole
1269 0 : // process, but we do want to log the error
1270 0 : error!("Failed to download extension: {}", err);
1271 1 : 0
1272 1 : }
1273 1 : };
1274 1 :
1275 1 : remote_ext_metrics.largest_ext_size =
1276 0 : std::cmp::max(remote_ext_metrics.largest_ext_size, download_size);
1277 0 : remote_ext_metrics.total_ext_download_size += download_size;
1278 1 : }
1279 1 : Ok(remote_ext_metrics)
1280 1 : }
1281 : }
|
