Advisory locks

Directory.Build.props

@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Project>
     <PropertyGroup>
-        <Version>7.6.0</Version>
+        <Version>7.7.0</Version>
         <LangVersion>12.0</LangVersion>
         <ImplicitUsings>enable</ImplicitUsings>
         <Nullable>enable</Nullable>

src/Weasel.Core/IAdvisoryLock.cs

@@ -0,0 +1,8 @@
+namespace Weasel.Core;
+
+public interface IAdvisoryLock: IAsyncDisposable
+{
+    bool HasLock(int lockId);
+    Task<bool> TryAttainLockAsync(int lockId, CancellationToken token);
+    Task ReleaseLockAsync(int lockId);
+}

src/Weasel.Postgresql.Tests/advisory_lock_usage.cs

@@ -1,11 +1,16 @@
-using Npgsql;
+using Microsoft.Extensions.Logging.Abstractions;
+using Npgsql;
 using Shouldly;
+using Weasel.Core;
+using Weasel.Core.Migrations;
 using Xunit;
 
 namespace Weasel.Postgresql.Tests;
 
 public class advisory_lock_usage
 {
+
+
     [Fact]
     public async Task explicitly_release_global_session_locks()
     {
@@ -133,3 +138,62 @@ public async Task tx_session_locks()
         await tx3.RollbackAsync();
     }
 }
+
+public class AdvisoryLockSpecs : IAsyncLifetime
+{
+    private SimplePostgresqlDatabase _database;
+    private AdvisoryLock theLock;
+
+    public Task InitializeAsync()
+    {
+        _database = new SimplePostgresqlDatabase(NpgsqlDataSource.Create(ConnectionSource.ConnectionString));
+        theLock = new AdvisoryLock(_database, NullLogger.Instance);
+        return Task.CompletedTask;
+    }
+
+    public async Task DisposeAsync()
+    {
+        await theLock.DisposeAsync();
+        await _database.DisposeAsync();
+    }
+
+
+    [Fact]
+    public async Task explicitly_release_global_session_locks()
+    {
+        await using var conn2 = new NpgsqlConnection(ConnectionSource.ConnectionString);
+        await using var conn3 = new NpgsqlConnection(ConnectionSource.ConnectionString);
+
+        await conn2.OpenAsync();
+        await conn3.OpenAsync();
+
+        await theLock.TryAttainLockAsync(10, CancellationToken.None);
+
+        // Cannot get the lock here
+        (await conn2.TryGetGlobalLock(10)).Succeeded.ShouldBeFalse();
+
+        await theLock.ReleaseLockAsync(10);
+
+        for (var j = 0; j < 5; j++)
+        {
+            if ((await conn2.TryGetGlobalLock(10)).Succeeded) return;
+
+            await Task.Delay(250);
+        }
+
+        throw new Exception("Advisory lock was not released");
+    }
+
+}
+
+public class SimplePostgresqlDatabase: PostgresqlDatabase
+{
+    public SimplePostgresqlDatabase(NpgsqlDataSource dataSource) : base(new DefaultMigrationLogger(), AutoCreate.CreateOrUpdate, new PostgresqlMigrator(), "Simple", dataSource)
+    {
+    }
+
+    public override IFeatureSchema[] BuildFeatureSchemas()
+    {
+        return Array.Empty<IFeatureSchema>();
+    }
+}

src/Weasel.Postgresql/AdvisoryLock.cs

