Fix cancellation being able to cancel prepended queries (#4907)

vonzshik · vonzshik · commit 23b61f535d0e · 2023-01-28T22:39:43.000+03:00
Fixes #4906 (cherry picked from commit 218a160)
diff --git a/src/Npgsql/Internal/NpgsqlConnector.cs b/src/Npgsql/Internal/NpgsqlConnector.cs
@@ -139,6 +139,13 @@ public sealed partial class NpgsqlConnector : IDisposable
     /// </summary>
     internal int PendingPrependedResponses { get; set; }
 
+    /// <summary>
+    /// A ManualResetEventSlim used to make sure a cancellation request doesn't run
+    /// while we're reading responses for the prepended query
+    /// as we can't gracefully handle their cancellation.
+    /// </summary>
+    readonly ManualResetEventSlim ReadingPrependedMessagesMRE = new(initialState: true);
+
     internal NpgsqlDataReader? CurrentReader;
 
     internal PreparedStatementManager PreparedStatementManager { get; }
@@ -221,7 +228,20 @@ internal void FlagAsWritableForMultiplexing()
     /// cancellation is delivered. This reduces the chance that a cancellation meant for a previous
     /// command will accidentally cancel a later one, see #615.
     /// </summary>
-    internal object CancelLock { get; }
+    object CancelLock { get; } = new();
+
+    /// <summary>
+    /// A lock that's taken to make sure no other concurrent operation is running.
+    /// Break takes it to set the state of the connector.
+    /// Anyone else should immediately check the state and exit
+    /// if the connector is closed.
+    /// </summary>
+    object SyncObj { get; } = new();
+
+    /// <summary>
+    /// A lock that's used to wait for the Cleanup to complete while breaking the connection.
+    /// </summary>
+    object CleanupLock { get; } = new();
 
     readonly bool _isKeepAliveEnabled;
     readonly Timer? _keepAliveTimer;
@@ -341,8 +361,6 @@ internal NpgsqlConnector(ConnectorSource connectorSource, NpgsqlConnection conn)
         Settings = connectorSource.Settings;
         PostgresParameters = new Dictionary<string, string>();
 
-        CancelLock = new object();
-
         _isKeepAliveEnabled = Settings.KeepAlive > 0;
         if (_isKeepAliveEnabled)
         {
@@ -504,7 +522,7 @@ internal async Task Open(NpgsqlTimeout timeout, bool async, CancellationToken ca
                 // Start the keep alive mechanism to work by scheduling the timer.
                 // Otherwise, it doesn't work for cases when no query executed during
                 // the connection lifetime in case of a new connector.
-                lock (this)
+                lock (SyncObj)
                 {
                     var keepAlive = Settings.KeepAlive * 1000;
                     _keepAliveTimer!.Change(keepAlive, keepAlive);
@@ -1298,9 +1316,14 @@ internal ValueTask<IBackendMessage> ReadMessage(bool async, DataRowLoadingMode d
                     connector.ReadBuffer.Timeout = TimeSpan.FromMilliseconds(connector.InternalCommandTimeout);
                     for (; connector.PendingPrependedResponses > 0; connector.PendingPrependedResponses--)
                         await ReadMessageLong(connector, async, DataRowLoadingMode.Skip, readingNotifications: false, isReadingPrependedMessage: true);
+                    // We've read all the prepended response.
+                    // Allow cancellation to proceed.
+                    connector.ReadingPrependedMessagesMRE.Set();
                 }
-                catch (PostgresException e)
+                catch (Exception e)
                 {
+                    // Prepended queries should never fail.
+                    // If they do, we're not even going to attempt to salvage the connector.
                     throw connector.Break(e);
                 }
             }
@@ -1663,18 +1686,39 @@ static RemoteCertificateValidationCallback SslRootValidation(string certRootPath
 
     #region Cancel
 
+    internal void ResetCancellation()
+    {
+        // If a cancellation is in progress, wait for it to "complete" before proceeding (#615)
+        lock (CancelLock)
+        {
+            if (PendingPrependedResponses > 0)
+                ReadingPrependedMessagesMRE.Reset();
+            Debug.Assert(ReadingPrependedMessagesMRE.IsSet || PendingPrependedResponses > 0);
+        }
+    }
+
     internal void PerformUserCancellation()
     {
         var connection = Connection;
         if (connection is null || connection.ConnectorBindingScope == ConnectorBindingScope.Reader)
             return;
 
-        // There's a subtle race condition where cancellation may be happening just as Break is called. Break takes the connector lock, and
-        // then ends the user action; this disposes the cancellation token registration, which waits until the cancellation callback
-        // completes. But the callback needs to take the connector lock below, which led to a deadlock (#4654).
-        // As a result, Break takes CancelLock, and we abort the cancellation attempt immediately if we can't get it here.
-        if (!Monitor.TryEnter(CancelLock))
-            return;
+        // Take the lock first to make sure there is no concurrent Break.
+        // We should be safe to take it as Break only take it to set the state.
+        lock (SyncObj)
+        {
+            // The connector is dead, exit gracefully.
+            if (!IsConnected)
+                return;
+            // The connector is still alive, take the CancelLock before exiting SingleUseLock.
+            // If a break will happen after, it's going to wait for the cancellation to complete.
+            Monitor.Enter(CancelLock);
+        }
+
+        // Wait before we've read all responses for the prepended queries
+        // as we can't gracefully handle their cancellation.
+        // Break makes sure that it's going to be set even if we fail while reading them.
+        ReadingPrependedMessagesMRE.Wait();
 
         try
         {
@@ -1687,28 +1731,18 @@ internal void PerformUserCancellation()
                 {
                     if (cancellationTimeout > 0)
                     {
-                        lock (this)
-                        {
-                            if (!IsConnected)
-                                return;
-                            UserTimeout = cancellationTimeout;
-                            ReadBuffer.Timeout = TimeSpan.FromMilliseconds(cancellationTimeout);
-                            ReadBuffer.Cts.CancelAfter(cancellationTimeout);
-                        }
+                        UserTimeout = cancellationTimeout;
+                        ReadBuffer.Timeout = TimeSpan.FromMilliseconds(cancellationTimeout);
+                        ReadBuffer.Cts.CancelAfter(cancellationTimeout);
                     }
 
                     return;
                 }
             }
 
-            lock (this)
-            {
-                if (!IsConnected)
-                    return;
-                UserTimeout = -1;
-                ReadBuffer.Timeout = _cancelImmediatelyTimeout;
-                ReadBuffer.Cts.Cancel();
-            }
+            UserTimeout = -1;
+            ReadBuffer.Timeout = _cancelImmediatelyTimeout;
+            ReadBuffer.Cts.Cancel();
         }
         finally
         {
@@ -1783,8 +1817,7 @@ void DoCancelRequest(int backendProcessId, int backendSecretKey)
         }
         finally
         {
-            lock (this)
-                FullCleanup();
+            FullCleanup();
         }
     }
 
@@ -1890,7 +1923,7 @@ copyOperation is NpgsqlCopyTextWriter ||
     // very unlikely to block (plus locking would need to be worked out)
     internal void Close()
     {
-        lock (this)
+        lock (SyncObj)
         {
             Log.Trace("Closing connector", Id);
 
@@ -1920,8 +1953,9 @@ internal void Close()
             }
 
             State = ConnectorState.Closed;
-            FullCleanup();
         }
+
+        FullCleanup();
     }
 
     internal bool TryRemovePendingEnlistedConnector(Transaction transaction)
@@ -1950,11 +1984,43 @@ internal Exception Break(Exception reason)
     {
         Debug.Assert(!IsClosed);
 
-        // See PerformUserCancellation on why we take CancelLock
-        lock (CancelLock)
-        lock (this)
+        Monitor.Enter(SyncObj);
+
+        if (State == ConnectorState.Broken)
         {
-            if (State != ConnectorState.Broken)
+            // We're already broken.
+            // Exit SingleUseLock to unblock other threads (like cancellation).
+            Monitor.Exit(SyncObj);
+            // Wait for the break to complete before going forward.
+            lock (CleanupLock) { }
+            return reason;
+        }
+        
+        try
+        {
+            // If we're broken while reading prepended messages
+            // the cancellation request might still be waiting on the MRE.
+            // Unblock it.
+            ReadingPrependedMessagesMRE.Set();
+
+            Log.Error("Breaking connector", reason, Id);
+
+            // Note that we may be reading and writing from the same connector concurrently, so safely set
+            // the original reason for the break before actually closing the socket etc.
+            Interlocked.CompareExchange(ref _breakReason, reason, null);
+            State = ConnectorState.Broken;
+            // Take the CleanupLock while in SingleUseLock to make sure concurrent Break doesn't take it first.
+            Monitor.Enter(CleanupLock);
+        }
+        finally
+        {
+            // Unblock other threads (like cancellation) to proceed and exit gracefully.
+            Monitor.Exit(SyncObj);
+        }
+
+        try
+        {
+            lock (CancelLock)
             {
                 // Note we only set the cluster to offline and clear the pool if the connection is being broken (we're in this method),
                 // *and* the exception indicates that the PG cluster really is down; the latter includes any IO/timeout issue,
@@ -1968,13 +2034,6 @@ internal Exception Break(Exception reason)
                     _connectorSource.Clear();
                 }
 
-                Log.Error("Breaking connector", reason, Id);
-
-                // Note that we may be reading and writing from the same connector concurrently, so safely set
-                // the original reason for the break before actually closing the socket etc.
-                Interlocked.CompareExchange(ref _breakReason, reason, null);
-                State = ConnectorState.Broken;
-
                 var connection = Connection;
 
                 FullCleanup();
@@ -1992,41 +2051,47 @@ internal Exception Break(Exception reason)
                         connection.Connector = null;
                         connection.ConnectorBindingScope = ConnectorBindingScope.None;
                     }
+
                     connection.FullState = ConnectionState.Broken;
                     connection.ReleaseCloseLock();
                 }
             }
 
             return reason;
         }
