diff --git a/src/builder.rs b/src/builder.rs index dc41aef1a..030026bff 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -683,6 +683,9 @@ impl NodeBuilder { /// Builds a [`Node`] instance with a [PostgreSQL] backend and according to the options /// previously configured. /// + /// This acquires an exclusive lease for the selected KV table before reading persisted node + /// state. Nodes may share a database when each node identity uses a distinct `kv_table_name`. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -696,13 +699,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 @@ -729,7 +729,13 @@ impl NodeBuilder { log_error!(logger, "Failed to set up Postgres store: {e}"); BuildError::KVStoreSetupFailed })?; - self.build_with_store_runtime_and_logger(node_entropy, kv_store, runtime, logger) + let node_lease = kv_store.node_lease(); + let mut node = + self.build_with_store_runtime_and_logger(node_entropy, kv_store, runtime, logger)?; + if !node.install_node_lease(node_lease) { + return Err(BuildError::KVStoreSetupFailed); + } + Ok(node) } /// Builds a [`Node`] instance with a [`FilesystemStoreV2`] backend and according to the options @@ -1225,6 +1231,9 @@ impl ArcedNodeBuilder { /// Builds a [`Node`] instance with a [PostgreSQL] backend and according to the options /// previously configured. /// + /// This acquires an exclusive lease for the selected KV table before reading persisted node + /// state. Nodes may share a database when each node identity uses a distinct `kv_table_name`. + /// /// Connects to the PostgreSQL database at the given `connection_string`, e.g., /// `"postgres://user:password@localhost/ldk_db"`. /// @@ -1238,13 +1247,10 @@ 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. + /// 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 @@ -2377,6 +2383,7 @@ fn build_with_store_internal( payment_store, lnurl_auth, is_running, + node_lease: None, node_metrics, om_mailbox, async_payments_role, diff --git a/src/io/mod.rs b/src/io/mod.rs index c70c68d96..6061b1fa8 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -7,6 +7,8 @@ //! Objects and traits for data persistence. +#[cfg_attr(not(feature = "postgres"), allow(dead_code))] +pub(crate) mod node_lease; #[cfg(feature = "postgres")] pub mod postgres_store; pub mod sqlite_store; diff --git a/src/io/node_lease.rs b/src/io/node_lease.rs new file mode 100644 index 000000000..6d5016073 --- /dev/null +++ b/src/io/node_lease.rs @@ -0,0 +1,173 @@ +// This file is Copyright its original authors, visible in version control history. +// +// This file is licensed under the Apache License, Version 2.0 or the MIT license , at your option. You may not use this file except in +// accordance with one or both of these licenses. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use lightning::io; + +pub(crate) const NODE_LEASE_DURATION: Duration = Duration::from_secs(30); +// Fail closed before the database lease expires, leaving time for process termination. +pub(crate) const NODE_LEASE_RENEWAL_DEADLINE: Duration = Duration::from_secs(20); +pub(crate) const NODE_LEASE_RENEWAL_INTERVAL: Duration = Duration::from_secs(10); +pub(crate) const NODE_LEASE_RETRY_INTERVAL: Duration = Duration::from_secs(1); +pub(crate) const NODE_LEASE_RELEASE_TIMEOUT: Duration = Duration::from_secs(5); + +type LeaseLossHandler = Box; + +pub(crate) struct NodeLease { + owner_id: [u8; 32], + lease_lost: AtomicBool, + last_confirmed_renewal: Mutex, + loss_sender: tokio::sync::watch::Sender, + loss_handler: Mutex>, +} + +impl NodeLease { + pub(crate) fn new() -> io::Result> { + let mut owner_id = [0u8; 32]; + getrandom::fill(&mut owner_id).map_err(|e| { + io::Error::new(io::ErrorKind::Other, format!("Failed to generate lease owner ID: {e}")) + })?; + let (loss_sender, _) = tokio::sync::watch::channel(false); + Ok(Arc::new(Self { + owner_id, + lease_lost: AtomicBool::new(false), + last_confirmed_renewal: Mutex::new(Instant::now()), + loss_sender, + loss_handler: Mutex::new(None), + })) + } + + pub(crate) fn owner_id(&self) -> &[u8; 32] { + &self.owner_id + } + + pub(crate) fn is_lost(&self) -> bool { + self.lease_lost.load(Ordering::Acquire) + } + + pub(crate) fn record_renewal_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 ensure_operation_active(&self) -> io::Result<()> { + if self.is_lost() || self.renewal_deadline_elapsed() { + self.mark_lost(); + Err(lease_lost_error()) + } else { + Ok(()) + } + } + + pub(crate) fn map_operation_error(&self, error: io::Error) -> io::Error { + // Preserve transient database errors until they outlive the local safety margin. + self.ensure_operation_active().err().unwrap_or(error) + } + + pub(crate) fn mark_lost(&self) { + if self + .lease_lost + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + + // Run any installed containment handler before publishing lease loss. + if let Some(handler) = self.loss_handler.lock().expect("lock").take() { + handler(); + } + self.loss_sender.send_replace(true); + } + + pub(crate) fn set_loss_handler(&self, handler: LeaseLossHandler) { + let mut locked_handler = self.loss_handler.lock().expect("lock"); + if self.is_lost() { + drop(locked_handler); + handler(); + } else { + *locked_handler = Some(handler); + } + } + + pub(crate) async fn wait_for_loss(self: Arc) { + let mut receiver = self.loss_sender.subscribe(); + let _ = receiver.wait_for(|lost| *lost).await; + } +} + +pub(crate) fn lease_lost_error() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "PostgreSQL node lease was lost") +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + + use super::*; + + #[test] + fn expired_operation_marks_loss_before_returning_error() { + let lease = NodeLease::new().unwrap(); + let handler_ran = Arc::new(AtomicBool::new(false)); + let handler_ran_ref = Arc::clone(&handler_ran); + lease.set_loss_handler(Box::new(move || { + handler_ran_ref.store(true, Ordering::Release); + })); + *lease.last_confirmed_renewal.lock().unwrap() = + Instant::now() - NODE_LEASE_RENEWAL_DEADLINE; + + let error = lease.map_operation_error(io::Error::from(io::ErrorKind::Other)); + + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert!(lease.is_lost()); + assert!(handler_ran.load(Ordering::Acquire)); + } + + #[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 c9add1c57..abb7e243f 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 4ffb948ff..62bada92b 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::{ + lease_lost_error, NodeLease, NODE_LEASE_DURATION, NODE_LEASE_RELEASE_TIMEOUT, + NODE_LEASE_RENEWAL_INTERVAL, NODE_LEASE_RETRY_INTERVAL, +}; use crate::io::utils::check_namespace_key_validity; -use crate::logger::{log_debug, log_info, LdkLogger, Logger}; +use crate::logger::{log_debug, log_error, log_info, LdkLogger, Logger}; use crate::runtime::StoreRuntime; mod migrations; @@ -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,26 +94,36 @@ fn sql_table_identifier(table_name: &str) -> io::Result { Ok(quoted_parts?.join(".")) } -/// 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` -/// 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 -/// twice (once normally, once on retry), so it must be side-effect-free outside of issuing the -/// query itself. +fn sql_node_lease_table_identifier(table_name: &str) -> io::Result { + sql_table_identifier(table_name)?; + let table_part = table_name.rsplit_once('.').map_or(table_name, |(_, table)| table); + if table_part.len() > MAX_KV_TABLE_NAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "PostgreSQL KV table name exceeds the maximum of {MAX_KV_TABLE_NAME_BYTES} bytes: {table_name}" + ), + )); + } + + sql_table_identifier(&format!("{table_name}{NODE_LEASE_TABLE_SUFFIX}")) +} + +/// Runs a tokio-postgres query and, if the connection dropped mid-flight, reconnects and retries +/// once. `$store` is the [`PostgresStoreInner`], `$locked` the held client slot guard, +/// `$err_map` an `Fn(PgError) -> io::Error` (called at most once), and `$query` an expression +/// 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 { Ok(v) => Ok(v), Err(e) if $locked.is_closed() || e.is_closed() => { - $store.assert_store_lock(); if let Some(logger) = $store.logger.as_ref() { log_debug!(logger, "Reconnecting to PostgreSQL after error: {e}"); } *$locked = make_config_connection(&$store.config, &$store.tls).await?; - // Recheck after awaiting the connection in case the store lock was lost while - // reconnecting. - $store.assert_store_lock(); $query.await.map_err($err_map) }, Err(e) => Err($err_map(e)), @@ -126,6 +147,7 @@ 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. /// /// [PostgreSQL]: https://www.postgresql.org pub struct PostgresStore { @@ -137,6 +159,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 +183,23 @@ 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::WouldBlock`] 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. + /// Use [`crate::Builder::build_with_postgres_store`] when building a node so ldk-node can install + /// its lease-loss handler before returning it. Passing this store to + /// [`crate::Builder::build_with_store`] directly does not support failover. pub async fn new( connection_string: String, db_name: Option, kv_table_name: Option, certificate_pem: Option, @@ -194,14 +222,67 @@ impl PostgresStore { )?); let tls = Self::build_tls_connector(certificate_pem)?; let task = internal_runtime.spawn(async move { - PostgresStoreInner::new(connection_string, db_name, kv_table_name, tls, logger).await + let (inner, renewal_started_at) = + PostgresStoreInner::new(connection_string, db_name, kv_table_name, tls, logger) + .await?; + Ok::<_, io::Error>((inner, renewal_started_at)) }); - let inner = task.await.map_err(|e| { + let (inner, renewal_started_at) = task.await.map_err(|e| { io::Error::new(io::ErrorKind::Other, format!("PostgreSQL runtime task failed: {}", e)) })??; let inner = Arc::new(inner); - let next_write_version = AtomicU64::new(1); - Ok(Self { inner, next_write_version, internal_runtime: Some(internal_runtime) }) + inner.node_lease.record_renewal_started_at(renewal_started_at); + inner.node_lease.ensure_operation_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.saturating_sub(renewal_started_at.elapsed()); + 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(); + return; + }, + 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(); + return; + }, + Err(e) => { + if let Some(logger) = inner_ref.logger.as_ref() { + log_error!(logger, "Failed to renew PostgreSQL node lease: {e}"); + } + if lease_ref.renewal_deadline_elapsed() { + lease_ref.mark_lost(); + return; + } + next_delay = NODE_LEASE_RETRY_INTERVAL; + }, + } + } + }); + + Ok(Self { + inner, + next_write_version: AtomicU64::new(1), + internal_runtime: Some(internal_runtime), + lease_renewal_task: Some(lease_renewal_task), + }) } fn build_tls_connector(certificate_pem: Option) -> io::Result { @@ -248,10 +329,41 @@ impl PostgresStore { fn internal_runtime(&self) -> Arc { Arc::clone(self.internal_runtime.as_ref().expect("PostgreSQL runtime must be available")) } + + pub(crate) fn node_lease(&self) -> Arc { + Arc::clone(&self.inner.node_lease) + } } impl Drop for PostgresStore { fn drop(&mut self) { + if let Some(internal_runtime) = self.internal_runtime.as_ref() { + let renewal_task = self.lease_renewal_task.take(); + if let Some(task) = renewal_task.as_ref() { + task.abort(); + } + + let runtime_handle = internal_runtime.handle().clone(); + let inner = Arc::clone(&self.inner); + let _ = 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(); @@ -270,7 +382,6 @@ impl KVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner.read_internal(&primary_namespace, &secondary_namespace, &key).await }); @@ -289,7 +400,6 @@ impl KVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner .write_internal( @@ -318,7 +428,6 @@ impl KVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner .remove_internal( @@ -343,7 +452,6 @@ impl KVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner.list_internal(&primary_namespace, &secondary_namespace).await }); @@ -361,7 +469,6 @@ impl PaginatedKVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner .list_paginated_internal(&primary_namespace, &secondary_namespace, page_token) @@ -379,7 +486,6 @@ impl MigratableKVStore for PostgresStore { let inner = Arc::clone(&self.inner); let runtime = self.internal_runtime(); async move { - inner.assert_store_lock(); let task = runtime.spawn(async move { inner.list_all_keys_internal().await }); handle_runtime_task_result(task.await) } @@ -388,11 +494,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>, @@ -402,9 +507,11 @@ impl PostgresStoreInner { async fn new( connection_string: String, db_name: Option, kv_table_name: Option, tls: PgTlsConnector, logger: Option>, - ) -> io::Result { + ) -> io::Result<(Self, Instant)> { 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,25 +550,75 @@ 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, + io::ErrorKind::WouldBlock, format!( "PostgreSQL store for database {db_name} and table {kv_table_name} is already in use" ), )); } + 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::WouldBlock, + "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 +633,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 +670,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,23 +694,56 @@ 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); let write_version_locks = Mutex::new(HashMap::new()); - Ok(Self { - pool, - lock_client: client, - config, - kv_table_name_sql, - tls, - write_version_locks, - logger, - }) + Ok(( + Self { + pool, + config, + kv_table_name_sql, + node_lease_table_name_sql, + node_lease, + tls, + write_version_locks, + logger, + }, + renewal_started_at, + )) } async fn create_database_if_not_exists( @@ -630,18 +825,114 @@ impl PostgresStoreInner { } async fn locked_client(&self) -> io::Result> { - let client = self.pool.get(&self.config, &self.tls, self.logger.as_deref()).await?; - // Recheck after any runtime queue, per-key write lock, and pool wait, immediately before - // the caller accesses PostgreSQL. - self.assert_store_lock(); - Ok(client) + self.pool.get(&self.config, &self.tls, self.logger.as_deref()).await } - fn assert_store_lock(&self) { - assert!( - !self.lock_client.is_closed(), - "PostgreSQL store lock connection closed; continuing may corrupt node state" + 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.ensure_operation_active()?; + + // The local check only fails early. This update is authoritative and holds the row lock + // through the caller's KV mutation and commit. + let lease_duration_secs = NODE_LEASE_DURATION.as_secs() as i64; + let lease_table = &self.node_lease_table_name_sql; + let update_sql = format!( + "UPDATE {lease_table} + SET expires_at = clock_timestamp() + ($2::bigint * interval '1 second') + WHERE id = 1 AND owner_id = $1 AND expires_at > clock_timestamp()" + ); + 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(); + return Err(lease_lost_error()); + } + 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}"), + )) + })?; + if node_lease.is_lost() { + return Err(lease_lost_error()); + } + node_lease.record_renewal_started_at(renewal_started_at); + Ok(()) } fn get_inner_lock_ref(&self, locking_key: String) -> Arc> { @@ -720,17 +1011,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 +1044,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 +1261,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,18 +1279,98 @@ mod tests { } #[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; + 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 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); + 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)); - cleanup_store(&store).await; + 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::WouldBlock); + 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")] @@ -1018,6 +1389,87 @@ mod tests { cleanup_store(&store_1).await; } + #[tokio::test] + async fn node_lease_contends_fences_and_releases() { + let leader = create_test_store("test_pg_lease_shared").await; + let leader_lease = leader.node_lease(); + let leader_lease_table = leader.inner.node_lease_table_name_sql.clone(); + let leader_index_name = "idx_test_pg_lease_shared_paginated"; + { + let client = leader.inner.pool.connections[0].lock().await; + client + .execute( + &format!("DROP INDEX IF EXISTS {}", sql_identifier(leader_index_name).unwrap()), + &[], + ) + .await + .unwrap(); + } + let contention_error = match PostgresStore::new( + test_connection_string(), + None, + Some("test_pg_lease_shared".to_string()), + None, + ) + .await + { + Ok(_) => panic!("a second store acquired the same node lease"), + Err(e) => e, + }; + assert_eq!(contention_error.kind(), io::ErrorKind::WouldBlock); + { + let client = leader.inner.pool.connections[0].lock().await; + let row = client + .query_one("SELECT to_regclass($1)::text", &[&leader_index_name]) + .await + .unwrap(); + assert_eq!(row.get::<_, Option>(0), None); + } + + let independent = create_test_store("test_pg_lease_independent").await; + let independent_lease_table = independent.inner.node_lease_table_name_sql.clone(); + + KVStore::write(&leader, "ns", "", "leader", vec![1]).await.unwrap(); + KVStore::write(&independent, "ns", "", "independent", vec![1]).await.unwrap(); + + { + let client = leader.inner.pool.connections[0].lock().await; + client + .execute( + &format!( + "UPDATE {leader_lease_table} SET expires_at = clock_timestamp() - interval '1 second' WHERE id = 1" + ), + &[], + ) + .await + .unwrap(); + } + + let fallback = create_test_store("test_pg_lease_shared").await; + let loss_notification = Arc::clone(&leader_lease).wait_for_loss(); + let write_error = KVStore::write(&leader, "ns", "", "stale", vec![2]).await.unwrap_err(); + assert_eq!(write_error.kind(), io::ErrorKind::PermissionDenied); + loss_notification.await; + assert!(leader_lease.is_lost()); + + let remove_error = KVStore::remove(&leader, "ns", "", "leader", false).await.unwrap_err(); + assert_eq!(remove_error.kind(), io::ErrorKind::PermissionDenied); + KVStore::write(&fallback, "ns", "", "fallback", vec![3]).await.unwrap(); + + cleanup_store(&leader).await; + drop(fallback); + let replacement = create_test_store("test_pg_lease_shared").await; + cleanup_store(&replacement).await; + drop(replacement); + cleanup_store(&independent).await; + drop(independent); + + let client = leader.inner.pool.connections[0].lock().await; + let _ = client.execute(&format!("DROP TABLE IF EXISTS {leader_lease_table}"), &[]).await; + let _ = + client.execute(&format!("DROP TABLE IF EXISTS {independent_lease_table}"), &[]).await; + } + #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_list_all_keys() { let store = create_test_store("test_pg_list_all_keys").await; @@ -1050,14 +1502,6 @@ 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; - } - } - #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_auto_reconnect() { let store = create_test_store("test_pg_reconnect").await; @@ -1082,53 +1526,6 @@ mod tests { cleanup_store(&store).await; } - #[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"; - 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); - let inner_lock = inner_lock_ref.lock().await; - - 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. - std::future::poll_fn(|cx| match write.as_mut().poll(cx) { - std::task::Poll::Pending => std::task::Poll::Ready(()), - std::task::Poll::Ready(result) => { - panic!("the write must be queued on the per-key lock, got {result:?}") - }, - }) - .await; - - kill_lock_connection(&store).await; - let second_store = create_test_store(table_name).await; - drop(inner_lock); - - let write_task = tokio::spawn(write); - let err = write_task.await.expect_err("the queued write must panic after the lock is lost"); - assert!(err.is_panic()); - let err = KVStore::read(&second_store, "test_ns", "test_sub", "key") - .await - .expect_err("the queued write must not access PostgreSQL"); - assert_eq!(err.kind(), io::ErrorKind::NotFound); - - cleanup_store(&second_store).await; - } - #[tokio::test(flavor = "multi_thread")] async fn test_postgres_store_paginated_listing() { let store = create_test_store("test_pg_paginated").await; diff --git a/src/lib.rs b/src/lib.rs index 14ee734a3..37a2679ba 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -110,6 +110,7 @@ mod util; mod wallet; use std::default::Default; +use std::future::Future; use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; #[cfg(cycle_tests)] @@ -145,6 +146,7 @@ use fee_estimator::{ use ffi::*; use gossip::GossipSource; use graph::NetworkGraph; +use io::node_lease::NodeLease; use io::utils::update_and_persist_node_metrics; pub use lightning; use lightning::chain::BlockLocator; @@ -267,6 +269,7 @@ pub struct Node { payment_store: Arc, lnurl_auth: Arc, is_running: Arc>, + node_lease: Option>, node_metrics: Arc, om_mailbox: Option>, async_payments_role: Option, @@ -277,6 +280,40 @@ pub struct Node { } impl Node { + fn node_lease_is_lost_or_expired(&self) -> bool { + self.node_lease.as_ref().is_some_and(|node_lease| { + if node_lease.renewal_deadline_elapsed() { + node_lease.mark_lost(); + } + node_lease.is_lost() + }) + } + + #[cfg(feature = "postgres")] + pub(crate) fn install_node_lease(&mut self, node_lease: Arc) -> bool { + debug_assert!(self.node_lease.is_none()); + let lease_stop_sender = self.stop_sender.clone(); + let lease_runtime = Arc::downgrade(&self.runtime); + let lease_peer_manager = Arc::downgrade(&self.peer_manager); + let lease_logger = Arc::downgrade(&self.logger); + // Store fencing is the safety boundary; this only limits activity until process exit. + node_lease.set_loss_handler(Box::new(move || { + if let Some(logger) = lease_logger.upgrade() { + log_error!(logger, "PostgreSQL node lease lost, stopping background processing"); + } + lease_stop_sender.send_replace(()); + if let Some(runtime) = lease_runtime.upgrade() { + runtime.abort_background_processor_task(); + } + if let Some(peer_manager) = lease_peer_manager.upgrade() { + peer_manager.disconnect_all_peers(); + } + })); + + self.node_lease = Some(node_lease); + !self.node_lease_is_lost_or_expired() + } + /// Starts the necessary background tasks, such as handling events coming from user input, /// LDK/BDK, and the peer-to-peer network. /// @@ -285,12 +322,17 @@ impl Node { /// /// After this returns, the [`Node`] instance can be controlled via the provided API methods in /// a thread-safe manner. + /// Returns [`Error::PersistenceFailed`] if its PostgreSQL lease is already lost or expires + /// during startup. pub fn start(&self) -> Result<(), Error> { // Acquire a run lock and hold it until we're setup. let mut is_running_lock = self.is_running.write().expect("lock"); if *is_running_lock { return Err(Error::AlreadyRunning); } + if self.node_lease_is_lost_or_expired() { + return Err(Error::PersistenceFailed); + } match self.start_inner(&mut is_running_lock) { Ok(()) => Ok(()), @@ -829,8 +871,18 @@ impl Node { } }); - log_info!(self.logger, "Startup complete."); + // The held write lock makes the lease check below the startup linearization point. *is_running_lock = true; + if self.node_lease_is_lost_or_expired() { + *is_running_lock = false; + // Repeat containment after all startup stop receivers have been created. + self.stop_sender.send_replace(()); + self.runtime.abort_background_processor_task(); + self.peer_manager.disconnect_all_peers(); + return Err(Error::PersistenceFailed); + } + + log_info!(self.logger, "Startup complete."); Ok(()) } @@ -900,6 +952,19 @@ impl Node { Ok(()) } + /// Returns a future which completes when this node loses its PostgreSQL lease. + /// + /// Returns [`None`] when the node was not built with the built-in PostgreSQL store. + /// Lease loss is terminal for this [`Node`]. The caller must immediately terminate the process + /// without calling [`Node::stop`] or dropping the node. A new process may reconstruct the node + /// from persistence. Store fencing cannot revoke an external effect from a process that pauses + /// after a fenced mutation, so every critical external effect must immediately follow one and + /// finish within that mutation's renewed lease interval. + pub fn wait_for_lease_loss(&self) -> Option + Send + 'static> { + let node_lease = self.node_lease.as_ref().map(Arc::clone)?; + Some(async move { node_lease.wait_for_loss().await }) + } + /// Returns the status of the [`Node`]. pub fn status(&self) -> NodeStatus { let is_running = *self.is_running.read().expect("lock"); diff --git a/src/runtime.rs b/src/runtime.rs index 7e29996e6..51743f7d8 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -246,6 +246,14 @@ impl Runtime { }; } + pub(crate) fn abort_background_processor_task(&self) { + if let Some(background_processor_task) = + self.background_processor_task.lock().expect("lock").as_ref() + { + background_processor_task.abort(); + } + } + #[cfg(tokio_unstable)] pub fn log_metrics(&self) { let runtime_handle = self.handle(); diff --git a/tests/integration_tests_migration.rs b/tests/integration_tests_migration.rs index 53f60d19f..dee077bb4 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 c9283039d..470e4220d 100644 --- a/tests/integration_tests_postgres.rs +++ b/tests/integration_tests_postgres.rs @@ -17,7 +17,9 @@ use rand::RngCore; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn channel_full_cycle_with_postgres_store() { drop_table("channel_cycle_a").await; + drop_table("channel_cycle_a_node_lease").await; drop_table("channel_cycle_b").await; + drop_table("channel_cycle_b_node_lease").await; let (bitcoind, electrsd) = common::setup_bitcoind_and_electrsd(); println!("== Node A =="); @@ -65,12 +67,15 @@ async fn channel_full_cycle_with_postgres_store() { .await; drop_table("channel_cycle_a").await; + drop_table("channel_cycle_a_node_lease").await; drop_table("channel_cycle_b").await; + drop_table("channel_cycle_b_node_lease").await; } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn postgres_node_restart() { drop_table("restart_test").await; + drop_table("restart_test_node_lease").await; let (bitcoind, electrsd) = common::setup_bitcoind_and_electrsd(); let esplora_url = format!("http://{}", electrsd.esplora_url.as_ref().unwrap()); @@ -141,4 +146,5 @@ async fn postgres_node_restart() { node.stop().unwrap(); drop_table("restart_test").await; + drop_table("restart_test_node_lease").await; }