@@ -0,0 +1,114 @@
+using System.Data;
+using JasperFx.Core;
+using Microsoft.Extensions.Logging;
+using Npgsql;
+using Weasel.Core;
+using Weasel.Core.Migrations;
+
+namespace Weasel.Postgresql;
+
+internal class AdvisoryLock : IAdvisoryLock
+{
+    private readonly PostgresqlDatabase _database;
+    private readonly ILogger _logger;
+    private NpgsqlConnection _conn;
+    private readonly List<int> _locks = new();
+
+    public AdvisoryLock(PostgresqlDatabase database, ILogger logger)
+    {
+        _database = database;
+        _logger = logger;
+    }
+
+    public bool HasLock(int lockId)
+    {
+        return _conn is not { State: ConnectionState.Closed } && _locks.Contains(lockId);
+    }
+
+    public async Task<bool> TryAttainLockAsync(int lockId, CancellationToken token)
+    {
+        if (_conn == null)
+        {
+            _conn = _database.CreateConnection();
+            await _conn.OpenAsync(token).ConfigureAwait(false);
+        }
+
+        if (_conn.State == ConnectionState.Closed)
+        {
+            try
+            {
+                await _conn.DisposeAsync().ConfigureAwait(false);
+            }
+            catch (Exception e)
+            {
+                _logger.LogError(e, "Error trying to clean up and restart an advisory lock connection");
+            }
+            finally
+            {
+                _conn = null;
+            }
+
+            return false;
+        }
+
+
+
+        var attained = await _conn.TryGetGlobalLock(lockId, cancellation: token).ConfigureAwait(false);
+        if (attained == AttainLockResult.Success)
+        {
+            _locks.Add(lockId);
+            return true;
+        }
+
+        return false;
+    }
+
+    public async Task ReleaseLockAsync(int lockId)
+    {
+        if (!_locks.Contains(lockId)) return;
+
+        if (_conn == null || _conn.State == ConnectionState.Closed)
+        {
+            _locks.Remove(lockId);
+            return;
+        }
+
+        var cancellation = new CancellationTokenSource();
+        cancellation.CancelAfter(1.Seconds());
+
+        await _conn.ReleaseGlobalLock(lockId, cancellation: cancellation.Token).ConfigureAwait(false);
+        _locks.Remove(lockId);
+
+        if (!_locks.Any())
+        {
+            await _conn.CloseAsync().ConfigureAwait(false);
+            await _conn.DisposeAsync().ConfigureAwait(false);
+            _conn = null;
+        }
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_conn == null) return;
+
+        try
+        {
+            foreach (var i in _locks)
+            {
+                await _conn.ReleaseGlobalLock(i, CancellationToken.None).ConfigureAwait(false);
+            }
+
+            await _conn.CloseAsync().ConfigureAwait(false);
+            await _conn.DisposeAsync().ConfigureAwait(false);
+        }
+        catch (Exception e)
+        {
+            _logger.LogError(e, "Error trying to dispose of advisory locks for database {Identifier}",
+                _database.Identifier);
+        }
+        finally
+        {
+            await _conn.DisposeAsync().ConfigureAwait(false);
+        }
+    }
+}

src/Weasel.Postgresql/PostgresqlDatabase.cs

@@ -5,7 +5,7 @@
 
 namespace Weasel.Postgresql;
 
-public abstract class PostgresqlDatabase: DatabaseBase<NpgsqlConnection>
+public abstract class PostgresqlDatabase: DatabaseBase<NpgsqlConnection>, IAsyncDisposable
 {
     protected PostgresqlDatabase(
         IMigrationLogger logger,
@@ -58,4 +58,9 @@ public NpgsqlConnection CreateConnection(TargetSessionAttributes targetSessionAt
 
         return base.CreateConnection();
     }
+
+    public ValueTask DisposeAsync()
+    {
+        return DataSource.DisposeAsync();
+    }
 }

src/Weasel.SqlServer.Tests/advisory_lock_usage.cs

@@ -1,4 +1,5 @@
 using Microsoft.Data.SqlClient;
+using Microsoft.Extensions.Logging.Abstractions;
 using Shouldly;
 using Xunit;
 
@@ -142,3 +143,48 @@ public async Task tx_session_locks()
         }
     }
 }
+
+public class AdvisoryLockSpecs : IAsyncLifetime
+{
+    private AdvisoryLock theLock;
+
+    public Task InitializeAsync()
+    {
+
+        theLock = new AdvisoryLock(() => new SqlConnection(ConnectionSource.ConnectionString), NullLogger.Instance, "Testing");
+        return Task.CompletedTask;
+    }
+
+    public async Task DisposeAsync()
+    {
+        await theLock.DisposeAsync();
+    }
+
+
+    [Fact]
+    public async Task explicitly_release_global_session_locks()
+    {
+        await using var conn2 = new SqlConnection(ConnectionSource.ConnectionString);
+        await using var conn3 = new SqlConnection(ConnectionSource.ConnectionString);
+
+        await conn2.OpenAsync();
+        await conn3.OpenAsync();
+
+        await theLock.TryAttainLockAsync(10, CancellationToken.None);
+
+        // Cannot get the lock here
+        (await conn2.TryGetGlobalLock(10.ToString())).ShouldBeFalse();
+
+        await theLock.ReleaseLockAsync(10);
+
+        for (var j = 0; j < 5; j++)
+        {
+            if ((await conn2.TryGetGlobalLock(10.ToString()))) return;
+
+            await Task.Delay(250);
+        }
+
+        throw new Exception("Advisory lock was not released");
+    }
+
+}