+        finally
+        {
+            Monitor.Exit(CleanupLock);
+        }
     }
         
     void FullCleanup()
     {
-        Debug.Assert(Monitor.IsEntered(this));
-
-        if (Settings.Multiplexing)
+        lock (CleanupLock)
         {
-            FlagAsNotWritableForMultiplexing();
+            if (Settings.Multiplexing)
+            {
+                FlagAsNotWritableForMultiplexing();
 
-            // Note that in multiplexing, this could be called from the read loop, while the write loop is
-            // writing into the channel. To make sure this race condition isn't a problem, the channel currently
-            // isn't set up with SingleWriter (since at this point it doesn't do anything).
-            CommandsInFlightWriter!.Complete();
+                // Note that in multiplexing, this could be called from the read loop, while the write loop is
+                // writing into the channel. To make sure this race condition isn't a problem, the channel currently
+                // isn't set up with SingleWriter (since at this point it doesn't do anything).
+                CommandsInFlightWriter!.Complete();
 
-            // The connector's read loop has a continuation to observe and log any exception coming out
-            // (see Open)
-        }
+                // The connector's read loop has a continuation to observe and log any exception coming out
+                // (see Open)
+            }
 
-        Log.Trace("Cleaning up connector", Id);
-        Cleanup();
+            Log.Trace("Cleaning up connector", Id);
+            Cleanup();
 
-        if (_isKeepAliveEnabled)
-        {
-            _userLock!.Dispose();
-            _userLock = null;
-            _keepAliveTimer!.Change(Timeout.Infinite, Timeout.Infinite);
-            _keepAliveTimer.Dispose();
+            if (_isKeepAliveEnabled)
+            {
+                _userLock!.Dispose();
+                _userLock = null;
+                _keepAliveTimer!.Change(Timeout.Infinite, Timeout.Infinite);
+                _keepAliveTimer.Dispose();
+            }
         }
     }
 
