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; }