npgsql · manandre · Jan 8, 2023 · Jan 11, 2023 · vonzshik · Jan 19, 2023
@@ -2005,11 +2005,28 @@ copyOperation is NpgsqlCopyTextWriter ||
         }
     }
 
-    // TODO in theory this should be async-optional, but the only I/O done here is the Terminate Flush, which is
-    // very unlikely to block (plus locking would need to be worked out)
     internal void Close()
     {
-        lock (SyncObj)
+        Close(async: false).GetAwaiter().GetResult();
+    }
+
+    readonly SemaphoreSlim _semaphore = new SemaphoreSlim(1, 1);
+
+    internal async Task Close(bool async, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            if (async)
+                await _semaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+            else
+                _semaphore.Wait(cancellationToken);
+        }
+        catch (ObjectDisposedException)
+        {
+            return;
+        }
+
+        try
         {
             if (IsReady)
             {
@@ -2021,7 +2038,7 @@ internal void Close()
                     // see https://github.com/npgsql/npgsql/issues/3592
                     WriteBuffer.Clear();
                     WriteTerminate();
-                    Flush();
+                    await Flush(async, cancellationToken).ConfigureAwait(false);
                 }
                 catch (Exception e)
                 {
@@ -2039,6 +2056,18 @@ internal void Close()
 
             State = ConnectorState.Closed;
         }
+        finally
+        {
+            try
+            {
+                _semaphore.Release();
+                _semaphore.Dispose();
+            }
+            catch (ObjectDisposedException)
+            {
+                // Ignore
+            }
+        }
 
         FullCleanup();
         LogMessages.ClosedPhysicalConnection(ConnectionLogger, Host, Port, Database, UserFacingConnectionString, Id);

@@ -28,7 +28,7 @@ static NpgsqlConnectionStringBuilder CloneSettingsForTargetSessionAttributes(
 
     internal override (int Total, int Idle, int Busy) Statistics => _wrappedSource.Statistics;
 
-    internal override void Clear() => _wrappedSource.Clear();
+    internal override Task Clear(bool async, CancellationToken cancellationToken = default) => _wrappedSource.Clear(async, cancellationToken);
     internal override ValueTask<NpgsqlConnector> Get(NpgsqlConnection conn, NpgsqlTimeout timeout, bool async, CancellationToken cancellationToken)
         => _wrappedSource.Get(conn, timeout, async, cancellationToken);
     internal override bool TryGetIdleConnector([NotNullWhen(true)] out NpgsqlConnector? connector)

@@ -1868,13 +1868,27 @@ public override void ChangeDatabase(string dbName)
     /// </summary>
     public static void ClearPool(NpgsqlConnection connection) => PoolManager.Clear(connection._connectionString);
 
+    /// <summary>
+    /// Clears the connection pool. All idle physical connections in the pool of the given connection are
+    /// immediately closed, and any busy connections which were opened before <see cref="ClearPool"/> was called
+    /// will be closed when returned to the pool.
+    /// </summary>
+    public static Task ClearPoolAsync(NpgsqlConnection connection, CancellationToken cancellationToken = default) => PoolManager.ClearAsync(connection._connectionString, cancellationToken);
+
     /// <summary>
     /// Clear all connection pools. All idle physical connections in all pools are immediately closed, and any busy
     /// connections which were opened before <see cref="ClearAllPools"/> was called will be closed when returned
     /// to their pool.
     /// </summary>
     public static void ClearAllPools() => PoolManager.ClearAll();
 
+    /// <summary>
+    /// Clear all connection pools. All idle physical connections in all pools are immediately closed, and any busy
+    /// connections which were opened before <see cref="ClearAllPools"/> was called will be closed when returned
+    /// to their pool.
+    /// </summary>
+    public static Task ClearAllPoolsAsync(CancellationToken cancellationToken = default) => PoolManager.ClearAllAsync(cancellationToken);
+
     /// <summary>
     /// Unprepares all prepared statements on this connection.
     /// </summary>

@@ -347,7 +347,11 @@ internal abstract ValueTask<NpgsqlConnector> Get(
 
     internal abstract void Return(NpgsqlConnector connector);
 
-    internal abstract void Clear();
+    internal void Clear() => Clear(async: false).GetAwaiter().GetResult();
+
+    internal Task ClearAsync(CancellationToken cancellationToken = default) => Clear(async: true, cancellationToken);
+
+    internal abstract Task Clear(bool async, CancellationToken cancellationToken = default);
 
     internal abstract bool OwnsConnectors { get; }
 
@@ -464,7 +468,6 @@ protected sealed override ValueTask DisposeAsyncCore()
         return default;
     }
 
-#pragma warning disable CS1998
     /// <inheritdoc cref="DisposeAsyncCore" />
     protected virtual async ValueTask DisposeAsyncBase()
     {
@@ -486,10 +489,8 @@ protected virtual async ValueTask DisposeAsyncBase()
 
         _setupMappingsSemaphore.Dispose();
 
-        // TODO: async Clear, #4499
-        Clear();
+        await ClearAsync().ConfigureAwait(false);
     }
-#pragma warning restore CS1998
 
     private protected void CheckDisposed()
     {

@@ -367,10 +367,10 @@ internal override bool TryGetIdleConnector([NotNullWhen(true)] out NpgsqlConnect
     internal override ValueTask<NpgsqlConnector?> OpenNewConnector(NpgsqlConnection conn, NpgsqlTimeout timeout, bool async, CancellationToken cancellationToken)
         => throw new NpgsqlException("Npgsql bug: trying to open a new connector from " + nameof(NpgsqlMultiHostDataSource));
 
-    internal override void Clear()
+    internal override async Task Clear(bool async, CancellationToken cancellationToken = default)
     {
         foreach (var pool in _pools)
-            pool.Clear();
+            await pool.Clear(async, cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>

@@ -1,5 +1,8 @@
-using System;
+using System;
 using System.Collections.Concurrent;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Npgsql;
 
@@ -20,13 +23,27 @@ internal static void Clear(string connString)
             pool.Clear();
     }
 
+    internal static async Task ClearAsync(string connString, CancellationToken cancellationToken = default)
+    {
+        // TODO: Actually remove the pools from here, #3387 (but be careful of concurrency)
+        if (Pools.TryGetValue(connString, out var pool))
+            await pool.ClearAsync(cancellationToken).ConfigureAwait(false);
+    }
+
     internal static void ClearAll()
     {
         // TODO: Actually remove the pools from here, #3387 (but be careful of concurrency)
         foreach (var pool in Pools.Values)
             pool.Clear();
     }
 
+    internal static async Task ClearAllAsync(CancellationToken cancellationToken = default)
+    {
+        // TODO: Actually remove the pools from here, #3387 (but be careful of concurrency)
+        foreach (var pool in Pools.Values)
+            await pool.ClearAsync(cancellationToken).ConfigureAwait(false);
+    }
+
     static PoolManager()
     {
         // When the appdomain gets unloaded (e.g. web app redeployment) attempt to nicely
@@ -45,4 +62,4 @@ internal static void Reset()
         ClearAll();
         Pools.Clear();
     }
-}
+}
@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
@@ -169,8 +169,8 @@ async ValueTask<NpgsqlConnector> RentAsync(
                             }
                         }
 
-                        if (CheckIdleConnector(connector))
-                            return connector;
+                        if (await CheckIdleConnector(connector, async, finalToken).ConfigureAwait(false))
+                            return connector!;
                     }
                     catch (OperationCanceledException)
                     {
@@ -217,8 +217,13 @@ internal sealed override bool TryGetIdleConnector([NotNullWhen(true)] out Npgsql
         return false;
     }
 
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
     bool CheckIdleConnector([NotNullWhen(true)] NpgsqlConnector? connector)
+    {
+        return CheckIdleConnector(connector, async: false).GetAwaiter().GetResult();
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    async Task<bool> CheckIdleConnector(NpgsqlConnector? connector, bool async, CancellationToken cancellationToken = default)
     {
         if (connector is null)
             return false;
@@ -232,14 +237,14 @@ bool CheckIdleConnector([NotNullWhen(true)] NpgsqlConnector? connector)
         // if keepalive isn't turned on.
         if (connector.IsBroken)
         {
-            CloseConnector(connector);
+            await CloseConnector(connector, async, cancellationToken).ConfigureAwait(false);
             return false;
         }
 
         if (_connectionLifetime != TimeSpan.Zero && DateTime.UtcNow > connector.OpenTimestamp + _connectionLifetime)
         {
             LogMessages.ConnectionExceededMaximumLifetime(_logger, _connectionLifetime, connector.Id);
-            CloseConnector(connector);
+            await CloseConnector(connector, async, cancellationToken).ConfigureAwait(false);
             return false;
         }
 
@@ -340,7 +345,7 @@ internal sealed override void Return(NpgsqlConnector connector)
         Debug.Assert(written);
     }
 
-    internal override void Clear()
+    internal override async Task Clear(bool async, CancellationToken cancellationToken = default)
     {
         Interlocked.Increment(ref _clearCounter);
 
@@ -352,9 +357,9 @@ internal override void Clear()
             var count = _idleCount;
             while (count > 0 && _idleConnectorReader.TryRead(out var connector))
             {
-                if (CheckIdleConnector(connector))
+                if (await CheckIdleConnector(connector, async, cancellationToken).ConfigureAwait(false))
                 {
-                    CloseConnector(connector);
+                    await CloseConnector(connector!, async, cancellationToken).ConfigureAwait(false);
                     count--;
                 }
             }
@@ -366,10 +371,15 @@ internal override void Clear()
     }
 
     void CloseConnector(NpgsqlConnector connector)
+    {
+        CloseConnector(connector, async: false).GetAwaiter().GetResult();
+    }
+
+    async Task CloseConnector(NpgsqlConnector connector, bool async, CancellationToken cancellationToken = default)
     {
         try
         {
-            connector.Close();
+            await connector.Close(async, cancellationToken).ConfigureAwait(false);
         }
         catch (Exception exception)
         {

@@ -1,4 +1,4 @@
-#nullable enable
+#nullable enable
 const Npgsql.PostgresErrorCodes.IdleSessionTimeout = "57P05" -> string!
 Npgsql.ChannelBinding
 Npgsql.ChannelBinding.Disable = 0 -> Npgsql.ChannelBinding
@@ -129,3 +129,5 @@ static NpgsqlTypes.NpgsqlCidr.implicit operator NpgsqlTypes.NpgsqlInet(NpgsqlTyp
 *REMOVED*Npgsql.PostgresTypes.PostgresType.PostgresType(string! ns, string! name, string! internalName, uint oid) -> void
 *REMOVED*Npgsql.PostgresTypes.PostgresType.PostgresType(string! ns, string! name, uint oid) -> void
 *REMOVED*Npgsql.TypeMapping.INpgsqlTypeMapper.AddTypeResolverFactory(Npgsql.Internal.TypeHandling.TypeHandlerResolverFactory! resolverFactory) -> void
+static Npgsql.NpgsqlConnection.ClearAllPoolsAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!
+static Npgsql.NpgsqlConnection.ClearPoolAsync(Npgsql.NpgsqlConnection! connection, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!
@@ -46,5 +46,5 @@ internal override void Return(NpgsqlConnector connector)
         connector.Close();
     }
 
-    internal override void Clear() {}
+    internal override Task Clear(bool async, CancellationToken cancellationToken = default) => Task.CompletedTask;
 }