AdvisoryLock: never strand a handle acquired after DisposeAsync drained (#396)

Directory.Packages.props

@@ -8,8 +8,8 @@
     <PackageVersion Include="Microsoft.VisualStudio.Threading.Analyzers" Version="17.12.19" />
     <PackageVersion Include="DistributedLock.Postgres" Version="1.3.0" />
     <PackageVersion Include="DotNet.ReproducibleBuilds" Version="1.2.39" />
-    <PackageVersion Include="JasperFx" Version="2.24.1" />
-    <PackageVersion Include="JasperFx.Events" Version="2.0.0" />
+    <PackageVersion Include="JasperFx" Version="2.36.3" />
+    <PackageVersion Include="JasperFx.Events" Version="2.36.3" />
     <PackageVersion Include="Microsoft.Data.Sqlite" Version="10.0.9" />
     <!-- Override the transitive SQLitePCLRaw 2.1.11 (vulnerable native lib,
          GHSA-2m69-gcr7-jv3q) with 3.0.3, which ships the patched SQLite build. -->

src/Weasel.Postgresql.Tests/advisory_lock_usage.cs

@@ -275,6 +275,86 @@ public async Task returns_false_once_the_lock_itself_has_begun_disposing()
     }
 }
 
+// weasel#396: TryAttainLockAsync checked _disposed only at method entry and then stored the winning
+// handle unconditionally. An acquire in flight when DisposeAsync drained _handles put its handle into
+// the drained dictionary, where nothing would ever dispose it — a granted advisory lock held until the
+// process exits. With transaction-scoped locks (Marten's default) that leaves a backend permanently
+// 'idle in transaction' on pg_try_advisory_xact_lock, which Marten's high-water gap detection reads as
+// a live pre-gap reserver and never advances past (marten#5090).
+public class advisory_lock_dispose_race
+{
+    // Distinct from every other lock id used in this file so parallel test classes cannot interfere.
+    private const int TheLockId = 3960001;
+
+    [Fact]
+    public async Task an_acquire_that_completes_after_disposal_does_not_strand_its_lock()
+    {
+        await using var source = NpgsqlDataSource.Create(ConnectionSource.ConnectionString);
+
+        // Start the acquire and dispose while it is still in flight — the acquire needs a database
+        // round trip, disposing an empty handle set does not, so the store reliably lands after the
+        // drain. Repeat so a single lucky interleaving cannot pass this by accident.
+        for (var i = 0; i < 20; i++)
+        {
+            var racedLock = new AdvisoryLock(source, NullLogger.Instance, "localhost",
+                new AdvisoryLockOptions { TransactionalLockEnabled = true });
+
+            var attain = racedLock.TryAttainLockAsync(TheLockId, CancellationToken.None);
+            await racedLock.DisposeAsync();
+
+            // Whatever it reports, it must not leave the lock held: either it stored the handle before
+            // the drain (and the drain disposed it) or it saw the disposal and disposed it itself.
+            await attain;
+        }
+
+        // The proof: a completely separate lock must be able to take the same key. Before the fix the
+        // stranded transaction-scoped handles still held it, so this came back false.
+        await using var otherSource = NpgsqlDataSource.Create(ConnectionSource.ConnectionString);
+        var successor = new AdvisoryLock(otherSource, NullLogger.Instance, "localhost",
+            new AdvisoryLockOptions { TransactionalLockEnabled = true });
+
+        try
+        {
+            (await successor.TryAttainLockAsync(TheLockId, CancellationToken.None))
+                .ShouldBeTrue("a disposed AdvisoryLock stranded the handle it acquired mid-disposal");
+        }
+        finally
+        {
+            await successor.DisposeAsync();
+        }
+    }
+
+    [Fact]
+    public async Task disposal_releases_a_lock_attained_before_it()
+    {
+        await using var source = NpgsqlDataSource.Create(ConnectionSource.ConnectionString);
+
+        var theLock = new AdvisoryLock(source, NullLogger.Instance, "localhost",
+            new AdvisoryLockOptions { TransactionalLockEnabled = true });
+
+        (await theLock.TryAttainLockAsync(TheLockId + 1, CancellationToken.None)).ShouldBeTrue();
+        theLock.HasLock(TheLockId + 1).ShouldBeTrue();
+
+        await theLock.DisposeAsync();
+
+        // The ordinary drain path still works — and a disposed lock reports no lock.
+        theLock.HasLock(TheLockId + 1).ShouldBeFalse();
+
+        await using var otherSource = NpgsqlDataSource.Create(ConnectionSource.ConnectionString);
+        var successor = new AdvisoryLock(otherSource, NullLogger.Instance, "localhost",
+            new AdvisoryLockOptions { TransactionalLockEnabled = true });
+
+        try
+        {
+            (await successor.TryAttainLockAsync(TheLockId + 1, CancellationToken.None)).ShouldBeTrue();
+        }
+        finally
+        {
+            await successor.DisposeAsync();
+        }
+    }
+}
+
 // FirstChanceException is an AppDomain-wide hook, so the latch test above cannot tolerate another
 // collection concurrently provoking a disposed-pool abort. Pin this class to a serial collection.
 [CollectionDefinition("advisory_lock_disposal_guard", DisableParallelization = true)]

