From 270a7d9314880c423b608197fa4ba5091d9730d0 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Mon, 3 Aug 2026 02:01:42 -0500 Subject: [PATCH] Lock Postgres stores on initialization Prevent multiple nodes from opening the same PostgreSQL database table at once while allowing separate database and table pairs to coexist. Retain the session-scoped advisory lock for the store lifetime. This change was created with OpenAI Codex. --- src/builder.rs | 16 +++ src/io/postgres_store/mod.rs | 191 +++++++++++++++++++++++++++++++++-- 2 files changed, 199 insertions(+), 8 deletions(-) diff --git a/src/builder.rs b/src/builder.rs index f117800996..e219a2bd7b 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -696,6 +696,14 @@ 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. + /// /// 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 @@ -1230,6 +1238,14 @@ impl ArcedNodeBuilder { /// 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. + /// /// 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 diff --git a/src/io/postgres_store/mod.rs b/src/io/postgres_store/mod.rs index 90b8cdc391..a6dfd94e0a 100644 --- a/src/io/postgres_store/mod.rs +++ b/src/io/postgres_store/mod.rs @@ -11,6 +11,7 @@ use std::future::Future; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; +use bitcoin::hashes::{sha256, Hash, HashEngine}; use lightning::io; use lightning::util::persist::{ KVStore, MigratableKVStore, PageToken, PaginatedKVStore, PaginatedListResponse, @@ -44,6 +45,18 @@ 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; +fn advisory_lock_id(db_name: &str, kv_table_name: &str) -> i64 { + let mut engine = sha256::Hash::engine(); + engine.input(b"ldk-node:postgres-store"); + for component in [db_name, kv_table_name] { + engine.input(&(component.len() as u64).to_be_bytes()); + engine.input(component.as_bytes()); + } + + let hash = sha256::Hash::from_engine(engine).to_byte_array(); + i64::from_be_bytes(hash[..8].try_into().expect("SHA-256 prefix has the expected length")) +} + fn sql_identifier(identifier: &str) -> io::Result { if identifier.is_empty() || identifier.contains('\0') { return Err(io::Error::new( @@ -71,11 +84,11 @@ fn sql_table_identifier(table_name: &str) -> io::Result { } /// 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 -/// that yields a fresh `Future>` each time it's evaluated. `$query` -/// may be evaluated up to twice (once normally, once on retry), so it must be side-effect-free -/// outside of issuing the query itself. +/// once after revalidating the store's advisory lock. `$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's evaluated. `$query` may be evaluated up to twice (once normally, once on retry), +/// so it must be side-effect-free outside of issuing the query itself. macro_rules! query_with_retry { ($store:expr, $locked:ident, $err_map:expr, $query:expr) => {{ match $query.await { @@ -85,6 +98,7 @@ macro_rules! query_with_retry { log_debug!(logger, "Reconnecting to PostgreSQL after error: {e}"); } *$locked = make_config_connection(&$store.config, &$store.tls).await?; + $store.revalidate_store_lock().await?; $query.await.map_err($err_map) }, Err(e) => Err($err_map(e)), @@ -128,6 +142,13 @@ 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. /// /// 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 @@ -373,6 +394,11 @@ 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. The mutex lets us replace the connection and reacquire + // the lock if the session is disconnected. + lock_client: tokio::sync::Mutex, + lock_id: i64, config: Config, kv_table_name_sql: String, tls: PgTlsConnector, @@ -426,6 +452,23 @@ impl PostgresStoreInner { Self::create_database_if_not_exists(&config, &tls, logger.as_deref()).await?; let client = make_config_connection(&config, &tls).await?; + 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) + }, + )?; + if !row.get::<_, bool>(0) { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + format!( + "PostgreSQL store for database {db_name} and table {kv_table_name} is already in use" + ), + )); + } // 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 @@ -502,12 +545,19 @@ impl PostgresStoreInner { io::Error::new(io::ErrorKind::Other, msg) })?; - // Drop the setup client; the pool builds its own POOL_SIZE fresh connections. - drop(client); 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, + lock_client: tokio::sync::Mutex::new(client), + lock_id, + config, + kv_table_name_sql, + tls, + write_version_locks, + logger, + }) } async fn create_database_if_not_exists( @@ -589,9 +639,61 @@ impl PostgresStoreInner { } async fn locked_client(&self) -> io::Result> { + self.ensure_store_lock().await?; self.pool.get(&self.config, &self.tls, self.logger.as_deref()).await } + async fn ensure_store_lock(&self) -> io::Result<()> { + let mut lock_client = self.lock_client.lock().await; + if !lock_client.is_closed() { + return Ok(()); + } + self.reacquire_store_lock(&mut lock_client).await + } + + async fn revalidate_store_lock(&self) -> io::Result<()> { + let mut lock_client = self.lock_client.lock().await; + if !lock_client.is_closed() && lock_client.simple_query("SELECT 1").await.is_ok() { + return Ok(()); + } + self.reacquire_store_lock(&mut lock_client).await + } + + async fn reacquire_store_lock(&self, lock_client: &mut ClientConnection) -> io::Result<()> { + if let Some(logger) = self.logger.as_ref() { + log_debug!(logger, "Reconnecting to PostgreSQL after store lock connection closed"); + } + let db_name = self + .config + .get_dbname() + .expect("database name must be set before reconnecting the store lock"); + let client = make_config_connection(&self.config, &self.tls).await?; + let row = client + .query_one("SELECT pg_try_advisory_lock($1)", &[&self.lock_id]) + .await + .map_err(|e| { + io::Error::new( + io::ErrorKind::Other, + format!( + "Failed to reacquire PostgreSQL store lock for database {} and table {}: {e}", + db_name, self.kv_table_name_sql + ), + ) + })?; + if !row.get::<_, bool>(0) { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + format!( + "PostgreSQL store for database {} and table {} is already in use", + db_name, self.kv_table_name_sql + ), + )); + } + + *lock_client = client; + 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()) @@ -927,6 +1029,29 @@ mod tests { assert!(sql_table_identifier("schema.").is_err()); } + #[test] + fn test_postgres_advisory_lock_id_uses_database_and_table() { + let lock_id = advisory_lock_id("database_a", "table_a"); + assert_eq!(lock_id, advisory_lock_id("database_a", "table_a")); + assert_ne!(lock_id, advisory_lock_id("database_b", "table_a")); + assert_ne!(lock_id, advisory_lock_id("database_a", "table_b")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_postgres_store_advisory_lock() { + let table_name = "test_pg_advisory_lock"; + let store = create_test_store(table_name).await; + + let err = + PostgresStore::new(test_connection_string(), None, Some(table_name.to_string()), None) + .await + .err() + .expect("a second store using the same database and table must fail"); + assert_eq!(err.kind(), io::ErrorKind::AlreadyExists); + + cleanup_store(&store).await; + } + #[tokio::test(flavor = "multi_thread")] async fn read_write_remove_list_persist() { let store = create_test_store("test_rwrl").await; @@ -975,6 +1100,14 @@ mod tests { } } + async fn kill_lock_connection(store: &PostgresStore) { + let client = store.inner.lock_client.lock().await; + let _ = client.execute("SELECT pg_terminate_backend(pg_backend_pid())", &[]).await; + while !client.is_closed() { + tokio::task::yield_now().await; + } + } + #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_auto_reconnect() { let store = create_test_store("test_pg_reconnect").await; @@ -999,6 +1132,48 @@ mod tests { cleanup_store(&store).await; } + #[tokio::test(flavor = "multi_thread")] + async fn test_postgres_store_lock_auto_reconnect() { + let table_name = "test_pg_lock_reconnect"; + let store = create_test_store(table_name).await; + + kill_lock_connection(&store).await; + KVStore::write(&store, "test_ns", "test_sub", "key", vec![1u8]).await.unwrap(); + + let err = + PostgresStore::new(test_connection_string(), None, Some(table_name.to_string()), None) + .await + .err() + .expect("the first store must reacquire its advisory lock"); + assert_eq!(err.kind(), io::ErrorKind::AlreadyExists); + + cleanup_store(&store).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_postgres_query_retry_revalidates_store_lock() -> io::Result<()> { + let table_name = "test_pg_retry_lock_reconnect"; + let store = create_test_store(table_name).await; + let mut locked = store.inner.locked_client().await.unwrap(); + + kill_lock_connection(&store).await; + let second_store = create_test_store(table_name).await; + + let _ = locked.execute("SELECT pg_terminate_backend(pg_backend_pid())", &[]).await; + while !locked.is_closed() { + tokio::task::yield_now().await; + } + + let err_map = |e: PgError| io::Error::new(io::ErrorKind::Other, e); + let err = query_with_retry!(store.inner, locked, err_map, locked.simple_query("SELECT 1")) + .err() + .expect("the query retry must fail after another store acquires the advisory lock"); + assert_eq!(err.kind(), io::ErrorKind::AlreadyExists); + + cleanup_store(&second_store).await; + Ok(()) + } + #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_paginated_listing() { let store = create_test_store("test_pg_paginated").await;