Support passive monitoring of Postgres locks

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. The monitored session therefore shows as `idle` in `pg_stat_activity`.
+
+Two things to be aware of:
+- Because the monitored session is idle, server-side idle-session reapers (`idle_session_timeout`, `idle_in_transaction_session_timeout`, or aggressive gateways) can kill it. If any of these are in play, set `KeepaliveCadence`, as when monitoring is active, the keepalive query will be interleaved with the passive wait.
+- 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:

src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs

@@ -357,22 +357,27 @@ private async Task<bool> 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<bool> DoMonitoringAsync(CancellationToken cancellationToken)
+ private async Task<bool> DoMonitoringAsync(TimeoutValue keepaliveCadence, CancellationToken cancellationToken)
  {
+ // 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.
+ TimeoutValue monitoringCadence = TimeSpan.FromMinutes(1);
+
  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),
- cancellationToken: cancellationToken,
+ // on cancellation (the connection is wanted for a real query) or connection loss, loop around and re-evaluate state
+ await connection.MonitorAsync(
+ monitoringCadence,
+ keepaliveCadence,
+ cancellationToken,
  executor: (command, token) => command.ExecuteNonQueryAsync(token, disallowAsyncCancellation: false, isConnectionMonitoringQuery: true)
  ).TryAwait();
 
@@ -393,11 +398,7 @@ private async Task<bool> DoKeepaliveAsync(TimeoutValue keepaliveCadence, Cancell
  using var connectionLockHandle = await this._connectionLock.TryAcquireAsync(TimeSpan.Zero, CancellationToken.None).ConfigureAwait(false);
  if (connectionLockHandle != null)
  {
- using var command = connection.CreateCommand();
- command.SetCommandText("SELECT 0 /* DistributedLock connection keepalive */");
- // Since this query is very fast and non-blocking, we don't bother trying to cancel it. This avoids having 
- // to deal with the overhead of throwing exceptions within ExecuteNonQueryAsync()
- await command.ExecuteNonQueryAsync(CancellationToken.None, disallowAsyncCancellation: false, isConnectionMonitoringQuery: true).AsTask().TryAwait();
+ await connection.ExecuteKeepaliveQueryAsync().ConfigureAwait(false);
  }
 
  return true;

src/DistributedLock.Core/Internal/Data/DatabaseConnection.cs

@@ -171,4 +171,20 @@ private async ValueTask DisposeTransactionAsync(bool isClosingOrDisposingConnect
  public abstract bool IsCommandCancellationException(Exception exception);
 
  public abstract Task SleepAsync(TimeSpan sleepTime, CancellationToken cancellationToken, Func<DatabaseCommand, CancellationToken, ValueTask<int>> executor);
+
+ public virtual Task MonitorAsync(
+ TimeoutValue monitoringCadence,
+ TimeoutValue keepaliveCadence,
+ CancellationToken cancellationToken,
+ Func<DatabaseCommand, CancellationToken, ValueTask<int>> executor) =>
+ this.SleepAsync(monitoringCadence.TimeSpan, cancellationToken, executor);
+
+ public async Task ExecuteKeepaliveQueryAsync()
+ {
+ using var command = this.CreateCommand();
+ command.SetCommandText("SELECT 0 /* DistributedLock connection keepalive */");
+ // Since this query is very fast and non-blocking, we don't bother trying to cancel it. This avoids having
+ // to deal with the overhead of throwing exceptions within ExecuteNonQueryAsync()
+ await command.ExecuteNonQueryAsync(CancellationToken.None, disallowAsyncCancellation: false, isConnectionMonitoringQuery: true).AsTask().TryAwait();
+ }
 }

src/DistributedLock.Postgres/PostgresDatabaseConnection.cs

@@ -2,14 +2,19 @@
 using Medallion.Threading.Internal.Data;
 using Npgsql;
 using System.Data;
-#if NET7_0_OR_GREATER
 using System.Data.Common;
-#endif
 
 namespace Medallion.Threading.Postgres;
 
 internal sealed class PostgresDatabaseConnection : DatabaseConnection
 {
+ /// <summary>
+ /// Set only for connections we own and can therefore run passive monitoring on
+ /// (the monitor never runs background work on externally-owned connections)
+ /// </summary>
+ private readonly NpgsqlConnection? _ownedNpgsqlConnection;
+ private bool? _supportsPassiveMonitoring;
+
  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._ownedNpgsqlConnection = ownedConnection as NpgsqlConnection;
  }
 
  // see https://www.npgsql.org/doc/prepare.html
@@ -40,6 +51,41 @@ exception is PostgresException postgresException
  // cancellation error code from https://www.postgresql.org/docs/10/errcodes-appendix.html
  && postgresException.SqlState == "57014";
 
+ public override async Task MonitorAsync(
+ TimeoutValue monitoringCadence,
+ TimeoutValue keepaliveCadence,
+ CancellationToken cancellationToken,
+ Func<DatabaseCommand, CancellationToken, ValueTask<int>> executor)
+ {
+ if (!this.SupportsPassiveMonitoring)
+ {
+ await base.MonitorAsync(monitoringCadence, keepaliveCadence, cancellationToken, executor).ConfigureAwait(false);
+ return;
+ }
+
+ // Passively wait for connection activity/failure without executing a query, leaving the
+ // session idle server-side. Cap the wait at the keepalive cadence so the session is never
+ // seen as idle for longer than that.
+ var maxWaitTime = !keepaliveCadence.IsInfinite && keepaliveCadence.CompareTo(monitoringCadence) < 0
+ ? keepaliveCadence
+ : monitoringCadence;
+ await this._ownedNpgsqlConnection!.WaitAsync(maxWaitTime.TimeSpan, cancellationToken).ConfigureAwait(false);
+
+ // the passive wait left the session idle; run the keepalive query to prevent idle session killing
+ if (!keepaliveCadence.IsInfinite && !cancellationToken.IsCancellationRequested)
+ {
+ await this.ExecuteKeepaliveQueryAsync().ConfigureAwait(false);
+ }
+ }
+
+ // NpgsqlConnection.Wait is unsupported with Npgsql multiplexing, and unsafe to cancel when Npgsql
+ // KeepAlive is enabled (cancellation mid-keepalive-exchange breaks the connection) — monitoring
+ // falls back to the pg_sleep query in those cases
+ private bool SupportsPassiveMonitoring =>
+ this._supportsPassiveMonitoring ??=
+ this._ownedNpgsqlConnection != null
+ && new NpgsqlConnectionStringBuilder(this._ownedNpgsqlConnection.ConnectionString) is { Multiplexing: false, KeepAlive: 0 };
+
  public override async Task SleepAsync(TimeSpan sleepTime, CancellationToken cancellationToken, Func<DatabaseCommand, CancellationToken, ValueTask<int>> executor)
  {
  Invariant.Require(sleepTime >= TimeSpan.Zero);

src/DistributedLock.Tests/Tests/Postgres/PostgresBehaviorTest.cs

@@ -1,4 +1,4 @@
-using Npgsql;
+using Npgsql;
 using NUnit.Framework;
 using System.Data;
 
@@ -135,7 +135,91 @@ public async Task TestDoesNotDetectConnectionBreakViaState()
 
  Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(.1)), Is.False);
 
