From 056be7a403cb5957e97694a93738df64778410c4 Mon Sep 17 00:00:00 2001 From: Joost Jager Date: Tue, 21 Jul 2026 12:47:46 +0200 Subject: [PATCH] Add PostgreSQL node failover leases Make PostgresStore acquire a table-scoped lease during construction and fence every mutation with it. Treat runtime lease loss as process-fatal so servers can exit without final persistence and restart from durable state. --- src/builder.rs | 50 +++- src/io/mod.rs | 2 + src/io/node_lease.rs | 138 ++++++++++ src/io/postgres_store/mod.rs | 391 +++++++++++++++++++++++++-- src/io/postgres_store/pool.rs | 4 +- src/lib.rs | 43 +++ src/runtime.rs | 9 + tests/integration_tests_migration.rs | 1 + tests/integration_tests_postgres.rs | 6 + 9 files changed, 609 insertions(+), 35 deletions(-) create mode 100644 src/io/node_lease.rs diff --git a/src/builder.rs b/src/builder.rs index a70b04b2ab..5261059b0d 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -190,6 +190,8 @@ pub enum BuildError { /// /// [`KVStore`]: lightning::util::persist::KVStore KVStoreSetupFailed, + /// Another node currently owns the PostgreSQL node lease. + NodeLeaseUnavailable, /// We failed to setup the onchain wallet. WalletSetupFailed, /// We failed to setup the logger. @@ -232,6 +234,7 @@ impl fmt::Display for BuildError { Self::WriteFailed => write!(f, "Failed to write to store."), Self::StoragePathAccessFailed => write!(f, "Failed to access the given storage path."), Self::KVStoreSetupFailed => write!(f, "Failed to setup KVStore."), + Self::NodeLeaseUnavailable => write!(f, "The PostgreSQL node lease is unavailable."), Self::WalletSetupFailed => write!(f, "Failed to setup onchain wallet."), Self::LoggerSetupFailed => write!(f, "Failed to setup the logger."), Self::ChainSourceSetupFailed => write!(f, "Failed to setup the chain source."), @@ -683,6 +686,9 @@ impl NodeBuilder { /// Builds a [`Node`] instance with a [PostgreSQL] backend and according to the options /// previously configured. /// + /// This acquires an exclusive lease for the selected KV table before reading persisted node + /// state. Nodes may share a database when each node identity uses a distinct `kv_table_name`. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -701,6 +707,9 @@ impl NodeBuilder { /// certificates (it does not replace them). If `certificate_pem` is `None`, connections /// will be unencrypted. /// + /// Returns [`BuildError::NodeLeaseUnavailable`] while another process owns the selected KV + /// table's lease. + /// /// [PostgreSQL]: https://www.postgresql.org #[cfg(feature = "postgres")] pub fn build_with_postgres_store( @@ -709,19 +718,29 @@ impl NodeBuilder { ) -> Result { let logger = setup_logger(&self.log_writer_config, &self.config)?; let runtime = self.setup_runtime(&logger)?; - let kv_store = runtime - .block_on(io::postgres_store::PostgresStore::new_with_logger( - connection_string, - db_name, - kv_table_name, - certificate_pem, - Some(Arc::clone(&logger)), - )) - .map_err(|e| { + let kv_store = match runtime.block_on(io::postgres_store::PostgresStore::new_with_logger( + connection_string, + db_name, + kv_table_name, + certificate_pem, + Some(Arc::clone(&logger)), + )) { + Ok(kv_store) => kv_store, + Err(e) if e.kind() == lightning::io::ErrorKind::WouldBlock => { + return Err(BuildError::NodeLeaseUnavailable); + }, + Err(e) => { log_error!(logger, "Failed to set up Postgres store: {e}"); - BuildError::KVStoreSetupFailed - })?; - self.build_with_store_runtime_and_logger(node_entropy, kv_store, runtime, logger) + return Err(BuildError::KVStoreSetupFailed); + }, + }; + let node_lease = kv_store.node_lease(); + let mut node = + self.build_with_store_runtime_and_logger(node_entropy, kv_store, runtime, logger)?; + if !node.install_node_lease(node_lease) { + return Err(BuildError::NodeLeaseUnavailable); + } + Ok(node) } /// Builds a [`Node`] instance with a [`FilesystemStoreV2`] backend and according to the options @@ -1217,6 +1236,9 @@ impl ArcedNodeBuilder { /// Builds a [`Node`] instance with a [PostgreSQL] backend and according to the options /// previously configured. /// + /// This acquires an exclusive lease for the selected KV table before reading persisted node + /// state. Nodes may share a database when each node identity uses a distinct `kv_table_name`. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -1235,6 +1257,9 @@ impl ArcedNodeBuilder { /// certificates (it does not replace them). If `certificate_pem` is `None`, connections /// will be unencrypted. /// + /// Returns [`BuildError::NodeLeaseUnavailable`] while another process owns the selected KV + /// table's lease. + /// /// [PostgreSQL]: https://www.postgresql.org #[cfg(feature = "postgres")] pub fn build_with_postgres_store( @@ -2335,6 +2360,7 @@ fn build_with_store_internal( payment_store, lnurl_auth, is_running, + node_lease: None, node_metrics, om_mailbox, async_payments_role, diff --git a/src/io/mod.rs b/src/io/mod.rs index a01aa59a83..ab6c5ba65d 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -7,6 +7,8 @@ //! Objects and traits for data persistence. +#[cfg_attr(not(feature = "postgres"), allow(dead_code))] +pub(crate) mod node_lease; #[cfg(feature = "postgres")] pub mod postgres_store; pub mod sqlite_store; diff --git a/src/io/node_lease.rs b/src/io/node_lease.rs new file mode 100644 index 0000000000..be35e711e6 --- /dev/null +++ b/src/io/node_lease.rs @@ -0,0 +1,138 @@ +// This file is Copyright its original authors, visible in version control history. +// +// This file is licensed under the Apache License, Version 2.0 or the MIT license , at your option. You may not use this file except in +// accordance with one or both of these licenses. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use lightning::io; + +pub(crate) const NODE_LEASE_DURATION: Duration = Duration::from_secs(30); +// Fail closed before the database lease expires, leaving time for process termination. +pub(crate) const NODE_LEASE_RENEWAL_DEADLINE: Duration = Duration::from_secs(20); +pub(crate) const NODE_LEASE_RENEWAL_INTERVAL: Duration = Duration::from_secs(10); +pub(crate) const NODE_LEASE_RETRY_INTERVAL: Duration = Duration::from_secs(1); +pub(crate) const NODE_LEASE_RELEASE_TIMEOUT: Duration = Duration::from_secs(5); + +type LeaseLossHandler = Box; + +pub(crate) struct NodeLease { + owner_id: [u8; 32], + lease_lost: AtomicBool, + last_confirmed_renewal: Mutex, + loss_sender: tokio::sync::watch::Sender, + loss_handler: Mutex>, +} + +impl NodeLease { + pub(crate) fn new() -> io::Result> { + let mut owner_id = [0u8; 32]; + getrandom::fill(&mut owner_id).map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to generate lease owner ID: {e}")) + })?; + let (loss_sender, _) = tokio::sync::watch::channel(false); + Ok(Arc::new(Self { + owner_id, + lease_lost: AtomicBool::new(false), + last_confirmed_renewal: Mutex::new(Instant::now()), + loss_sender, + loss_handler: Mutex::new(None), + })) + } + + pub(crate) fn owner_id(&self) -> &[u8; 32] { + &self.owner_id + } + + pub(crate) fn is_lost(&self) -> bool { + self.lease_lost.load(Ordering::Acquire) + } + + pub(crate) fn record_renewal(&self) { + if !self.is_lost() { + *self.last_confirmed_renewal.lock().expect("lock") = Instant::now(); + } + } + + pub(crate) fn renewal_deadline_elapsed(&self) -> bool { + self.last_confirmed_renewal.lock().expect("lock").elapsed() >= NODE_LEASE_RENEWAL_DEADLINE + } + + pub(crate) fn ensure_operation_active(&self) -> io::Result<()> { + if self.is_lost() || self.renewal_deadline_elapsed() { + self.mark_lost(); + Err(lease_lost_error()) + } else { + Ok(()) + } + } + + pub(crate) fn map_operation_error(&self, error: io::Error) -> io::Error { + // Preserve transient database errors until they outlive the local safety margin. + self.ensure_operation_active().err().unwrap_or(error) + } + + pub(crate) fn mark_lost(&self) { + if self + .lease_lost + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + + // Run any installed containment handler before publishing lease loss. + if let Some(handler) = self.loss_handler.lock().expect("lock").take() { + handler(); + } + self.loss_sender.send_replace(true); + } + + pub(crate) fn set_loss_handler(&self, handler: LeaseLossHandler) { + let mut locked_handler = self.loss_handler.lock().expect("lock"); + if self.is_lost() { + drop(locked_handler); + handler(); + } else { + *locked_handler = Some(handler); + } + } + + pub(crate) async fn wait_for_loss(self: Arc) { + let mut receiver = self.loss_sender.subscribe(); + let _ = receiver.wait_for(|lost| *lost).await; + } +} + +pub(crate) fn lease_lost_error() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "PostgreSQL node lease was lost") +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + + use super::*; + + #[test] + fn expired_operation_marks_loss_before_returning_error() { + let lease = NodeLease::new().unwrap(); + let handler_ran = Arc::new(AtomicBool::new(false)); + let handler_ran_ref = Arc::clone(&handler_ran); + lease.set_loss_handler(Box::new(move || { + handler_ran_ref.store(true, Ordering::Release); + })); + *lease.last_confirmed_renewal.lock().unwrap() = + Instant::now() - NODE_LEASE_RENEWAL_DEADLINE; + + let error = lease.map_operation_error(io::Error::from(io::ErrorKind::Other)); + + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert!(lease.is_lost()); + assert!(handler_ran.load(Ordering::Acquire)); + } +} diff --git a/src/io/postgres_store/mod.rs b/src/io/postgres_store/mod.rs index 90b8cdc391..704e7335b0 100644 --- a/src/io/postgres_store/mod.rs +++ b/src/io/postgres_store/mod.rs @@ -19,11 +19,16 @@ use lightning_types::string::PrintableString; use native_tls::TlsConnector; use postgres_native_tls::MakeTlsConnector; use tokio_postgres::config::SslMode; +use tokio_postgres::types::ToSql; use tokio_postgres::{Config, Error as PgError}; use self::pool::{make_config_connection, ClientConnection, PgTlsConnector, SmallPool}; +use crate::io::node_lease::{ + lease_lost_error, NodeLease, NODE_LEASE_DURATION, NODE_LEASE_RELEASE_TIMEOUT, + NODE_LEASE_RENEWAL_INTERVAL, NODE_LEASE_RETRY_INTERVAL, +}; use crate::io::utils::check_namespace_key_validity; -use crate::logger::{log_debug, log_info, LdkLogger, Logger}; +use crate::logger::{log_debug, log_error, log_info, LdkLogger, Logger}; use crate::runtime::StoreRuntime; mod migrations; @@ -44,6 +49,11 @@ const PAGE_SIZE: usize = 50; // Keep this small while still allowing progress if one runtime worker blocks on sync store access. const INTERNAL_RUNTIME_WORKERS: usize = 2; +const NODE_LEASE_TABLE_SUFFIX: &str = "_node_lease"; +const POSTGRES_IDENTIFIER_MAX_BYTES: usize = 63; +const MAX_KV_TABLE_NAME_BYTES: usize = + POSTGRES_IDENTIFIER_MAX_BYTES - NODE_LEASE_TABLE_SUFFIX.len(); + fn sql_identifier(identifier: &str) -> io::Result { if identifier.is_empty() || identifier.contains('\0') { return Err(io::Error::new( @@ -70,6 +80,21 @@ fn sql_table_identifier(table_name: &str) -> io::Result { Ok(quoted_parts?.join(".")) } +fn sql_node_lease_table_identifier(table_name: &str) -> io::Result { + sql_table_identifier(table_name)?; + let table_part = table_name.rsplit_once('.').map_or(table_name, |(_, table)| table); + if table_part.len() > MAX_KV_TABLE_NAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "PostgreSQL KV table name exceeds the maximum of {MAX_KV_TABLE_NAME_BYTES} bytes: {table_name}" + ), + )); + } + + sql_table_identifier(&format!("{table_name}{NODE_LEASE_TABLE_SUFFIX}")) +} + /// Runs a tokio-postgres query and, if the connection dropped mid-flight, reconnects and retries /// once. `$store` is the [`PostgresStoreInner`], `$locked` the held client slot guard, /// `$err_map` an `Fn(PgError) -> io::Error` (called at most once), and `$query` an expression @@ -95,6 +120,7 @@ macro_rules! query_with_retry { /// A [`KVStore`] implementation that writes to and reads from a [PostgreSQL] database. /// /// Maintains an internal runtime for the underlying tokio-postgres connection drivers. +/// Each instance exclusively leases its configured KV table and fences every mutation. /// /// [PostgreSQL]: https://www.postgresql.org pub struct PostgresStore { @@ -106,6 +132,9 @@ pub struct PostgresStore { // A store-internal runtime that drives PostgreSQL I/O independently from the node runtime. internal_runtime: Option>, + + node_lease: Arc, + lease_renewal_task: Option>, } // tokio::sync::Mutex (used for the DB client) contains UnsafeCell which opts out of @@ -128,11 +157,17 @@ impl PostgresStore { /// the default `postgres` database to create it. /// /// The given `kv_table_name` will be used or default to [`DEFAULT_KV_TABLE_NAME`]. + /// A companion lease table is created by appending `_node_lease` to this name. /// /// If `certificate_pem` is `Some`, TLS will be used for database connections and the /// provided PEM-encoded CA certificate will be added to the system's default root /// certificates (it does not replace them). If `certificate_pem` is `None`, connections /// will be unencrypted. + /// + /// Construction acquires an exclusive lease for the selected KV table. Returns an error with + /// [`io::ErrorKind::WouldBlock`] while another store owns the lease. + /// Use [`crate::Builder::build_with_postgres_store`] when building a node so ldk-node can install + /// its lease-loss handler before returning it. pub async fn new( connection_string: String, db_name: Option, kv_table_name: Option, certificate_pem: Option, @@ -154,15 +189,63 @@ impl PostgresStore { "PostgreSQL", )?); let tls = Self::build_tls_connector(certificate_pem)?; + let node_lease = NodeLease::new()?; + let lease_ref = Arc::clone(&node_lease); let task = internal_runtime.spawn(async move { - PostgresStoreInner::new(connection_string, db_name, kv_table_name, tls, logger).await + let inner = + PostgresStoreInner::new(connection_string, db_name, kv_table_name, tls, logger) + .await?; + inner.acquire_node_lease(&lease_ref).await?; + Ok::<_, io::Error>(inner) }); let inner = task.await.map_err(|e| { io::Error::new(io::ErrorKind::Other, format!("PostgreSQL runtime task failed: {}", e)) })??; let inner = Arc::new(inner); - let next_write_version = AtomicU64::new(1); - Ok(Self { inner, next_write_version, internal_runtime: Some(internal_runtime) }) + node_lease.record_renewal(); + + let inner_ref = Arc::clone(&inner); + let lease_ref = Arc::clone(&node_lease); + let lease_renewal_task = internal_runtime.spawn(async move { + let mut next_delay = NODE_LEASE_RENEWAL_INTERVAL; + loop { + tokio::time::sleep(next_delay).await; + + if lease_ref.renewal_deadline_elapsed() { + lease_ref.mark_lost(); + return; + } + + match inner_ref.renew_node_lease(&lease_ref).await { + Ok(true) => { + lease_ref.record_renewal(); + next_delay = NODE_LEASE_RENEWAL_INTERVAL; + }, + Ok(false) => { + lease_ref.mark_lost(); + return; + }, + Err(e) => { + if let Some(logger) = inner_ref.logger.as_ref() { + log_error!(logger, "Failed to renew PostgreSQL node lease: {e}"); + } + if lease_ref.renewal_deadline_elapsed() { + lease_ref.mark_lost(); + return; + } + next_delay = NODE_LEASE_RETRY_INTERVAL; + }, + } + } + }); + + Ok(Self { + inner, + next_write_version: AtomicU64::new(1), + internal_runtime: Some(internal_runtime), + node_lease, + lease_renewal_task: Some(lease_renewal_task), + }) } fn build_tls_connector(certificate_pem: Option) -> io::Result { @@ -209,10 +292,42 @@ impl PostgresStore { fn internal_runtime(&self) -> Arc { Arc::clone(self.internal_runtime.as_ref().expect("PostgreSQL runtime must be available")) } + + pub(crate) fn node_lease(&self) -> Arc { + Arc::clone(&self.node_lease) + } } impl Drop for PostgresStore { fn drop(&mut self) { + if let Some(internal_runtime) = self.internal_runtime.as_ref() { + let renewal_task = self.lease_renewal_task.take(); + if let Some(task) = renewal_task.as_ref() { + task.abort(); + } + + let runtime_handle = internal_runtime.handle().clone(); + let inner = Arc::clone(&self.inner); + let node_lease = Arc::clone(&self.node_lease); + let _ = std::thread::spawn(move || { + runtime_handle.block_on(async move { + if let Some(task) = renewal_task { + let _ = task.await; + } + + // Never run clean-release I/O after the terminal loss path has begun. + if !node_lease.is_lost() { + let _ = tokio::time::timeout( + NODE_LEASE_RELEASE_TIMEOUT, + inner.release_node_lease(&node_lease), + ) + .await; + } + }); + }) + .join(); + } + if let Some(internal_runtime) = self.internal_runtime.take() { if let Ok(internal_runtime) = Arc::try_unwrap(internal_runtime) { internal_runtime.shutdown_background(); @@ -252,6 +367,7 @@ impl KVStore for PostgresStore { let secondary_namespace = secondary_namespace.to_string(); let key = key.to_string(); let inner = Arc::clone(&self.inner); + let node_lease = Arc::clone(&self.node_lease); let runtime = self.internal_runtime(); async move { let task = runtime.spawn(async move { @@ -264,6 +380,7 @@ impl KVStore for PostgresStore { &secondary_namespace, &key, buf, + node_lease, ) .await }); @@ -285,6 +402,7 @@ impl KVStore for PostgresStore { let secondary_namespace = secondary_namespace.to_string(); let key = key.to_string(); let inner = Arc::clone(&self.inner); + let node_lease = Arc::clone(&self.node_lease); let runtime = self.internal_runtime(); async move { let task = runtime.spawn(async move { @@ -296,6 +414,7 @@ impl KVStore for PostgresStore { &primary_namespace, &secondary_namespace, &key, + node_lease, ) .await }); @@ -375,6 +494,7 @@ struct PostgresStoreInner { pool: SmallPool, config: Config, kv_table_name_sql: String, + node_lease_table_name_sql: String, tls: PgTlsConnector, write_version_locks: Mutex>>>, logger: Option>, @@ -387,6 +507,7 @@ impl PostgresStoreInner { ) -> io::Result { let kv_table_name = kv_table_name.unwrap_or(DEFAULT_KV_TABLE_NAME.to_string()); let kv_table_name_sql = sql_table_identifier(&kv_table_name)?; + let node_lease_table_name_sql = sql_node_lease_table_identifier(&kv_table_name)?; let mut config: Config = connection_string.parse().map_err(|e: PgError| { let msg = format!("Failed to parse PostgreSQL connection string: {e}"); @@ -446,6 +567,17 @@ impl PostgresStoreInner { io::Error::new(io::ErrorKind::Other, msg) })?; + let sql = format!( + "CREATE TABLE IF NOT EXISTS {node_lease_table_name_sql} ( + id SMALLINT PRIMARY KEY CHECK (id = 1), + owner_id BYTEA NOT NULL, + expires_at TIMESTAMPTZ NOT NULL + )" + ); + client.execute(&sql, &[]).await.map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to create node lease table: {e}")) + })?; + // Read the schema version from the table comment (analogous to SQLite's PRAGMA user_version). let row = client .query_one("SELECT obj_description(to_regclass($1), 'pg_class')", &[&kv_table_name_sql]) @@ -507,7 +639,15 @@ impl PostgresStoreInner { let pool = SmallPool::new(&config, &tls).await?; let write_version_locks = Mutex::new(HashMap::new()); - Ok(Self { pool, config, kv_table_name_sql, tls, write_version_locks, logger }) + Ok(Self { + pool, + config, + kv_table_name_sql, + node_lease_table_name_sql, + tls, + write_version_locks, + logger, + }) } async fn create_database_if_not_exists( @@ -592,6 +732,142 @@ impl PostgresStoreInner { self.pool.get(&self.config, &self.tls, self.logger.as_deref()).await } + async fn acquire_node_lease(&self, node_lease: &NodeLease) -> io::Result<()> { + let locked = self.locked_client().await?; + let lease_table = &self.node_lease_table_name_sql; + let lease_duration_secs = NODE_LEASE_DURATION.as_secs() as i64; + // A single conditional upsert lets PostgreSQL arbitrate contenders using its own clock. + let acquire_sql = format!( + "INSERT INTO {lease_table} (id, owner_id, expires_at) + VALUES (1, $1, clock_timestamp() + ($2::bigint * interval '1 second')) + ON CONFLICT (id) DO UPDATE SET + owner_id = EXCLUDED.owner_id, + expires_at = EXCLUDED.expires_at + WHERE {lease_table}.expires_at <= clock_timestamp() + OR {lease_table}.owner_id = EXCLUDED.owner_id + RETURNING id" + ); + let row = locked + .query_opt(&acquire_sql, &[&node_lease.owner_id().as_slice(), &lease_duration_secs]) + .await + .map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to acquire node lease: {e}")) + })?; + match row { + Some(_) => Ok(()), + None => Err(io::Error::new( + io::ErrorKind::WouldBlock, + "PostgreSQL node lease is unavailable", + )), + } + } + + async fn renew_node_lease(&self, node_lease: &NodeLease) -> io::Result { + if node_lease.is_lost() { + return Ok(false); + } + + let lease_duration_secs = NODE_LEASE_DURATION.as_secs() as i64; + let lease_table = &self.node_lease_table_name_sql; + let sql = format!( + "UPDATE {lease_table} + SET expires_at = clock_timestamp() + ($2::bigint * interval '1 second') + WHERE id = 1 AND owner_id = $1 AND expires_at > clock_timestamp()" + ); + let locked = self.locked_client().await?; + let updated = locked + .execute(&sql, &[&node_lease.owner_id().as_slice(), &lease_duration_secs]) + .await + .map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to renew node lease: {e}")) + })?; + Ok(updated == 1) + } + + async fn release_node_lease(&self, node_lease: &NodeLease) -> io::Result<()> { + let lease_table = &self.node_lease_table_name_sql; + let sql = format!("DELETE FROM {lease_table} WHERE id = 1 AND owner_id = $1"); + let locked = self.locked_client().await?; + locked.execute(&sql, &[&node_lease.owner_id().as_slice()]).await.map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to release node lease: {e}")) + })?; + Ok(()) + } + + async fn renew_node_lease_in_transaction( + &self, transaction: &tokio_postgres::Transaction<'_>, node_lease: &NodeLease, + ) -> io::Result<()> { + node_lease.ensure_operation_active()?; + + // The local check only fails early. This update is authoritative and holds the row lock + // through the caller's KV mutation and commit. + let lease_duration_secs = NODE_LEASE_DURATION.as_secs() as i64; + let lease_table = &self.node_lease_table_name_sql; + let update_sql = format!( + "UPDATE {lease_table} + SET expires_at = clock_timestamp() + ($2::bigint * interval '1 second') + WHERE id = 1 AND owner_id = $1 AND expires_at > clock_timestamp() + RETURNING id" + ); + let row = transaction + .query_opt(&update_sql, &[&node_lease.owner_id().as_slice(), &lease_duration_secs]) + .await + .map_err(|e| { + node_lease.map_operation_error(io::Error::new( + io::ErrorKind::Other, + format!("Failed to check and renew node lease: {e}"), + )) + })?; + if row.is_none() { + node_lease.mark_lost(); + return Err(lease_lost_error()); + } + Ok(()) + } + + async fn execute_fenced_mutation io::Error>( + &self, node_lease: &NodeLease, sql: &str, params: &[&(dyn ToSql + Sync)], err_map: F, + ) -> io::Result<()> { + let mut locked = + self.locked_client().await.map_err(|e| node_lease.map_operation_error(e))?; + let transaction_result = locked.transaction().await; + let reconnect = transaction_result.as_ref().is_err_and(PgError::is_closed); + let transaction_result = if reconnect { + if let (Some(logger), Err(e)) = (self.logger.as_ref(), &transaction_result) { + log_debug!(logger, "Reconnecting to PostgreSQL after error: {e}"); + } + drop(transaction_result); + *locked = make_config_connection(&self.config, &self.tls) + .await + .map_err(|e| node_lease.map_operation_error(e))?; + locked.transaction().await + } else { + transaction_result + }; + let transaction = transaction_result.map_err(|e| { + node_lease.map_operation_error(io::Error::new( + io::ErrorKind::Other, + format!("Failed to start fenced mutation transaction: {e}"), + )) + })?; + self.renew_node_lease_in_transaction(&transaction, node_lease).await?; + transaction + .execute(sql, params) + .await + .map_err(|e| node_lease.map_operation_error(err_map(e)))?; + transaction.commit().await.map_err(|e| { + node_lease.map_operation_error(io::Error::new( + io::ErrorKind::Other, + format!("Failed to commit fenced mutation transaction: {e}"), + )) + })?; + if node_lease.is_lost() { + return Err(lease_lost_error()); + } + node_lease.record_renewal(); + Ok(()) + } + fn get_inner_lock_ref(&self, locking_key: String) -> Arc> { let mut outer_lock = self.write_version_locks.lock().unwrap(); Arc::clone(&outer_lock.entry(locking_key).or_default()) @@ -646,6 +922,7 @@ impl PostgresStoreInner { async fn write_internal( &self, inner_lock_ref: Arc>, locking_key: String, version: u64, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec, + node_lease: Arc, ) -> io::Result<()> { check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "write")?; @@ -668,24 +945,20 @@ impl PostgresStoreInner { io::Error::new(io::ErrorKind::Other, msg) }; - let mut locked = self.locked_client().await?; - query_with_retry!( - self, - locked, + self.execute_fenced_mutation( + &node_lease, + sql.as_str(), + &[&primary_namespace, &secondary_namespace, &key, &buf], err_map, - locked.execute( - sql.as_str(), - &[&primary_namespace, &secondary_namespace, &key, &buf], - ) ) - .map(|_| ()) + .await }) .await } async fn remove_internal( &self, inner_lock_ref: Arc>, locking_key: String, version: u64, - primary_namespace: &str, secondary_namespace: &str, key: &str, + primary_namespace: &str, secondary_namespace: &str, key: &str, node_lease: Arc, ) -> io::Result<()> { check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "remove")?; @@ -706,14 +979,13 @@ impl PostgresStoreInner { io::Error::new(io::ErrorKind::Other, msg) }; - let mut locked = self.locked_client().await?; - query_with_retry!( - self, - locked, + self.execute_fenced_mutation( + &node_lease, + sql.as_str(), + &[&primary_namespace, &secondary_namespace, &key], err_map, - locked.execute(sql.as_str(), &[&primary_namespace, &secondary_namespace, &key]) ) - .map(|_| ()) + .await }) .await } @@ -913,6 +1185,12 @@ mod tests { let _ = client.execute(&format!("DROP TABLE IF EXISTS {kv_table}"), &[]).await; } + async fn cleanup_lease_table(store: &PostgresStore) { + let lease_table = store.inner.node_lease_table_name_sql.clone(); + let client = store.inner.pool.connections[0].lock().await; + let _ = client.execute(&format!("DROP TABLE IF EXISTS {lease_table}"), &[]).await; + } + #[test] fn test_postgres_identifier_quoting() { assert_eq!(sql_identifier("tenant-1").unwrap(), "\"tenant-1\""); @@ -925,6 +1203,13 @@ mod tests { assert!(sql_identifier("").is_err()); assert!(sql_table_identifier("too.many.parts").is_err()); assert!(sql_table_identifier("schema.").is_err()); + assert_eq!(sql_node_lease_table_identifier("tenant-1").unwrap(), "\"tenant-1_node_lease\""); + assert_eq!( + sql_node_lease_table_identifier("tenant.select").unwrap(), + "\"tenant\".\"select_node_lease\"" + ); + assert!(sql_node_lease_table_identifier(&"a".repeat(MAX_KV_TABLE_NAME_BYTES)).is_ok()); + assert!(sql_node_lease_table_identifier(&"a".repeat(MAX_KV_TABLE_NAME_BYTES + 1)).is_err()); } #[tokio::test(flavor = "multi_thread")] @@ -943,6 +1228,70 @@ mod tests { cleanup_store(&store_1).await; } + #[tokio::test] + async fn node_lease_contends_fences_and_releases() { + let leader = create_test_store("test_pg_lease_shared").await; + let leader_lease = leader.node_lease(); + let contention_error = match PostgresStore::new( + test_connection_string(), + None, + Some("test_pg_lease_shared".to_string()), + None, + ) + .await + { + Ok(_) => panic!("a second store acquired the same node lease"), + Err(e) => e, + }; + assert_eq!(contention_error.kind(), io::ErrorKind::WouldBlock); + + let independent = create_test_store("test_pg_lease_independent").await; + let independent_lease_table = independent.inner.node_lease_table_name_sql.clone(); + + KVStore::write(&leader, "ns", "", "leader", vec![1]).await.unwrap(); + KVStore::write(&independent, "ns", "", "independent", vec![1]).await.unwrap(); + + { + let lease_table = leader.inner.node_lease_table_name_sql.clone(); + let client = leader.inner.pool.connections[0].lock().await; + client + .execute( + &format!( + "UPDATE {lease_table} SET expires_at = clock_timestamp() - interval '1 second' WHERE id = 1" + ), + &[], + ) + .await + .unwrap(); + } + + let fallback = create_test_store("test_pg_lease_shared").await; + let fallback_lease = fallback.node_lease(); + let loss_notification = Arc::clone(&leader_lease).wait_for_loss(); + let write_error = KVStore::write(&leader, "ns", "", "stale", vec![2]).await.unwrap_err(); + assert_eq!(write_error.kind(), io::ErrorKind::PermissionDenied); + loss_notification.await; + assert!(leader_lease.is_lost()); + + let remove_error = KVStore::remove(&leader, "ns", "", "leader", false).await.unwrap_err(); + assert_eq!(remove_error.kind(), io::ErrorKind::PermissionDenied); + KVStore::write(&fallback, "ns", "", "fallback", vec![3]).await.unwrap(); + + cleanup_store(&leader).await; + drop(fallback); + let replacement = create_test_store("test_pg_lease_shared").await; + cleanup_store(&replacement).await; + drop(replacement); + assert!(!fallback_lease.is_lost()); + cleanup_store(&independent).await; + drop(independent); + + cleanup_lease_table(&leader).await; + let client = leader.inner.pool.connections[0].lock().await; + let _ = + client.execute(&format!("DROP TABLE IF EXISTS {independent_lease_table}"), &[]).await; + } + #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_list_all_keys() { let store = create_test_store("test_pg_list_all_keys").await; diff --git a/src/io/postgres_store/pool.rs b/src/io/postgres_store/pool.rs index 0f93f30c65..b8a6aaf8da 100644 --- a/src/io/postgres_store/pool.rs +++ b/src/io/postgres_store/pool.rs @@ -62,8 +62,8 @@ pub(super) const POOL_SIZE: usize = 2; /// A fixed-size pool of mutex-guarded PostgreSQL clients. /// /// `get` awaits on whichever slot lock becomes available first via [`tokio::select!`]. Up to -/// [`POOL_SIZE`] operations can run concurrently against the database. Reconnection only happens -/// when the chosen client's local mpsc sender is already closed — no per-acquire health-check +/// [`POOL_SIZE`] operations can run concurrently against the database. `get` reconnects only when +/// the chosen client's local mpsc sender is already closed, without a per-acquire health-check /// round trip. pub(super) struct SmallPool { pub(super) connections: [tokio::sync::Mutex; POOL_SIZE], diff --git a/src/lib.rs b/src/lib.rs index aa234bb639..aef7007c69 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -110,6 +110,7 @@ mod util; mod wallet; use std::default::Default; +use std::future::Future; use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; #[cfg(cycle_tests)] @@ -145,6 +146,7 @@ use fee_estimator::{ use ffi::*; use gossip::GossipSource; use graph::NetworkGraph; +use io::node_lease::NodeLease; use io::utils::update_and_persist_node_metrics; pub use lightning; use lightning::chain::BlockLocator; @@ -266,6 +268,7 @@ pub struct Node { payment_store: Arc, lnurl_auth: Arc, is_running: Arc>, + node_lease: Option>, node_metrics: Arc, om_mailbox: Option>, async_payments_role: Option, @@ -276,6 +279,35 @@ pub struct Node { } impl Node { + #[cfg(feature = "postgres")] + pub(crate) fn install_node_lease(&mut self, node_lease: Arc) -> bool { + debug_assert!(self.node_lease.is_none()); + if node_lease.is_lost() { + return false; + } + + let lease_stop_sender = self.stop_sender.clone(); + let lease_runtime = Arc::downgrade(&self.runtime); + let lease_peer_manager = Arc::downgrade(&self.peer_manager); + let lease_logger = Arc::downgrade(&self.logger); + // Store fencing is the safety boundary; this only limits activity until process exit. + node_lease.set_loss_handler(Box::new(move || { + if let Some(logger) = lease_logger.upgrade() { + log_error!(logger, "PostgreSQL node lease lost, stopping background processing"); + } + lease_stop_sender.send_replace(()); + if let Some(runtime) = lease_runtime.upgrade() { + runtime.abort_background_processor_task(); + } + if let Some(peer_manager) = lease_peer_manager.upgrade() { + peer_manager.disconnect_all_peers(); + } + })); + + self.node_lease = Some(node_lease); + !self.node_lease.as_ref().unwrap().is_lost() + } + /// Starts the necessary background tasks, such as handling events coming from user input, /// LDK/BDK, and the peer-to-peer network. /// @@ -882,6 +914,17 @@ impl Node { Ok(()) } + /// Returns a future which completes when this node loses its PostgreSQL lease. + /// + /// Returns [`None`] when the node was not built with the built-in PostgreSQL store. + /// Lease loss is terminal for this [`Node`]. The caller must immediately terminate the process + /// without calling [`Node::stop`] or dropping the node. A new process may reconstruct the node + /// from persistence. + pub fn wait_for_lease_loss(&self) -> Option + Send + 'static> { + let node_lease = self.node_lease.as_ref().map(Arc::clone)?; + Some(async move { node_lease.wait_for_loss().await }) + } + /// Returns the status of the [`Node`]. pub fn status(&self) -> NodeStatus { let is_running = *self.is_running.read().expect("lock"); diff --git a/src/runtime.rs b/src/runtime.rs index 7e29996e62..bcf80fe64a 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -246,6 +246,15 @@ impl Runtime { }; } + #[cfg(feature = "postgres")] + pub(crate) fn abort_background_processor_task(&self) { + if let Some(background_processor_task) = + self.background_processor_task.lock().expect("lock").as_ref() + { + background_processor_task.abort(); + } + } + #[cfg(tokio_unstable)] pub fn log_metrics(&self) { let runtime_handle = self.handle(); diff --git a/tests/integration_tests_migration.rs b/tests/integration_tests_migration.rs index 7e5767dca6..246cb856da 100644 --- a/tests/integration_tests_migration.rs +++ b/tests/integration_tests_migration.rs @@ -29,6 +29,7 @@ use rand::seq::SliceRandom; async fn drop_tables<'a>(table_names: impl IntoIterator) { for table_name in table_names { drop_table(table_name).await; + drop_table(&format!("{table_name}_node_lease")).await; } } diff --git a/tests/integration_tests_postgres.rs b/tests/integration_tests_postgres.rs index 889d681ba4..a0757fce49 100644 --- a/tests/integration_tests_postgres.rs +++ b/tests/integration_tests_postgres.rs @@ -17,7 +17,9 @@ use rand::RngCore; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn channel_full_cycle_with_postgres_store() { drop_table("channel_cycle_a").await; + drop_table("channel_cycle_a_node_lease").await; drop_table("channel_cycle_b").await; + drop_table("channel_cycle_b_node_lease").await; let (bitcoind, electrsd) = common::setup_bitcoind_and_electrsd(); println!("== Node A =="); @@ -64,12 +66,15 @@ async fn channel_full_cycle_with_postgres_store() { .await; drop_table("channel_cycle_a").await; + drop_table("channel_cycle_a_node_lease").await; drop_table("channel_cycle_b").await; + drop_table("channel_cycle_b_node_lease").await; } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn postgres_node_restart() { drop_table("restart_test").await; + drop_table("restart_test_node_lease").await; let (bitcoind, electrsd) = common::setup_bitcoind_and_electrsd(); let esplora_url = format!("http://{}", electrsd.esplora_url.as_ref().unwrap()); @@ -140,4 +145,5 @@ async fn postgres_node_restart() { node.stop().unwrap(); drop_table("restart_test").await; + drop_table("restart_test_node_lease").await; }