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.

Author: benthecarman

src/builder.rs

@@ -695,6 +695,8 @@ impl NodeBuilder {
 	///
 	/// The given `kv_table_name` will be used or default to
 	/// [`DEFAULT_KV_TABLE_NAME`](io::postgres_store::DEFAULT_KV_TABLE_NAME).
+	/// Building fails if another PostgreSQL-backed node using the same database and table is still
+	/// alive. 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
@@ -1229,6 +1231,8 @@ impl ArcedNodeBuilder {
 	///
 	/// The given `kv_table_name` will be used or default to
 	/// [`DEFAULT_KV_TABLE_NAME`](io::postgres_store::DEFAULT_KV_TABLE_NAME).
+	/// Building fails if another PostgreSQL-backed node using the same database and table is still
+	/// alive. 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

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<String> {
 	if identifier.is_empty() || identifier.contains('\0') {
 		return Err(io::Error::new(
@@ -128,6 +141,8 @@ impl PostgresStore {
 	/// the default `postgres` database to create it.
 	///
 	/// The given `kv_table_name` will be used or default to [`DEFAULT_KV_TABLE_NAME`].
+	/// Construction fails if another [`PostgresStore`] using the same database and table is still
+	/// alive. 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 +388,9 @@ 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,
 	tls: PgTlsConnector,
@@ -426,6 +444,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 +537,18 @@ 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: client,
+			config,
+			kv_table_name_sql,
+			tls,
+			write_version_locks,
+			logger,
+		})
 	}
 
 	async fn create_database_if_not_exists(
@@ -927,6 +968,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;