- Assert.Throws<NpgsqlException>(() => 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<NpgsqlException>(() => getPidCommand.ExecuteScalar());
+ Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(5)), Is.True);
+ }
+
+ /// <summary>
+ /// Demonstrates that a timed-out <see cref="NpgsqlConnection.WaitAsync(TimeSpan, CancellationToken)"/> is
+ /// non-destructive: it returns false and the connection (including an open transaction) remains usable.
+ /// Passive connection monitoring relies on this.
+ /// </summary>
+ [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);
+ }
+ }
+
+ /// <summary>
+ /// Demonstrates that canceling <see cref="NpgsqlConnection.WaitAsync(TimeSpan, CancellationToken)"/> is
+ /// non-destructive: it throws <see cref="OperationCanceledException"/> 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.
+ /// </summary>
+ [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<OperationCanceledException>(() => 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);
+ }
+
+ /// <summary>
+ /// Demonstrates that a connection killed during <see cref="NpgsqlConnection.WaitAsync(TimeSpan, CancellationToken)"/>
+ /// throws and fires <see cref="System.Data.Common.DbConnection.StateChange"/>, which is what drives
+ /// <see cref="IDistributedSynchronizationHandle.HandleLostToken"/> under passive monitoring.
+ /// </summary>
+ [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<NpgsqlException>(() => waitTask);
+ Assert.That(connection.State, Is.Not.EqualTo(ConnectionState.Open));
  Assert.That(stateChangedEvent.Wait(TimeSpan.FromSeconds(5)), Is.True);
  }