diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs index 9295f091d5..e013f269ed 100644 --- a/compute_tools/src/bin/compute_ctl.rs +++ b/compute_tools/src/bin/compute_ctl.rs @@ -45,6 +45,7 @@ use std::{thread, time::Duration}; use anyhow::{Context, Result}; use chrono::Utc; use clap::Arg; +use compute_tools::lsn_lease::launch_lsn_lease_loop_for_static; use signal_hook::consts::{SIGQUIT, SIGTERM}; use signal_hook::{consts::SIGINT, iterator::Signals}; use tracing::{error, info, warn}; @@ -363,6 +364,8 @@ fn wait_spec( state.start_time = now; } + launch_lsn_lease_loop_for_static(&compute); + Ok(WaitSpecResult { compute, http_port, diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 1a8050e5ef..40060f4117 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -10,12 +10,8 @@ use std::sync::atomic::AtomicU32; use std::sync::atomic::Ordering; use std::sync::{Condvar, Mutex, RwLock}; use std::thread; -use std::time; -use std::time::Duration; use std::time::Instant; -use std::time::SystemTime; -use anyhow::bail; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use futures::future::join_all; @@ -23,7 +19,6 @@ use futures::stream::FuturesUnordered; use futures::StreamExt; use nix::unistd::Pid; use postgres::error::SqlState; -use postgres::SimpleQueryMessage; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; use utils::id::{TenantId, TimelineId}; @@ -1089,8 +1084,6 @@ impl ComputeNode { }; info!(?metrics, "compute start finished"); - start_lsn_lease_for_static(&compute_state); - Ok(pg_process) } @@ -1392,91 +1385,3 @@ pub fn forward_termination_signal() { kill(pg_pid, Signal::SIGQUIT).ok(); } } - -// Start a thread that will periodically renew LSN lease for static compute -fn start_lsn_lease_for_static(compute_state: &ComputeState) { - let spec = compute_state.pspec.as_ref().expect("spec must be set"); - - let lsn = match spec.spec.mode { - ComputeMode::Static(lsn) => lsn, - _ => return, - }; - - let conn_strings = spec.pageserver_connstr.split(','); - - let configs = conn_strings - .map(|connstr| { - let mut config = postgres::Config::from_str(connstr).expect("invalid connstr"); - if let Some(storage_auth_token) = &spec.storage_auth_token { - info!("Got storage auth token from spec file"); - config.password(storage_auth_token.clone()); - } else { - info!("Storage auth token not set"); - } - config - }) - .collect::>(); - - let cmd = format!("lease lsn {} {} {} ", spec.tenant_id, spec.timeline_id, lsn); - - thread::spawn(move || lsn_lease_loop(configs, cmd)); -} - -fn lsn_lease_loop(configs: Vec, cmd: String) { - loop { - match lsn_lease_request(&configs, &cmd) { - Ok(valid_until) => { - let valid_until_duration = Duration::from_millis(valid_until as u64); - - let sleep_duration = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or(Duration::ZERO) - .checked_sub(valid_until_duration) - .unwrap_or(Duration::ZERO) - .checked_sub(Duration::from_secs(60)) - .unwrap_or(Duration::ZERO); - - // Ensure the sleep duration is at least 60 seconds to avoid busy loops - let sleep_duration = std::cmp::max(sleep_duration, Duration::from_secs(60)); - - info!( - "lsn_lease_request succeeded, sleeping for {} seconds", - sleep_duration.as_secs() - ); - thread::sleep(sleep_duration); - } - Err(e) => { - error!("lsn_lease_request failed: {:#}", e); - thread::sleep(time::Duration::from_secs(10)); - } - } - } -} - -fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result { - info!("lsn_lease_request: {}", cmd); - let valid_until = configs - .iter() - .map(|config| { - let mut client = config.connect(NoTls)?; - let res = client.simple_query(cmd)?; - let msg = match res.first() { - Some(msg) => msg, - None => bail!("empty response"), - }; - let row = match msg { - SimpleQueryMessage::Row(row) => row, - _ => bail!("error parsing lsn lease response"), - }; - let valid_until_str = match row.get("valid_until") { - Some(valid_until) => valid_until, - None => bail!("valid_until not found"), - }; - Ok(u128::from_str(valid_until_str)?) - }) - .collect::>>()? - .into_iter() - .min() - .unwrap(); - Ok(valid_until) -} diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs index 18c228ba54..041cb650f4 100644 --- a/compute_tools/src/lib.rs +++ b/compute_tools/src/lib.rs @@ -11,6 +11,7 @@ pub mod logger; pub mod catalog; pub mod compute; pub mod extension_server; +pub mod lsn_lease; pub mod monitor; pub mod params; pub mod pg_helpers; diff --git a/compute_tools/src/lsn_lease.rs b/compute_tools/src/lsn_lease.rs new file mode 100644 index 0000000000..386f842361 --- /dev/null +++ b/compute_tools/src/lsn_lease.rs @@ -0,0 +1,112 @@ +use anyhow::bail; +use anyhow::Result; +use postgres::{NoTls, SimpleQueryMessage}; +use std::{ + str::FromStr, + sync::Arc, + thread, + time::{Duration, SystemTime}, +}; + +use compute_api::spec::ComputeMode; +use tracing::{error, info}; +use utils::lsn::Lsn; + +use crate::compute::{ComputeNode, ComputeState}; + +pub fn launch_lsn_lease_loop_for_static(compute: &Arc) { + let lsn = { + let state = compute.state.lock().unwrap(); + let spec = state.pspec.as_ref().expect("spec must be set"); + match spec.spec.mode { + ComputeMode::Static(lsn) => lsn, + _ => return, + } + }; + let compute = compute.clone(); + thread::spawn(move || lsn_lease_loop(compute, lsn)); +} + +fn postgres_configs_from_state(compute_state: &ComputeState) -> Vec { + let spec = compute_state.pspec.as_ref().expect("spec must be set"); + let conn_strings = spec.pageserver_connstr.split(','); + + conn_strings + .map(|connstr| { + let mut config = postgres::Config::from_str(connstr).expect("invalid connstr"); + if let Some(storage_auth_token) = &spec.storage_auth_token { + info!("Got storage auth token from spec file"); + config.password(storage_auth_token.clone()); + } else { + info!("Storage auth token not set"); + } + config + }) + .collect::>() +} + +fn lsn_lease_loop(compute: Arc, lsn: Lsn) { + loop { + let state = compute.state.lock().unwrap(); + + let spec = state.pspec.as_ref().expect("spec must be set"); + let configs = postgres_configs_from_state(&state); + + let cmd = format!("lease lsn {} {} {} ", spec.tenant_id, spec.timeline_id, lsn); + + match lsn_lease_request(&configs, &cmd) { + Ok(valid_until) => { + let valid_until_duration = Duration::from_millis(valid_until as u64); + + let sleep_duration = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or(Duration::ZERO) + .checked_sub(valid_until_duration) + .unwrap_or(Duration::ZERO) + .checked_sub(Duration::from_secs(60)) + .unwrap_or(Duration::ZERO); + + // Ensure the sleep duration is at least 60 seconds to avoid busy loops + let sleep_duration = std::cmp::max(sleep_duration, Duration::from_secs(60)); + + info!( + "lsn_lease_request succeeded, sleeping for {} seconds", + sleep_duration.as_secs() + ); + thread::sleep(sleep_duration); + } + Err(e) => { + error!("lsn_lease_request failed: {:#}", e); + thread::sleep(Duration::from_secs(10)); + } + } + } +} + +fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result { + info!("lsn_lease_request: {}", cmd); + let valid_until = configs + .iter() + .map(|config| { + let mut client = config.connect(NoTls)?; + let res = client.simple_query(cmd)?; + let msg = match res.first() { + Some(msg) => msg, + None => bail!("empty response"), + }; + let row = match msg { + SimpleQueryMessage::Row(row) => row, + _ => bail!("error parsing lsn lease response"), + }; + let valid_until_str = match row.get("valid_until") { + Some(valid_until) => valid_until, + None => bail!("valid_until not found"), + }; + Ok(u128::from_str(valid_until_str)?) + }) + .collect::>>()? + .into_iter() + .min() + .unwrap(); + Ok(valid_until) +}