src/Weasel.Postgresql/AdvisoryLock.cs

@@ -28,9 +28,15 @@ public class AdvisoryLock : IAdvisoryLock
     private readonly string _databaseName;
     private readonly AdvisoryLockOptions _options;
     private readonly ILogger _logger;
+
+    // weasel#396: every read and write of _handles — and every read and write of _disposed that has to
+    // agree with them — goes through this lock. The dictionary is touched by the caller's leadership
+    // poll, by ReleaseLockAsync, and by DisposeAsync, which can run concurrently; and disposal has to
+    // be atomic with respect to storing a freshly acquired handle (see TryAttainLockAsync).
+    private readonly object _handlesLock = new();
     private readonly Dictionary<int, PostgresDistributedLockHandle> _handles = new();
     private readonly LightweightCache<int, PostgresDistributedLock> _distributedLockProviders;
-    private volatile bool _disposed;
+    private bool _disposed;
 
     public AdvisoryLock(NpgsqlDataSource dataSource, ILogger logger, string databaseName, AdvisoryLockOptions options)
     {
@@ -46,6 +52,17 @@ public AdvisoryLock(NpgsqlDataSource dataSource, ILogger logger, string database
         _options = options;
     }
 
+    private bool IsDisposed
+    {
+        get
+        {
+            lock (_handlesLock)
+            {
+                return _disposed;
+            }
+        }
+    }
+
     private static NpgsqlDataSource EnsurePrimaryWhenMultiHost(NpgsqlDataSource source)
     {
         if (source is NpgsqlMultiHostDataSource multiHostDataSource)
@@ -56,13 +73,21 @@ private static NpgsqlDataSource EnsurePrimaryWhenMultiHost(NpgsqlDataSource sour
 
     public bool HasLock(int lockId)
     {
-        var lockState = _handles.TryGetValue(lockId, out var handle);
-        if (lockState && _options.LockMonitoringEnabled)
+        PostgresDistributedLockHandle? handle;
+        lock (_handlesLock)
+        {
+            if (!_handles.TryGetValue(lockId, out handle))
+            {
+                return false;
+            }
+        }
+
+        if (_options.LockMonitoringEnabled)
         {
-            return !handle!.HandleLostToken.IsCancellationRequested;
+            return !handle.HandleLostToken.IsCancellationRequested;
         }
 
-        return lockState;
+        return true;
     }
 
     /// <summary>
@@ -78,18 +103,48 @@ public async Task<bool> TryAttainLockAsync(int lockId, CancellationToken token)
         // weasel#349: never start a new acquire once disposal has begun. On a HotCold cold/standby node the
         // coordinator polls this on a cadence, and during host shutdown the owned NpgsqlDataSource races with
         // disposal — an in-flight OpenAsync aborts with ObjectDisposedException: 'Npgsql.PoolingDataSource'.
-        if (_disposed) return false;
+        if (IsDisposed) return false;
 
         try
         {
             var locker = _distributedLockProviders[lockId];
             var handle = await locker.TryAcquireAsync(cancellationToken: token).ConfigureAwait(false);
-            if (handle is not null)
+            if (handle is null) return false;
+
+            // weasel#396: the entry check above is not enough on its own. DisposeAsync can drain
+            // _handles while this acquire is in flight, and the handle would then be stored into a
+            // dictionary nothing will ever dispose — a granted advisory lock held for the life of the
+            // process. With transaction-scoped locks that is a permanent 'idle in transaction' backend
+            // on pg_try_advisory_xact_lock, which Marten's high-water gap detection reads as a live
+            // pre-gap reserver and never advances past (marten#5090). Storing under the same lock that
+            // DisposeAsync drains under makes the two orderings exhaustive: either the handle lands
+            // before the drain and the drain disposes it, or it observes the disposal and disposes
+            // itself here.
+            PostgresDistributedLockHandle? orphaned = null;
+            var stored = false;
+
+            lock (_handlesLock)
             {
-                _handles[lockId] = handle;
-                return true;
+                if (_disposed)
+                {
+                    orphaned = handle;
+                }
+                else
+                {
+                    // A handle already sitting in this slot is one whose lock we lost (monitored mode)
+                    // and re-attained; it is displaced, not released, so it has to be disposed too.
+                    _handles.Remove(lockId, out orphaned);
+                    _handles[lockId] = handle;
+                    stored = true;
+                }
             }
-            return false;
+
+            if (orphaned is not null)
+            {
+                await disposeHandleSafelyAsync(orphaned).ConfigureAwait(false);
+            }
+
+            return stored;
         }
         catch (ObjectDisposedException)
         {
@@ -105,10 +160,14 @@ public async Task<bool> TryAttainLockAsync(int lockId, CancellationToken token)
             //
             // Callers that would rather not see it can check HasLock, or simply poll again — the latch guarantees
             // the second call returns false quietly.
-            _disposed = true;
+            lock (_handlesLock)
+            {
+                _disposed = true;
+            }
+
             throw;
         }
-        catch (Exception e) when (_disposed && e is NpgsqlException or InvalidOperationException)
+        catch (Exception e) when (IsDisposed && e is NpgsqlException or InvalidOperationException)
         {
             // Same shutdown race, surfaced as a disposed-pool NpgsqlException / InvalidOperationException.
             return false;
@@ -117,35 +176,53 @@ public async Task<bool> TryAttainLockAsync(int lockId, CancellationToken token)
 
     public async Task ReleaseLockAsync(int lockId)
     {
-        if (_handles.Remove(lockId, out var handle))
+        PostgresDistributedLockHandle? handle;
+        lock (_handlesLock)
+        {
+            _handles.Remove(lockId, out handle);
+        }
+
+        if (handle is not null)
         {
             await handle.DisposeAsync().ConfigureAwait(false);
         }
     }
 
     public async ValueTask DisposeAsync()
     {
-        // Set first, before disposing handles, so any concurrent TryAttainLockAsync short-circuits (weasel#349).
-        _disposed = true;
+        PostgresDistributedLockHandle[] handles;
 
-        foreach (var i in _handles.Keys)
+        lock (_handlesLock)
         {
-            if (_handles.Remove(i, out var handle))
-            {
-                try
-                {
-                    await handle.DisposeAsync().ConfigureAwait(false);
-                }
-                catch (InvalidOperationException)
-                {
-                    // Underlying connection is already closed and there's nothing to dispose. ObjectDisposedException
-                    // derives from this, so a data source that went first lands here too — nothing worth logging.
-                }
-                catch (Exception e)
-                {
-                    _logger.LogError(e, "Error trying to dispose of advisory locks for database {Identifier}", _databaseName);
-                }
-            }
+            // Latch and drain atomically (weasel#349 for the latch, weasel#396 for the drain): a
+            // concurrent TryAttainLockAsync either got its handle into _handles before this snapshot —
+            // in which case it is disposed below — or it will see _disposed set when it tries to store
+            // and dispose the handle itself. Nothing can land in the dictionary after this point.
+            _disposed = true;
+            handles = _handles.Values.ToArray();
+            _handles.Clear();
+        }
+
+        foreach (var handle in handles)
+        {
+            await disposeHandleSafelyAsync(handle).ConfigureAwait(false);
+        }
+    }
+
+    private async Task disposeHandleSafelyAsync(PostgresDistributedLockHandle handle)
+    {
+        try
+        {
+            await handle.DisposeAsync().ConfigureAwait(false);
+        }
+        catch (InvalidOperationException)
+        {
+            // Underlying connection is already closed and there's nothing to dispose. ObjectDisposedException
+            // derives from this, so a data source that went first lands here too — nothing worth logging.
+        }
+        catch (Exception e)
+        {
+            _logger.LogError(e, "Error trying to dispose of advisory locks for database {Identifier}", _databaseName);
         }
     }
 }