diff --git a/docs/DistributedLock.Postgres.md b/docs/DistributedLock.Postgres.md index 6ff9403..55ecae6 100644 --- a/docs/DistributedLock.Postgres.md +++ b/docs/DistributedLock.Postgres.md @@ -41,6 +41,14 @@ Under the hood, [Postgres advisory locks can be based on either one 64-bit integ In addition to specifying the `key`, Postgres-based locks allow you to specify either a `connectionString`, an `IDbConnection`, or a `DbDataSource` as a means of connecting to the database. In most cases, using a `connectionString` is preferred because it allows for the library to efficiently multiplex connections under the hood and, in the case of `IDbConnection`, eliminates the risk that the passed-in `IDbConnection` gets used in a way that disrupts the locking process. **NOTE that since `IDbConnection` objects are not thread-safe, lock objects constructed with them can only be used by one thread at a time.** +### Connection monitoring (`HandleLostToken`) + +When `HandleLostToken` is used on a lock backed by a library-owned connection, the library monitors the connection passively using `NpgsqlConnection.WaitAsync`, which uses a blocking socket read that detects connection loss (e.g. a database restart or `pg_terminate_backend`) as soon as the socket breaks, without executing any query. A cheap keepalive query is executed once per passive wait, which happens every `KeepaliveCadence` or every minute, whichever is shorter. Between keepalives the monitored session shows as `idle` in `pg_stat_activity`. + +Two things to be aware of: +- Because the monitored session is idle between keepalives, server-side idle-session reapers (`idle_session_timeout`, `idle_in_transaction_session_timeout`, or aggressive gateways) configured with a timeout under one minute can kill it. If any of these are in play, set `KeepaliveCadence` below the reaper timeout. +- If the connection string enables Npgsql `Multiplexing` (where `Wait` is unsupported) or Npgsql `KeepAlive` (where interrupting `Wait` is not safe), monitoring falls back to parking a `pg_sleep` query on the connection, which shows as an active long-running query. + ## Options In addition to specifying the `key`, several tuning options are available for `connectionString`-based locks: diff --git a/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs b/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs index 41d796f..27b7ddb 100644 --- a/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs +++ b/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs @@ -357,21 +357,23 @@ private async Task TryKeepaliveOrMonitorAsync() stateChangedToken = this._monitorStateChangedTokenSource!.Token; } - return await (isMonitoring ? this.DoMonitoringAsync(stateChangedToken) : this.DoKeepaliveAsync(keepaliveCadence, stateChangedToken)).ConfigureAwait(false); + return await (isMonitoring ? this.DoMonitoringAsync(keepaliveCadence, stateChangedToken) : this.DoKeepaliveAsync(keepaliveCadence, stateChangedToken)).ConfigureAwait(false); } - private async Task DoMonitoringAsync(CancellationToken cancellationToken) + private async Task DoMonitoringAsync(TimeoutValue keepaliveCadence, CancellationToken cancellationToken) { if (!this._weakConnection.TryGetTarget(out var connection)) { return false; } // don't pass token here: this should finish quickly and we don't want to throw using var _ = await this._connectionLock.AcquireAsync(CancellationToken.None).ConfigureAwait(false); - // 1-min increments is kind of an arbitrary choice. We want to avoid this being too short since each time - // we "come up to breathe" that's a waste of resources. We also want to avoid this being too long since - // in case people have some kind of monitoring set up for hanging queries await connection.SleepAsync( - sleepTime: TimeSpan.FromMinutes(1), + // 1-min increments is kind of an arbitrary choice. We want to avoid this being too short since each time + // we "come up to breathe" that's a waste of resources. We also want to avoid this being too long since + // in case people have some kind of monitoring set up for hanging queries. Coming up to breathe also + // re-resolves the weak connection reference so that this loop never roots an abandoned connection for long. + // Capped at keepaliveCadence so that keepalive queries keep firing on cadence while monitoring. + sleepTime: (keepaliveCadence.CompareTo(TimeSpan.FromMinutes(1)) < 0 ? keepaliveCadence : TimeSpan.FromMinutes(1)).TimeSpan, cancellationToken: cancellationToken, executor: (command, token) => command.ExecuteNonQueryAsync(token, disallowAsyncCancellation: false, isConnectionMonitoringQuery: true) ).TryAwait(); diff --git a/src/DistributedLock.Postgres/PostgresDatabaseConnection.cs b/src/DistributedLock.Postgres/PostgresDatabaseConnection.cs index 1ce3ada..52120ba 100644 --- a/src/DistributedLock.Postgres/PostgresDatabaseConnection.cs +++ b/src/DistributedLock.Postgres/PostgresDatabaseConnection.cs @@ -1,15 +1,20 @@ -using Medallion.Threading.Internal; +using Medallion.Threading.Internal; using Medallion.Threading.Internal.Data; using Npgsql; using System.Data; -#if NET7_0_OR_GREATER using System.Data.Common; -#endif +using System.Diagnostics; namespace Medallion.Threading.Postgres; internal sealed class PostgresDatabaseConnection : DatabaseConnection { + /// + /// Only safe to use inside , where the connection monitor holds the connection lock. + /// Non-default only for connections we own (passive monitoring never touches externally-owned connections). + /// + private readonly WaitAsyncWrapper _unsafeWaitAsyncWrapper; + public PostgresDatabaseConnection(IDbConnection connection) : base(connection, isExternallyOwned: true) { @@ -22,14 +27,20 @@ public PostgresDatabaseConnection(IDbTransaction transaction) #if NET7_0_OR_GREATER public PostgresDatabaseConnection(DbDataSource dbDataSource) - : base(dbDataSource.CreateConnection(), isExternallyOwned: false) + : this(dbDataSource.CreateConnection()) { } #endif public PostgresDatabaseConnection(string connectionString) - : base(new NpgsqlConnection(connectionString), isExternallyOwned: false) + : this(new NpgsqlConnection(connectionString)) + { + } + + private PostgresDatabaseConnection(DbConnection ownedConnection) + : base(ownedConnection, isExternallyOwned: false) { + this._unsafeWaitAsyncWrapper = new(ownedConnection); } // see https://www.npgsql.org/doc/prepare.html @@ -44,6 +55,18 @@ public override async Task SleepAsync(TimeSpan sleepTime, CancellationToken canc { Invariant.Require(sleepTime >= TimeSpan.Zero); + // Where supported, "sleep" by passively waiting for connection activity/failure without executing a + // query. This detects connection loss as soon as the socket breaks rather than when the sleep query + // errors, and leaves the session idle server-side. + if (await this._unsafeWaitAsyncWrapper.TryWaitAsync(sleepTime, cancellationToken).ConfigureAwait(false)) + { + // the passive wait left the session idle; run a keepalive query to prevent idle session reaping + using var keepaliveCommand = this.CreateCommand(); + keepaliveCommand.SetCommandText("SELECT 0 /* DistributedLock connection keepalive */"); + await executor(keepaliveCommand, cancellationToken).ConfigureAwait(false); + return; + } + // if we're in a transaction, we need to establish a savepoint so that we can roll back if we // get canceled without the whole transaction being aborted const string SavePointName = "medallion_threading_postgres_database_connection_sleep"; @@ -74,4 +97,40 @@ public override async Task SleepAsync(TimeSpan sleepTime, CancellationToken canc } } } -} \ No newline at end of file + + /// + /// Exposes only from the wrapped + /// connection, keeping the rest of the (non-thread-safe) connection surface inaccessible. + /// + private readonly struct WaitAsyncWrapper(DbConnection dbConnection) + { + // Null when passive waiting is unsupported: non-Npgsql connection, Npgsql multiplexing (Wait is + // unsupported), or Npgsql KeepAlive (cancellation mid-keepalive-exchange breaks the connection) + private readonly NpgsqlConnection? _connection = + dbConnection is NpgsqlConnection npgsqlConnection + && new NpgsqlConnectionStringBuilder(npgsqlConnection.ConnectionString) is { Multiplexing: false, KeepAlive: 0 } + ? npgsqlConnection + : null; + + public async ValueTask TryWaitAsync(TimeSpan timeout, CancellationToken cancellationToken) + { + if (this._connection is null) { return false; } + + // WaitAsync completes when ANY message arrives (e.g. a notification), not just on timeout, + // so loop until the full timeout has elapsed + var startTimestamp = Stopwatch.GetTimestamp(); + var remaining = timeout; + while (await this._connection.WaitAsync(remaining, cancellationToken).ConfigureAwait(false)) + { +#if NET7_0_OR_GREATER + var elapsed = Stopwatch.GetElapsedTime(startTimestamp); +#else + var elapsed = TimeSpan.FromSeconds((Stopwatch.GetTimestamp() - startTimestamp) / (double)Stopwatch.Frequency); +#endif + remaining = timeout - elapsed; + if (remaining <= TimeSpan.Zero) { break; } + } + return true; + } + } +} diff --git a/src/DistributedLock.Tests/Tests/Postgres/PostgresBehaviorTest.cs b/src/DistributedLock.Tests/Tests/Postgres/PostgresBehaviorTest.cs index 74ff8e6..42526ad 100644 --- a/src/DistributedLock.Tests/Tests/Postgres/PostgresBehaviorTest.cs +++ b/src/DistributedLock.Tests/Tests/Postgres/PostgresBehaviorTest.cs @@ -1,4 +1,4 @@ -using Npgsql; +using Npgsql; using NUnit.Framework; using System.Data; @@ -135,7 +135,121 @@ public async Task TestDoesNotDetectConnectionBreakViaState() Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(.1)), Is.False); - Assert.Throws(() => getPidCommand.ExecuteScalar()); + // Catch rather than Throws because whether this surfaces as NpgsqlException (broken connection) + // or the derived PostgresException (the server's 57P01 error message was read first) is timing-dependent + Assert.Catch(() => getPidCommand.ExecuteScalar()); + Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(5)), Is.True); + } + + /// + /// Demonstrates that a timed-out is + /// non-destructive: it returns false and the connection (including an open transaction) remains usable. + /// Passive connection monitoring relies on this. + /// + [Test] + public async Task TestWaitAsyncTimeoutDoesNotBreakConnection() + { + using var connection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await connection.OpenAsync(); + + Assert.That(await connection.WaitAsync(TimeSpan.FromMilliseconds(100), CancellationToken.None), Is.False); + + using var command = connection.CreateCommand(); + command.CommandText = "SELECT 1"; + (await command.ExecuteScalarAsync()).ShouldEqual(1); + + using (var transaction = connection.BeginTransaction()) + { + Assert.That(await connection.WaitAsync(TimeSpan.FromMilliseconds(100), CancellationToken.None), Is.False); + + // the transaction was not aborted by the timed-out wait + command.Transaction = transaction; + command.CommandText = "SELECT 2"; + (await command.ExecuteScalarAsync()).ShouldEqual(2); + } + } + + /// + /// Demonstrates that canceling is + /// non-destructive: it throws and the connection remains usable. + /// Passive connection monitoring relies on this because the monitor's wait is canceled whenever the + /// connection is needed for a real query. + /// + [Test] + public async Task TestWaitAsyncCancellationDoesNotBreakConnection() + { + using var connection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await connection.OpenAsync(); + + using var cancellationTokenSource = new CancellationTokenSource(); + cancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(.5)); + Assert.CatchAsync(() => connection.WaitAsync(TimeSpan.FromSeconds(30), cancellationTokenSource.Token)); + + Assert.That(connection.State, Is.EqualTo(ConnectionState.Open)); + using var command = connection.CreateCommand(); + command.CommandText = "SELECT 1"; + (await command.ExecuteScalarAsync()).ShouldEqual(1); + } + + /// + /// Demonstrates that returns true + /// as soon as any message (e.g. a notification) arrives, before the timeout elapses. Passive connection + /// monitoring accounts for this by looping until its full wait time has elapsed. + /// + [Test] + public async Task TestWaitAsyncReturnsTrueWhenMessageArrives() + { + var channelName = $"wait_test_{Guid.NewGuid():N}"; + + using var connection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await connection.OpenAsync(); + using (var listenCommand = connection.CreateCommand()) + { + listenCommand.CommandText = $"LISTEN {channelName}"; + await listenCommand.ExecuteNonQueryAsync(); + } + + var waitTask = connection.WaitAsync(TimeSpan.FromSeconds(30), CancellationToken.None); + + using var notifyingConnection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await notifyingConnection.OpenAsync(); + using var notifyCommand = notifyingConnection.CreateCommand(); + notifyCommand.CommandText = $"NOTIFY {channelName}"; + await notifyCommand.ExecuteNonQueryAsync(); + + Assert.That(await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(10))), Is.SameAs(waitTask), "wait should complete when the notification arrives"); + Assert.That(await waitTask, Is.True); + } + + /// + /// Demonstrates that a connection killed during + /// throws and fires , which is what drives + /// under passive monitoring. + /// + [Test] + public async Task TestWaitAsyncOnKilledConnectionFiresStateChanged() + { + using var stateChangedEvent = new ManualResetEventSlim(initialState: false); + + using var connection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await connection.OpenAsync(); + connection.StateChange += (o, e) => stateChangedEvent.Set(); + + using var getPidCommand = connection.CreateCommand(); + getPidCommand.CommandText = "SELECT pg_backend_pid()"; + var pid = (int)(await getPidCommand.ExecuteScalarAsync())!; + + var waitTask = connection.WaitAsync(TimeSpan.FromSeconds(30), CancellationToken.None); + + // kill the connection from the back end + using var killingConnection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await killingConnection.OpenAsync(); + using var killCommand = killingConnection.CreateCommand(); + killCommand.CommandText = $"SELECT pg_terminate_backend({pid})"; + await killCommand.ExecuteNonQueryAsync(); + + Assert.CatchAsync(() => waitTask); + Assert.That(connection.State, Is.Not.EqualTo(ConnectionState.Open)); Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(5)), Is.True); } diff --git a/src/DistributedLock.Tests/Tests/Postgres/PostgresConnectionMonitoringTest.cs b/src/DistributedLock.Tests/Tests/Postgres/PostgresConnectionMonitoringTest.cs new file mode 100644 index 0000000..9e1d12b --- /dev/null +++ b/src/DistributedLock.Tests/Tests/Postgres/PostgresConnectionMonitoringTest.cs @@ -0,0 +1,126 @@ +using Medallion.Threading.Postgres; +using Medallion.Threading.Tests.Data; +using Npgsql; +using NUnit.Framework; + +namespace Medallion.Threading.Tests.Postgres; + +/// +/// Tests for passive connection monitoring on Postgres: when +/// is used on an owned connection, monitoring uses +/// (the session stays idle) rather than parking a pg_sleep query on the connection. +/// +public class PostgresConnectionMonitoringTest +{ + private readonly TestingPostgresDb _db = new(); + + [Test] + public async Task TestMonitoringSessionIsIdleWithoutSleepQuery() + { + var applicationName = UniqueApplicationName(); + var @lock = CreateLock(applicationName); + await using var handle = await @lock.AcquireAsync(); + + Assert.That(handle.HandleLostToken.CanBeCanceled, Is.True); // starts monitoring + + // give the monitoring worker time to engage + await Task.Delay(TimeSpan.FromSeconds(1)); + + using var connection = new NpgsqlConnection(TestingPostgresDb.DefaultConnectionString); + await connection.OpenAsync(); + using var command = connection.CreateCommand(); + command.CommandText = @" + SELECT state, query + FROM pg_stat_activity + WHERE application_name = @applicationName"; + command.Parameters.AddWithValue("applicationName", applicationName); + + var sessions = new List<(string State, string Query)>(); + using (var reader = await command.ExecuteReaderAsync()) + { + while (await reader.ReadAsync()) + { + sessions.Add(( + reader.IsDBNull(0) ? string.Empty : reader.GetString(0), + reader.IsDBNull(1) ? string.Empty : reader.GetString(1) + )); + } + } + + Assert.That(sessions, Is.Not.Empty); + Assert.That(sessions, Has.All.Matches<(string State, string Query)>(s => s.State == "idle"), "monitored sessions should not be running a query"); + Assert.That(sessions, Has.None.Matches<(string State, string Query)>(s => s.Query.Contains("pg_sleep")), "monitoring should not use pg_sleep"); + } + + [Test] + public async Task TestHandleLostTokenFiresOnKilledConnectionWithPassiveMonitoring() + { + var applicationName = UniqueApplicationName(); + var @lock = CreateLock(applicationName); + var handle = await @lock.AcquireAsync(); + + using var handleLostEvent = new ManualResetEventSlim(initialState: false); + using var registration = handle.HandleLostToken.Register(handleLostEvent.Set); + + await this._db.KillSessionsAsync(applicationName, idleSince: null); + + Assert.That(handleLostEvent.Wait(TimeSpan.FromSeconds(10)), Is.True); + + // dispose may throw since the underlying connection is broken + try { handle.Dispose(); } catch { } + } + + [Test] + [NonParallelizable, Retry(5)] // timing-sensitive + public async Task TestMonitoringWithKeepaliveCadenceSurvivesIdleSessionKiller() + { + var applicationName = UniqueApplicationName(); + var @lock = CreateLock(applicationName, options => options.KeepaliveCadence(TimeSpan.FromSeconds(.05))); + + var handle = await @lock.AcquireAsync(); + Assert.That(handle.HandleLostToken.CanBeCanceled, Is.True); // monitoring + keepalive cadence => keepalive interleave + + using var idleSessionKiller = new IdleSessionKiller(this._db, applicationName, idleTimeout: TimeSpan.FromSeconds(.5)); + await Task.Delay(TimeSpan.FromSeconds(2)); + + Assert.That(handle.HandleLostToken.IsCancellationRequested, Is.False); + Assert.DoesNotThrow(handle.Dispose); + } + + /// + /// Npgsql KeepAlive is incompatible with canceling , + /// so monitoring falls back to the pg_sleep query for such connection strings. This verifies the fallback end-to-end. + /// + [Test] + public async Task TestHandleLostTokenWorksWithNpgsqlKeepAliveFallback() + { + var applicationName = UniqueApplicationName(); + var @lock = CreateLock(applicationName, connectionStringOptions: builder => builder.KeepAlive = 1); + var handle = await @lock.AcquireAsync(); + + using var handleLostEvent = new ManualResetEventSlim(initialState: false); + Assert.That(handle.HandleLostToken.CanBeCanceled, Is.True); + using var registration = handle.HandleLostToken.Register(handleLostEvent.Set); + + await this._db.KillSessionsAsync(applicationName, idleSince: null); + + Assert.That(handleLostEvent.Wait(TimeSpan.FromSeconds(10)), Is.True); + + // dispose may throw since the underlying connection is broken + try { handle.Dispose(); } catch { } + } + + private static string UniqueApplicationName() => $"monitoring_test_{Guid.NewGuid():N}"; + + private static PostgresDistributedLock CreateLock( + string applicationName, + Action? options = null, + Action? connectionStringOptions = null) + { + var connectionStringBuilder = new NpgsqlConnectionStringBuilder(TestingPostgresDb.DefaultConnectionString) { ApplicationName = applicationName }; + connectionStringOptions?.Invoke(connectionStringBuilder); + + // use a unique lock name since advisory lock keys are global to the database (and some tests retry) + return new PostgresDistributedLock(new(Guid.NewGuid().ToString(), allowHashing: true), connectionStringBuilder.ConnectionString, options); + } +}