From a964bfcc2fff808a2ee82e62f51b6fedee8227c9 Mon Sep 17 00:00:00 2001 From: Joost Jager Date: Tue, 21 Jul 2026 12:47:46 +0200 Subject: [PATCH] Replace PostgreSQL advisory locking with fenced leases Acquire a table-scoped lease before loading persisted state and validate and renew it atomically with each mutation. Renew idle leases in the background and permanently invalidate the store on detected lease loss, preserving panic behavior and the existing lock-contention error kind. Coordinate the one-way schema-v1 upgrade with the legacy advisory lock and commit the schema-v2 marker with the initial lease. Older releases then reject the upgraded store instead of bypassing lease fencing. Keep lease handling inside PostgresStore, with node lifecycle and recovery unchanged. Cover contention, migration, queued-write fencing, lease release, and panic propagation in the store tests. --- src/builder.rs | 30 +- src/io/mod.rs | 2 + src/io/node_lease.rs | 129 +++++++ src/io/postgres_store/migrations.rs | 33 +- src/io/postgres_store/mod.rs | 546 +++++++++++++++++++++++---- tests/integration_tests_migration.rs | 1 + tests/integration_tests_postgres.rs | 6 + 7 files changed, 649 insertions(+), 98 deletions(-) create mode 100644 src/io/node_lease.rs diff --git a/src/builder.rs b/src/builder.rs index 1158044e47..bb909142da 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -722,6 +722,10 @@ 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`. + /// The store panics on detected lease loss; node recovery is not handled automatically. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -735,13 +739,10 @@ impl NodeBuilder { /// The given `kv_table_name` will be used or default to /// [`DEFAULT_KV_TABLE_NAME`](io::postgres_store::DEFAULT_KV_TABLE_NAME). /// - /// # Warning - /// - /// Do not point multiple [`Node`] instances at the same database and table. Concurrent access is - /// unsafe and can corrupt node state. You must make sure that only one node accesses each - /// database and table. The store uses a PostgreSQL advisory lock to reduce this risk. This lock - /// is only a temporary safeguard and does not make concurrent access safe. - /// Nodes using a different database or table on the same server may coexist. + /// Opening a schema-v1 store upgrades it to the lease-aware schema v2. Stop all processes using + /// the v1 store before upgrading. For the first v2 open, use the same resolved database name and + /// byte-for-byte same `kv_table_name` spelling, including schema qualification, so its transition + /// lock matches v1. Older releases cannot reopen a v2 store, so downgrading is unsupported. /// /// 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 @@ -1319,6 +1320,10 @@ impl Builder { /// 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`. + /// The store panics on detected lease loss; node recovery is not handled automatically. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -1332,13 +1337,10 @@ impl Builder { /// The given `kv_table_name` will be used or default to /// [`DEFAULT_KV_TABLE_NAME`](io::postgres_store::DEFAULT_KV_TABLE_NAME). /// - /// # Warning - /// - /// Do not point multiple [`Node`] instances at the same database and table. Concurrent access is - /// unsafe and can corrupt node state. You must make sure that only one node accesses each - /// database and table. The store uses a PostgreSQL advisory lock to reduce this risk. This lock - /// is only a temporary safeguard and does not make concurrent access safe. - /// Nodes using a different database or table on the same server may coexist. + /// Opening a schema-v1 store upgrades it to the lease-aware schema v2. Stop all processes using + /// the v1 store before upgrading. For the first v2 open, use the same resolved database name and + /// byte-for-byte same `kv_table_name` spelling, including schema qualification, so its transition + /// lock matches v1. Older releases cannot reopen a v2 store, so downgrading is unsupported. /// /// 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 diff --git a/src/io/mod.rs b/src/io/mod.rs index b7e4d2131f..70e400c87d 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -10,6 +10,8 @@ #[cfg(feature = "storage-filesystem")] pub(crate) mod fs_store; #[cfg(feature = "storage-postgres")] +pub(crate) mod node_lease; +#[cfg(feature = "storage-postgres")] pub mod postgres_store; #[cfg(feature = "storage-sqlite")] pub mod sqlite_store; diff --git a/src/io/node_lease.rs b/src/io/node_lease.rs new file mode 100644 index 0000000000..3ea9ce235b --- /dev/null +++ b/src/io/node_lease.rs @@ -0,0 +1,129 @@ +// 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); +// Detect loss before the database lease expires, keeping a margin for delayed renewals. +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); + +pub(crate) struct NodeLease { + owner_id: [u8; 32], + lease_lost: AtomicBool, + last_confirmed_renewal: 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}")) + })?; + Ok(Arc::new(Self { + owner_id, + lease_lost: AtomicBool::new(false), + last_confirmed_renewal: Mutex::new(Instant::now()), + })) + } + + 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_started_at(&self, renewal_started_at: Instant) { + if !self.is_lost() { + let mut last_confirmed_renewal = self.last_confirmed_renewal.lock().expect("lock"); + *last_confirmed_renewal = (*last_confirmed_renewal).max(renewal_started_at); + } + } + + pub(crate) fn renewal_deadline_elapsed(&self) -> bool { + self.last_confirmed_renewal.lock().expect("lock").elapsed() >= NODE_LEASE_RENEWAL_DEADLINE + } + + pub(crate) async fn wait_for_renewal_deadline(&self) { + loop { + let last_confirmed_renewal = *self.last_confirmed_renewal.lock().expect("lock"); + let deadline = last_confirmed_renewal + NODE_LEASE_RENEWAL_DEADLINE; + tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; + if self.renewal_deadline_elapsed() { + return; + } + } + } + + pub(crate) fn assert_active(&self) { + if self.is_lost() || self.renewal_deadline_elapsed() { + self.mark_lost(); + } + } + + pub(crate) fn map_operation_error(&self, error: io::Error) -> io::Error { + // Preserve transient database errors until they outlive the local safety margin. + self.assert_active(); + error + } + + pub(crate) fn mark_lost(&self) -> ! { + self.lease_lost.store(true, Ordering::Release); + panic!("PostgreSQL node lease was lost; continuing may corrupt node state"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn expired_operation_permanently_marks_loss_before_panicking() { + let lease = NodeLease::new().unwrap(); + *lease.last_confirmed_renewal.lock().unwrap() = + Instant::now() - NODE_LEASE_RENEWAL_DEADLINE; + + assert!(std::panic::catch_unwind(|| { + lease.map_operation_error(io::Error::from(io::ErrorKind::Other)); + }) + .is_err()); + assert!(lease.is_lost()); + lease.record_renewal_started_at(Instant::now()); + assert!(std::panic::catch_unwind(|| lease.assert_active()).is_err()); + } + + #[test] + fn confirmed_renewal_uses_attempt_time_and_does_not_regress() { + let lease = NodeLease::new().unwrap(); + let renewal_started_at = Instant::now() - Duration::from_secs(1); + *lease.last_confirmed_renewal.lock().unwrap() = renewal_started_at - Duration::from_secs(1); + + lease.record_renewal_started_at(renewal_started_at); + lease.record_renewal_started_at(renewal_started_at - Duration::from_secs(1)); + + assert_eq!(*lease.last_confirmed_renewal.lock().unwrap(), renewal_started_at); + } + + #[tokio::test] + async fn expired_renewal_deadline_completes_immediately() { + let lease = NodeLease::new().unwrap(); + *lease.last_confirmed_renewal.lock().unwrap() = + Instant::now() - NODE_LEASE_RENEWAL_DEADLINE; + + tokio::time::timeout(Duration::from_secs(1), lease.wait_for_renewal_deadline()) + .await + .unwrap(); + } +} diff --git a/src/io/postgres_store/migrations.rs b/src/io/postgres_store/migrations.rs index c9add1c57c..abb7e243ff 100644 --- a/src/io/postgres_store/migrations.rs +++ b/src/io/postgres_store/migrations.rs @@ -6,16 +6,35 @@ // accordance with one or both of these licenses. use lightning::io; -use tokio_postgres::Client; +use tokio_postgres::Transaction; pub(super) async fn migrate_schema( - _client: &Client, _kv_table_name: &str, from_version: u16, to_version: u16, + transaction: &Transaction<'_>, kv_table_name: &str, mut from_version: u16, to_version: u16, ) -> io::Result<()> { assert!(from_version < to_version); - // Future migrations go here, e.g.: - // if from_version == 1 && to_version >= 2 { - // migrate_v1_to_v2(client, kv_table_name).await?; - // from_version = 2; - // } + if from_version == 1 && to_version >= 2 { + migrate_v1_to_v2(transaction, kv_table_name).await?; + from_version = 2; + } + + if from_version != to_version { + return Err(io::Error::new( + io::ErrorKind::Other, + format!("No PostgreSQL schema migration from version {from_version} to {to_version}"), + )); + } + Ok(()) +} + +async fn migrate_v1_to_v2(transaction: &Transaction<'_>, kv_table_name: &str) -> io::Result<()> { + // Schema v2 marks the transition from the legacy session advisory lock to fenced node leases. + // Older releases reject this version instead of reopening the store without lease fencing. + let sql = format!("COMMENT ON TABLE {kv_table_name} IS '2'"); + transaction.execute(&sql, &[]).await.map_err(|e| { + io::Error::new( + io::ErrorKind::Other, + format!("Failed to set PostgreSQL schema version 2: {e}"), + ) + })?; Ok(()) } diff --git a/src/io/postgres_store/mod.rs b/src/io/postgres_store/mod.rs index 4ffb948ff9..ffba44e4c2 100644 --- a/src/io/postgres_store/mod.rs +++ b/src/io/postgres_store/mod.rs @@ -10,6 +10,7 @@ use std::collections::HashMap; use std::future::Future; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; +use std::time::Instant; use bitcoin::hashes::{sha256, Hash, HashEngine}; use lightning::io; @@ -20,11 +21,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::{ + 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; @@ -37,7 +43,7 @@ pub const DEFAULT_DB_NAME: &str = "ldk_db"; pub const DEFAULT_KV_TABLE_NAME: &str = "ldk_data"; // The current schema version for the PostgreSQL store. -const SCHEMA_VERSION: u16 = 1; +const SCHEMA_VERSION: u16 = 2; // The number of entries returned per page in paginated list operations. const PAGE_SIZE: usize = 50; @@ -57,6 +63,11 @@ fn advisory_lock_id(db_name: &str, kv_table_name: &str) -> i64 { i64::from_be_bytes(hash[..8].try_into().expect("SHA-256 prefix has the expected length")) } +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( @@ -83,8 +94,23 @@ 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 pooled connection dropped mid-flight, reconnects and -/// retries once after asserting that the store's advisory-lock connection is still open. `$store` +/// retries once after asserting that the store's lease is still active. `$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 that yields a fresh /// `Future>` each time it is evaluated. `$query` may be evaluated up to @@ -126,6 +152,9 @@ fn handle_runtime_task_result( /// 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. +/// Lease loss panics in the task that detects it and permanently invalidates the store. +/// Subsequent store operations also panic. This does not provide node shutdown or recovery. /// /// [PostgreSQL]: https://www.postgresql.org pub struct PostgresStore { @@ -137,6 +166,8 @@ pub struct PostgresStore { // A store-internal runtime that drives PostgreSQL I/O independently from the node runtime. internal_runtime: Option>, + + lease_renewal_task: Option>, } // tokio::sync::Mutex (used for the DB client) contains UnsafeCell which opts out of @@ -159,19 +190,20 @@ impl PostgresStore { /// the default `postgres` database to create it. /// /// The given `kv_table_name` will be used or default to [`DEFAULT_KV_TABLE_NAME`]. - /// - /// # Warning - /// - /// Do not point multiple [`PostgresStore`] instances at the same database and table. Concurrent - /// access is unsafe and can corrupt stored data. You must make sure that only one store accesses - /// each database and table. The store uses a PostgreSQL advisory lock to reduce this risk. This - /// lock is only a temporary safeguard and does not make concurrent access safe. - /// Stores using a different database or table on the same PostgreSQL server may coexist. + /// 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::AlreadyExists`] while another store owns the lease. + /// + /// Opening a schema-v1 store upgrades it to the lease-aware schema v2. Stop all processes using + /// the v1 store before upgrading. For the first v2 open, use the same resolved database name and + /// byte-for-byte same `kv_table_name` spelling, including schema qualification, so its transition + /// lock matches v1. Older releases cannot reopen a v2 store, so downgrading is unsupported. pub async fn new( connection_string: String, db_name: Option, kv_table_name: Option, certificate_pem: Option, @@ -200,8 +232,53 @@ impl PostgresStore { 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) }) + inner.node_lease.assert_active(); + + let inner_ref = Arc::clone(&inner); + let lease_ref = Arc::clone(&inner.node_lease); + let lease_renewal_task = internal_runtime.spawn(async move { + let mut next_delay = NODE_LEASE_RENEWAL_INTERVAL; + loop { + let renewal_attempt = async { + tokio::time::sleep(next_delay).await; + let started_at = Instant::now(); + (started_at, inner_ref.renew_node_lease().await) + }; + let (renewal_started_at, renewal_result) = tokio::select! { + biased; + _ = lease_ref.wait_for_renewal_deadline() => { + lease_ref.mark_lost(); + }, + result = renewal_attempt => result, + }; + match renewal_result { + Ok(true) => { + lease_ref.record_renewal_started_at(renewal_started_at); + next_delay = NODE_LEASE_RENEWAL_INTERVAL + .saturating_sub(renewal_started_at.elapsed()); + }, + Ok(false) => { + lease_ref.mark_lost(); + }, + 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(); + } + next_delay = NODE_LEASE_RETRY_INTERVAL; + }, + } + } + }); + + Ok(Self { + inner, + next_write_version: AtomicU64::new(1), + internal_runtime: Some(internal_runtime), + lease_renewal_task: Some(lease_renewal_task), + }) } fn build_tls_connector(certificate_pem: Option) -> io::Result { @@ -252,6 +329,33 @@ impl PostgresStore { 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 _ = 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 !inner.node_lease.is_lost() { + let _ = tokio::time::timeout( + NODE_LEASE_RELEASE_TIMEOUT, + inner.release_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(); @@ -388,11 +492,10 @@ impl MigratableKVStore for PostgresStore { struct PostgresStoreInner { pool: SmallPool, - // PostgreSQL advisory locks are session-scoped, so keep the connection that acquired our lock - // alive for the lifetime of the store. - lock_client: ClientConnection, config: Config, kv_table_name_sql: String, + node_lease_table_name_sql: String, + node_lease: Arc, tls: PgTlsConnector, write_version_locks: Mutex>>>, logger: Option>, @@ -405,6 +508,8 @@ 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 node_lease = NodeLease::new()?; let mut config: Config = connection_string.parse().map_err(|e: PgError| { let msg = format!("Failed to parse PostgreSQL connection string: {e}"); @@ -443,16 +548,30 @@ impl PostgresStoreInner { Self::create_database_if_not_exists(&config, &tls, logger.as_deref()).await?; - let client = make_config_connection(&config, &tls).await?; + let mut client = make_config_connection(&config, &tls).await?; + let pool = SmallPool::new(&config, &tls).await?; + let transaction = client.transaction().await.map_err(|e| { + io::Error::new( + io::ErrorKind::Other, + format!("Failed to start PostgreSQL schema setup transaction: {e}"), + ) + })?; + + // Releases using schema v1 hold this advisory lock for their lifetime. Take its + // transaction-scoped counterpart before changing anything so a v1 process and the v2 lease + // protocol cannot be active at the same time during a coordinated upgrade. let lock_id = advisory_lock_id(&db_name, &kv_table_name); - let row = client.query_one("SELECT pg_try_advisory_lock($1)", &[&lock_id]).await.map_err( - |e| { - let msg = format!( - "Failed to acquire PostgreSQL store lock for database {db_name} and table {kv_table_name}: {e}" - ); - io::Error::new(io::ErrorKind::Other, msg) - }, - )?; + let row = transaction + .query_one("SELECT pg_try_advisory_xact_lock($1)", &[&lock_id]) + .await + .map_err(|e| { + io::Error::new( + io::ErrorKind::Other, + format!( + "Failed to acquire PostgreSQL migration lock for database {db_name} and table {kv_table_name}: {e}" + ), + ) + })?; if !row.get::<_, bool>(0) { return Err(io::Error::new( io::ErrorKind::AlreadyExists, @@ -462,6 +581,42 @@ impl PostgresStoreInner { )); } + 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 + )" + ); + transaction.execute(&sql, &[]).await.map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to create node lease table: {e}")) + })?; + + // Keep the lease row locked until the schema marker and all setup changes are committed. + let lease_duration_secs = NODE_LEASE_DURATION.as_secs() as i64; + let acquire_sql = format!( + "INSERT INTO {node_lease_table_name_sql} (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 {node_lease_table_name_sql}.expires_at <= clock_timestamp() + OR {node_lease_table_name_sql}.owner_id = EXCLUDED.owner_id + RETURNING id" + ); + let row = transaction + .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}")) + })?; + if row.is_none() { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + "PostgreSQL node lease is unavailable", + )); + } + // Create the KV data table if it doesn't exist. `sort_order` uses BIGSERIAL so // the database assigns a fresh, monotonically increasing value on each INSERT and // keeps the previous value untouched on UPSERT-update; the sequence persists across @@ -476,13 +631,13 @@ impl PostgresStoreInner { PRIMARY KEY (primary_namespace, secondary_namespace, key) )" ); - client.execute(sql.as_str(), &[]).await.map_err(|e| { + transaction.execute(sql.as_str(), &[]).await.map_err(|e| { let msg = format!("Failed to create table {kv_table_name}: {e}"); io::Error::new(io::ErrorKind::Other, msg) })?; // Read the schema version from the table comment (analogous to SQLite's PRAGMA user_version). - let row = client + let row = transaction .query_one("SELECT obj_description(to_regclass($1), 'pg_class')", &[&kv_table_name_sql]) .await .map_err(|e| { @@ -513,13 +668,18 @@ impl PostgresStoreInner { if version_res == 0 { // New table, set our SCHEMA_VERSION. let sql = format!("COMMENT ON TABLE {kv_table_name_sql} IS '{SCHEMA_VERSION}'"); - client.execute(sql.as_str(), &[]).await.map_err(|e| { + transaction.execute(sql.as_str(), &[]).await.map_err(|e| { let msg = format!("Failed to set schema version: {e}"); io::Error::new(io::ErrorKind::Other, msg) })?; } else if version_res < SCHEMA_VERSION { - migrations::migrate_schema(&client, &kv_table_name_sql, version_res, SCHEMA_VERSION) - .await?; + migrations::migrate_schema( + &transaction, + &kv_table_name_sql, + version_res, + SCHEMA_VERSION, + ) + .await?; } else if version_res > SCHEMA_VERSION { let msg = format!( "Failed to open database: incompatible schema version {version_res}. Expected: {SCHEMA_VERSION}" @@ -532,19 +692,50 @@ impl PostgresStoreInner { let sql = format!( "CREATE INDEX IF NOT EXISTS {index_name_sql} ON {kv_table_name_sql} (primary_namespace, secondary_namespace, sort_order DESC, key ASC)" ); - client.execute(sql.as_str(), &[]).await.map_err(|e| { + transaction.execute(sql.as_str(), &[]).await.map_err(|e| { let msg = format!("Failed to create index on table {kv_table_name}: {e}"); io::Error::new(io::ErrorKind::Other, msg) })?; - let pool = SmallPool::new(&config, &tls).await?; + let renewal_started_at = Instant::now(); + let renew_sql = format!( + "UPDATE {node_lease_table_name_sql} + SET expires_at = clock_timestamp() + ($2::bigint * interval '1 second') + WHERE id = 1 AND owner_id = $1" + ); + let updated = transaction + .execute(&renew_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 after schema setup: {e}"), + ) + })?; + if updated != 1 { + return Err(io::Error::new( + io::ErrorKind::Other, + "Failed to renew node lease after schema setup", + )); + } + transaction.commit().await.map_err(|e| { + io::Error::new( + io::ErrorKind::Other, + format!("Failed to commit PostgreSQL schema setup transaction: {e}"), + ) + })?; + // Drop the setup client; the pool has its own connections. + drop(client); + + node_lease.record_renewal_started_at(renewal_started_at); let write_version_locks = Mutex::new(HashMap::new()); Ok(Self { pool, - lock_client: client, config, kv_table_name_sql, + node_lease_table_name_sql, + node_lease, tls, write_version_locks, logger, @@ -638,10 +829,111 @@ impl PostgresStoreInner { } fn assert_store_lock(&self) { - assert!( - !self.lock_client.is_closed(), - "PostgreSQL store lock connection closed; continuing may corrupt node state" + self.node_lease.assert_active(); + } + + async fn renew_node_lease(&self) -> io::Result { + if self.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, &[&self.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) -> 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, &[&self.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<'_>, + ) -> io::Result<()> { + self.node_lease.assert_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()" ); + let updated = transaction + .execute(&update_sql, &[&self.node_lease.owner_id().as_slice(), &lease_duration_secs]) + .await + .map_err(|e| { + self.node_lease.map_operation_error(io::Error::new( + io::ErrorKind::Other, + format!("Failed to check and renew node lease: {e}"), + )) + })?; + if updated != 1 { + self.node_lease.mark_lost(); + } + Ok(()) + } + + async fn execute_fenced_mutation io::Error>( + &self, sql: &str, params: &[&(dyn ToSql + Sync)], err_map: F, + ) -> io::Result<()> { + let node_lease = &self.node_lease; + 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}"), + )) + })?; + let renewal_started_at = Instant::now(); + self.renew_node_lease_in_transaction(&transaction).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}"), + )) + })?; + node_lease.assert_active(); + node_lease.record_renewal_started_at(renewal_started_at); + Ok(()) } fn get_inner_lock_ref(&self, locking_key: String) -> Arc> { @@ -720,17 +1012,12 @@ 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( + 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 } @@ -758,14 +1045,12 @@ 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( + sql.as_str(), + &[&primary_namespace, &secondary_namespace, &key], err_map, - locked.execute(sql.as_str(), &[&primary_namespace, &secondary_namespace, &key]) ) - .map(|_| ()) + .await }) .await } @@ -977,6 +1262,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()); } #[test] @@ -988,8 +1280,8 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn test_postgres_store_advisory_lock() { - let table_name = "test_pg_advisory_lock"; + async fn test_postgres_store_lease() { + let table_name = "test_pg_lease"; let store = create_test_store(table_name).await; let err = @@ -1002,6 +1294,101 @@ mod tests { cleanup_store(&store).await; } + #[tokio::test(flavor = "multi_thread")] + async fn test_postgres_store_migrates_advisory_lock_to_lease() { + let table_name = "test_pg_advisory_lock_migration"; + let kv_table_name_sql = sql_table_identifier(table_name).unwrap(); + let node_lease_table_name_sql = sql_node_lease_table_identifier(table_name).unwrap(); + let mut config: Config = test_connection_string().parse().unwrap(); + let db_name = config + .get_dbname() + .map(ToOwned::to_owned) + .unwrap_or_else(|| DEFAULT_DB_NAME.to_string()); + config.dbname(&db_name); + let client = make_config_connection(&config, &PgTlsConnector::Plain).await.unwrap(); + client + .execute(&format!("DROP TABLE IF EXISTS {node_lease_table_name_sql}"), &[]) + .await + .unwrap(); + client.execute(&format!("DROP TABLE IF EXISTS {kv_table_name_sql}"), &[]).await.unwrap(); + client + .execute( + &format!( + "CREATE TABLE {kv_table_name_sql} ( + primary_namespace TEXT NOT NULL, + secondary_namespace TEXT NOT NULL DEFAULT '', + key TEXT NOT NULL CHECK (key <> ''), + value BYTEA, + sort_order BIGSERIAL CHECK (sort_order >= 0), + PRIMARY KEY (primary_namespace, secondary_namespace, key) + )" + ), + &[], + ) + .await + .unwrap(); + client.execute(&format!("COMMENT ON TABLE {kv_table_name_sql} IS '1'"), &[]).await.unwrap(); + client + .execute( + &format!( + "INSERT INTO {kv_table_name_sql} \ + (primary_namespace, secondary_namespace, key, value) VALUES ($1, $2, $3, $4)" + ), + &[&"ns", &"", &"preserved", &vec![42u8]], + ) + .await + .unwrap(); + + let lock_id = advisory_lock_id(&db_name, table_name); + let row = client.query_one("SELECT pg_try_advisory_lock($1)", &[&lock_id]).await.unwrap(); + assert!(row.get::<_, bool>(0)); + + let error = match PostgresStore::new( + test_connection_string(), + None, + Some(table_name.to_string()), + None, + ) + .await + { + Ok(_) => panic!("a schema-v1 advisory lock must block migration"), + Err(error) => error, + }; + assert_eq!(error.kind(), io::ErrorKind::AlreadyExists); + let row = client + .query_one("SELECT obj_description(to_regclass($1), 'pg_class')", &[&kv_table_name_sql]) + .await + .unwrap(); + assert_eq!(row.get::<_, Option<&str>>(0), Some("1")); + let row = client + .query_one("SELECT to_regclass($1)::text", &[&node_lease_table_name_sql]) + .await + .unwrap(); + assert_eq!(row.get::<_, Option>(0), None); + + let row = client.query_one("SELECT pg_advisory_unlock($1)", &[&lock_id]).await.unwrap(); + assert!(row.get::<_, bool>(0)); + let store = create_test_store(table_name).await; + assert_eq!(KVStore::read(&store, "ns", "", "preserved").await.unwrap(), vec![42u8]); + + let row = client + .query_one("SELECT obj_description(to_regclass($1), 'pg_class')", &[&kv_table_name_sql]) + .await + .unwrap(); + assert_eq!(row.get::<_, Option<&str>>(0), Some("2")); + let row = client.query_one("SELECT pg_try_advisory_lock($1)", &[&lock_id]).await.unwrap(); + assert!(row.get::<_, bool>(0)); + let row = client.query_one("SELECT pg_advisory_unlock($1)", &[&lock_id]).await.unwrap(); + assert!(row.get::<_, bool>(0)); + + drop(store); + client + .execute(&format!("DROP TABLE IF EXISTS {node_lease_table_name_sql}"), &[]) + .await + .unwrap(); + client.execute(&format!("DROP TABLE IF EXISTS {kv_table_name_sql}"), &[]).await.unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn read_write_remove_list_persist() { let store = create_test_store("test_rwrl").await; @@ -1050,12 +1437,18 @@ mod tests { } } - async fn kill_lock_connection(store: &PostgresStore) { - let client = &store.inner.lock_client; - let _ = client.execute("SELECT pg_terminate_backend(pg_backend_pid())", &[]).await; - while !client.is_closed() { - tokio::task::yield_now().await; - } + async fn expire_lease(store: &PostgresStore) { + let client = store.inner.pool.connections[0].lock().await; + let lease_table = &store.inner.node_lease_table_name_sql; + client + .execute( + &format!( + "UPDATE {lease_table} SET expires_at = clock_timestamp() - interval '1 second' WHERE id = 1" + ), + &[], + ) + .await + .unwrap(); } #[tokio::test(flavor = "multi_thread")] @@ -1083,21 +1476,8 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - #[should_panic( - expected = "PostgreSQL store lock connection closed; continuing may corrupt node state" - )] - async fn test_postgres_store_panics_when_lock_connection_closes() { - let table_name = "test_pg_lock_connection_closed"; - let store = create_test_store(table_name).await; - - kill_lock_connection(&store).await; - cleanup_store(&store).await; - KVStore::write(&store, "test_ns", "test_sub", "key", vec![1u8]).await.unwrap(); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_queued_write_rechecks_closed_lock_connection() { - let table_name = "test_pg_queued_write_lock_connection_closed"; + async fn test_queued_write_rechecks_lease() { + let table_name = "test_pg_queued_write_lease"; let store = create_test_store(table_name).await; let locking_key = store.build_locking_key("test_ns", "test_sub", "key"); let inner_lock_ref = store.inner.get_inner_lock_ref(locking_key); @@ -1105,7 +1485,7 @@ mod tests { let mut write = Box::pin(KVStore::write(&store, "test_ns", "test_sub", "key", vec![1u8])); // Poll the public method once so its initial check passes and its internal write is queued - // on the per-key lock before closing the store lock connection. + // on the per-key lock before expiring the store lease. std::future::poll_fn(|cx| match write.as_mut().poll(cx) { std::task::Poll::Pending => std::task::Poll::Ready(()), std::task::Poll::Ready(result) => { @@ -1114,7 +1494,7 @@ mod tests { }) .await; - kill_lock_connection(&store).await; + expire_lease(&store).await; let second_store = create_test_store(table_name).await; drop(inner_lock); @@ -1126,7 +1506,19 @@ mod tests { .expect_err("the queued write must not access PostgreSQL"); assert_eq!(err.kind(), io::ErrorKind::NotFound); - cleanup_store(&second_store).await; + assert!(store.inner.node_lease.is_lost()); + let remove = tokio::spawn(KVStore::remove(&store, "test_ns", "test_sub", "key", false)); + assert!(remove.await.unwrap_err().is_panic()); + let read = tokio::spawn(KVStore::read(&store, "test_ns", "test_sub", "key")); + assert!(read.await.unwrap_err().is_panic()); + KVStore::write(&second_store, "test_ns", "test_sub", "key", vec![2]).await.unwrap(); + drop(second_store); + let replacement = create_test_store(table_name).await; + assert_eq!( + KVStore::read(&replacement, "test_ns", "test_sub", "key").await.unwrap(), + vec![2] + ); + cleanup_store(&replacement).await; } #[tokio::test(flavor = "multi_thread")] diff --git a/tests/integration_tests_migration.rs b/tests/integration_tests_migration.rs index 84e332bce2..96dfd9e681 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 280c11de5d..8141649749 100644 --- a/tests/integration_tests_postgres.rs +++ b/tests/integration_tests_postgres.rs @@ -15,7 +15,9 @@ use ldk_node::Builder; #[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(); let chain_source = random_chain_source(&bitcoind, &electrsd); @@ -63,12 +65,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 chain_source = random_chain_source(&bitcoind, &electrsd); @@ -135,4 +140,5 @@ async fn postgres_node_restart() { node.stop().unwrap(); drop_table("restart_test").await; + drop_table("restart_test_node_lease").await; }