@@ -2270,7 +2335,7 @@ internal UserAction StartUserAction(
         if (!_isKeepAliveEnabled)
             return DoStartUserAction();
 
-        lock (this)
+        lock (SyncObj)
         {
             if (!IsConnected)
             {
@@ -2353,7 +2418,7 @@ internal void EndUserAction()
 
         if (_isKeepAliveEnabled)
         {
-            lock (this)
+            lock (SyncObj)
             {
                 if (IsReady || !IsConnected)
                     return;
@@ -2399,7 +2464,7 @@ void PerformKeepAlive(object? state)
 
         // SemaphoreSlim.Dispose() isn't thread-safe - it may be in progress so we shouldn't try to wait on it;
         // we need a standard lock to protect it.
-        if (!Monitor.TryEnter(this))
+        if (!Monitor.TryEnter(SyncObj))
             return;
 
         try
@@ -2432,7 +2497,7 @@ void PerformKeepAlive(object? state)
         }
         finally
         {
-            Monitor.Exit(this);
+            Monitor.Exit(SyncObj);
         }
     }
 #pragma warning restore CA1801 // Review unused parameters
diff --git a/src/Npgsql/NpgsqlCommand.cs b/src/Npgsql/NpgsqlCommand.cs
@@ -1400,9 +1400,7 @@ internal async ValueTask<NpgsqlDataReader> ExecuteReader(CommandBehavior behavio
                     TraceCommandStart(connector);
 
                     // If a cancellation is in progress, wait for it to "complete" before proceeding (#615)
-                    lock (connector.CancelLock)
-                    {
-                    }
+                    connector.ResetCancellation();
 
                     // We do not wait for the entire send to complete before proceeding to reading -
                     // the sending continues in parallel with the user's reading. Waiting for the
diff --git a/test/Npgsql.Tests/CommandTests.cs b/test/Npgsql.Tests/CommandTests.cs
