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. 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:

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

@@ -357,21 +357,23 @@ 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)
     {
         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();

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
 {
+    /// <summary>
+    /// Only safe to use inside <see cref="SleepAsync"/>, where the connection monitor holds the connection lock.
+    /// Non-default only for connections we own (passive monitoring never touches externally-owned connections).
+    /// </summary>
+    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
             }
         }
     }
-}
+
+    /// <summary>
+    /// Exposes only <see cref="NpgsqlConnection.WaitAsync(TimeSpan, CancellationToken)"/> from the wrapped
+    /// connection, keeping the rest of the (non-thread-safe) connection surface inaccessible.
+    /// </summary>
+    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<bool> 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;
+        }
+    }
+}

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<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 <see cref="NpgsqlConnection.WaitAsync(TimeSpan, CancellationToken)"/> 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.
+    /// </summary>
+    [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);
+    }
+
+    /// <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);
     }