From c09d817c983403425bbe484f2ce5dc599b4cb569 Mon Sep 17 00:00:00 2001 From: Gleb Novikov Date: Wed, 15 Jan 2025 17:13:52 +0000 Subject: [PATCH 1/3] review comments --- compute_tools/src/bin/fast_import.rs | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/compute_tools/src/bin/fast_import.rs b/compute_tools/src/bin/fast_import.rs index dd1480eeff..8bab37d5d9 100644 --- a/compute_tools/src/bin/fast_import.rs +++ b/compute_tools/src/bin/fast_import.rs @@ -87,6 +87,9 @@ pub(crate) async fn main() -> anyhow::Result<()> { if args.s3_prefix.is_none() && args.source_connection_string.is_none() { anyhow::bail!("either s3_prefix or source_connection_string must be specified"); } + if args.s3_prefix.is_some() && args.source_connection_string.is_some() { + anyhow::bail!("only one of s3_prefix or source_connection_string can be specified"); + } let working_directory = args.working_directory; let pg_bin_dir = args.pg_bin_dir; @@ -227,13 +230,14 @@ pub(crate) async fn main() -> anyhow::Result<()> { // Create neondb database in the running postgres let restore_pg_connstring = format!("host=localhost port=5432 user={superuser} dbname=postgres"); + loop { match tokio_postgres::connect(&restore_pg_connstring, tokio_postgres::NoTls).await { Ok((client, connection)) => { // Spawn the connection handling task to maintain the connection tokio::spawn(async move { if let Err(e) = connection.await { - eprintln!("connection error: {}", e); + warn!("connection error: {}", e); } }); @@ -243,12 +247,16 @@ pub(crate) async fn main() -> anyhow::Result<()> { break; } Err(e) => { - info!("failed to create database: {}", e); - break; + warn!("failed to create database: {}", e); + continue; } } } - Err(_) => continue, + Err(_) => { + info!("postgres not ready yet, retrying in 0.3s"); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.3)).await; + continue; + }, } } From 4c2ee6a011edd1f3df7982eeb492f91217bc869d Mon Sep 17 00:00:00 2001 From: Gleb Novikov Date: Wed, 15 Jan 2025 17:15:58 +0000 Subject: [PATCH 2/3] added 10 min timeout on waiting loop --- compute_tools/src/bin/fast_import.rs | 29 +++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/compute_tools/src/bin/fast_import.rs b/compute_tools/src/bin/fast_import.rs index 8bab37d5d9..c2de4c3c0a 100644 --- a/compute_tools/src/bin/fast_import.rs +++ b/compute_tools/src/bin/fast_import.rs @@ -31,7 +31,7 @@ use camino::{Utf8Path, Utf8PathBuf}; use clap::Parser; use compute_tools::extension_server::{get_pg_version, PostgresMajorVersion}; use nix::unistd::Pid; -use tracing::{info, info_span, warn, Instrument}; +use tracing::{error, info, info_span, warn, Instrument}; use utils::fs_ext::is_directory_empty; #[path = "fast_import/aws_s3_sync.rs"] @@ -231,7 +231,18 @@ pub(crate) async fn main() -> anyhow::Result<()> { let restore_pg_connstring = format!("host=localhost port=5432 user={superuser} dbname=postgres"); + let timeout_duration = std::time::Duration::from_secs(600); // 10 minutes + let start_time = std::time::Instant::now(); + let retry_interval = std::time::Duration::from_secs_f32(0.3); + loop { + if start_time.elapsed() > timeout_duration { + error!( + "timeout exceeded: failed to poll postgres and create database within 10 minutes" + ); + std::process::exit(1); + } + match tokio_postgres::connect(&restore_pg_connstring, tokio_postgres::NoTls).await { Ok((client, connection)) => { // Spawn the connection handling task to maintain the connection @@ -247,16 +258,24 @@ pub(crate) async fn main() -> anyhow::Result<()> { break; } Err(e) => { - warn!("failed to create database: {}", e); + warn!( + "failed to create database: {}, retying in {}s", + e, + retry_interval.as_secs_f32() + ); + tokio::time::sleep(retry_interval).await; continue; } } } Err(_) => { - info!("postgres not ready yet, retrying in 0.3s"); - tokio::time::sleep(std::time::Duration::from_secs_f32(0.3)).await; + info!(format!( + "postgres not ready yet, retrying in {}s", + retry_interval.as_secs_f32() + )); + tokio::time::sleep(retry_interval).await; continue; - }, + } } } From ebc4735bf48a9b5bcdab03bc2b88a75799599b3f Mon Sep 17 00:00:00 2001 From: Gleb Novikov Date: Wed, 15 Jan 2025 17:45:21 +0000 Subject: [PATCH 3/3] postgres waiting timeout & retry as constants --- compute_tools/src/bin/fast_import.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/compute_tools/src/bin/fast_import.rs b/compute_tools/src/bin/fast_import.rs index c2de4c3c0a..5b008f8182 100644 --- a/compute_tools/src/bin/fast_import.rs +++ b/compute_tools/src/bin/fast_import.rs @@ -41,6 +41,9 @@ mod child_stdio_to_log; #[path = "fast_import/s3_uri.rs"] mod s3_uri; +const PG_WAIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); +const PG_WAIT_RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(300); + #[derive(clap::Parser)] struct Args { #[clap(long)] @@ -231,12 +234,10 @@ pub(crate) async fn main() -> anyhow::Result<()> { let restore_pg_connstring = format!("host=localhost port=5432 user={superuser} dbname=postgres"); - let timeout_duration = std::time::Duration::from_secs(600); // 10 minutes let start_time = std::time::Instant::now(); - let retry_interval = std::time::Duration::from_secs_f32(0.3); loop { - if start_time.elapsed() > timeout_duration { + if start_time.elapsed() > PG_WAIT_TIMEOUT { error!( "timeout exceeded: failed to poll postgres and create database within 10 minutes" ); @@ -261,19 +262,19 @@ pub(crate) async fn main() -> anyhow::Result<()> { warn!( "failed to create database: {}, retying in {}s", e, - retry_interval.as_secs_f32() + PG_WAIT_RETRY_INTERVAL.as_secs_f32() ); - tokio::time::sleep(retry_interval).await; + tokio::time::sleep(PG_WAIT_RETRY_INTERVAL).await; continue; } } } Err(_) => { - info!(format!( + info!( "postgres not ready yet, retrying in {}s", - retry_interval.as_secs_f32() - )); - tokio::time::sleep(retry_interval).await; + PG_WAIT_RETRY_INTERVAL.as_secs_f32() + ); + tokio::time::sleep(PG_WAIT_RETRY_INTERVAL).await; continue; } }