v4: a close with commands in flight no longer completes a call with another call's reply (#3329)

src/RESPite/Operations/RespConnection.cs

@@ -2,6 +2,7 @@
 using System.Buffers;
 using System.Buffers.Binary;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using RESPite.Buffers;
@@ -57,6 +58,21 @@ internal class RespConnection : TransportReceiver, IAsyncDisposable
     /// <summary>Operations written and awaiting a reply, in wire order.</summary>
     private readonly ConcurrentQueue<PendingEntry> _pending = new();
 
+    /// <summary>
+    /// Serialises the two things that take from <see cref="_pending"/>: the reader matching a reply, and
+    /// <see cref="Close"/> failing what is left.
+    /// </summary>
+    /// <remarks>
+    /// <b>Concurrent is not the same as coordinated.</b> The queue itself is safe to dequeue from two threads; what
+    /// it cannot promise is which entry each gets. Close draining it while the reader was still matching replies
+    /// from the receive buffer took entry N, and the reader then took entry N+1 for reply N: that command completed,
+    /// successfully, with another command's answer, and every reply after it was off by one - 24% of the in-flight
+    /// commands on a <c>CloseAsync(allowCommandsToComplete: false)</c> (#3329). So: the reader checks closed and
+    /// dequeues as one step under this lock, Close sets closed and then empties the queue under it, and from then
+    /// on a reply that arrives has nothing to match and dies with the connection.
+    /// </remarks>
+    private readonly object _pendingLock = new();
+
     /// <summary>A queued operation, and WHICH life of it was written.</summary>
     /// <remarks>
     /// <b>The token is the point.</b> An operation that ends early - cancelled, timed out, faulted - keeps its
@@ -1391,47 +1407,48 @@ private void Drain()
                 continue;
             }
 
-            if (_pending.TryDequeue(out var entry))
+            if (!TryTakePending(out var entry, out var closed))
             {
-                // offered first: a redirect is an instruction, not an answer, and completing the
-                // operation with it would report a failure for something routine. Only to the life that was
-                // written: a stale entry's frame is simply consumed (see PendingEntry)
-                if (!(entry.IsCurrent && TryHandOff(frame, entry.Message)))
+                // Close owns the queue now and is failing what was in it (see _pendingLock); matching anything
+                // here would hand a reply to the wrong command
+                if (closed) return;
+
+                // a reply with nothing pending is a protocol break; there is no correct recovery, and
+                // pretending otherwise would mis-address every reply after it
+                Close(new InvalidOperationException("A reply arrived with no operation pending."));
+                return;
+            }
+
+            // offered first: a redirect is an instruction, not an answer, and completing the
+            // operation with it would report a failure for something routine. Only to the life that was
+            // written: a stale entry's frame is simply consumed (see PendingEntry)
+            if (!(entry.IsCurrent && TryHandOff(frame, entry.Message)))
+            {
+                if (entry.SequenceChecked)
                 {
-                    if (entry.SequenceChecked)
-                    {
-                        // a checked reply waits for its check; a stale life has nobody waiting, so its frame is
-                        // simply consumed and only the check is validated
-                        switch (entry.IsCurrent ? LookAhead(in entry, frame, buffer) : LookAheadResult.NotYet)
-                        {
-                            case LookAheadResult.Closed: return;
-                            case LookAheadResult.Delivered: break;
-                            default:
-                                if (entry.IsCurrent) Park(in entry, frame, buffer);
-                                _awaitingSequenceCheck = true;
-                                break;
-                        }
-                    }
-                    else
+                    // a checked reply waits for its check; a stale life has nobody waiting, so its frame is
+                    // simply consumed and only the check is validated
+                    switch (entry.IsCurrent ? LookAhead(in entry, frame, buffer) : LookAheadResult.NotYet)
                     {
-                        // the buffer goes along with the bytes, so a message whose result IS the frame can
-                        // retain it rather than copy it out
-                        entry.Message.TrySetResult(entry.Token, frame, buffer);
+                        case LookAheadResult.Closed: return;
+                        case LookAheadResult.Delivered: break;
+                        default:
+                            if (entry.IsCurrent) Park(in entry, frame, buffer);
+                            _awaitingSequenceCheck = true;
+                            break;
                     }
                 }
-                else if (entry.SequenceChecked)
+                else
                 {
-                    // handed off (a redirect, a receipt): nothing to hold, but the check still follows
-                    _awaitingSequenceCheck = true;
+                    // the buffer goes along with the bytes, so a message whose result IS the frame can
+                    // retain it rather than copy it out
+                    entry.Message.TrySetResult(entry.Token, frame, buffer);
                 }
             }
-
-            // a reply with nothing pending is a protocol break; there is no correct recovery, and
-            // pretending otherwise would mis-address every reply after it
-            else
+            else if (entry.SequenceChecked)
             {
-                Close(new InvalidOperationException("A reply arrived with no operation pending."));
-                return;
+                // handed off (a redirect, a receipt): nothing to hold, but the check still follows
+                _awaitingSequenceCheck = true;
             }
         }
 
@@ -1447,6 +1464,23 @@ private void Drain()
         // anyone else is holding the buffer before it moves anything.
     }
 
+    /// <summary>Take the next unanswered operation for the reply in hand; see <see cref="_pendingLock"/>.</summary>
+    /// <returns>False with <paramref name="closed"/> set when the connection is closed, false when nothing is pending.</returns>
+    private bool TryTakePending(out PendingEntry entry, out bool closed)
+    {
+        lock (_pendingLock)
+        {
+            closed = Volatile.Read(ref _closed) != 0;
+            if (closed)
+            {
+                entry = default;
+                return false;
+            }
+
+            return _pending.TryDequeue(out entry);
+        }
+    }
+
     private void Close(Exception? fault)
     {
         if (Interlocked.Exchange(ref _closed, 1) != 0) return;
@@ -1458,12 +1492,27 @@ private void Close(Exception? fault)
         var buffer = Interlocked.Exchange(ref _inbound, null);
         buffer?.Release();
 
+        // Taken under the reader's lock (see _pendingLock): with _closed already set, the reader takes nothing
+        // once we hold it, and what we take is everything still unanswered. Failed outside the lock, since a
+        // completion can run whatever continuation is waiting on it.
+        List<PendingEntry>? unanswered = null;
+        lock (_pendingLock)
+        {
+            while (_pending.TryDequeue(out var entry))
+            {
+                (unanswered ??= new()).Add(entry);
+            }
+        }
+
         // INDEFINITE, and this is the case that distinction exists for: a write may have reached the
         // server and had its reply lost with the connection, so these operations are not provably
         // unapplied and their instances must not go back to a pool
-        while (_pending.TryDequeue(out var entry))
+        if (unanswered is not null)
         {
-            entry.Message.TrySetException(entry.Token, reason, definite: false);
+            foreach (var entry in unanswered)
+            {
+                entry.Message.TrySetException(entry.Token, reason, definite: false);
+            }
         }
 
         // a reply held for a check that will now never arrive is pending in the same sense: answered, unverified
@@ -1504,8 +1553,15 @@ private void Close(Exception? fault)
     /// <returns>The exception pending commands fail with.</returns>
     protected virtual Exception CreateRemoteClosedFault() => new System.IO.IOException("The connection was closed by the remote end.");
 
+    /// <summary>The fault for commands failed by a close this side asked for (dispose), when no other fault is recorded.</summary>
+    /// <remarks>
+    /// The counterpart of <see cref="CreateRemoteClosedFault"/> for our own close: a host whose callers expect a
+    /// particular type for "the connection went away under my command" (as StackExchange.Redis's do) overrides both.
+    /// </remarks>
+    protected virtual Exception CreateClosedFault() => new InvalidOperationException("The connection is closed.");
+
     private Exception ClosedFault()
-        => _fault ?? new InvalidOperationException("The connection is closed.");
+        => _fault ?? CreateClosedFault();
 
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()

src/StackExchange.Redis/Contexts/RespClientConnection.cs

@@ -94,6 +94,21 @@ protected override Exception CreateRemoteClosedFault()
             return new RedisConnectionException(ConnectionFailureType.SocketClosed, CommandFlags.None, message, null, CommandStatus.Unknown);
         }
 
+        /// <inheritdoc/>
+        /// <remarks>
+        /// Our own close - a dispose with commands still in flight - failed those commands with
+        /// <see cref="ConnectionFailureType.ConnectionDisposed"/> under the same shipped type on v3; a caller catching
+        /// <see cref="RedisConnectionException"/> around a shutdown still does, rather than meeting RESPite's own
+        /// <see cref="InvalidOperationException"/>.
+        /// </remarks>
+        protected override Exception CreateClosedFault()
+        {
+            var message = Server?.EndPoint is { } endpoint
+                ? $"This client closed the connection to {Format.ToString(endpoint)}."
+                : "This client closed the connection.";
+            return new RedisConnectionException(ConnectionFailureType.ConnectionDisposed, CommandFlags.None, message, null, CommandStatus.Unknown);
+        }
+
         /// <inheritdoc/>
         /// <remarks>
         /// High-integrity mode's failure, under the shipped type and <see cref="ConnectionFailureType.ResponseIntegrityFailure"/>:

tests/StackExchange.Redis.Tests/CloseInFlightTests.cs

@@ -0,0 +1,70 @@
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace StackExchange.Redis.Tests;
+
+public class CloseInFlightTests(ITestOutputHelper output) : TestBase(output)
+{
+    // #3329: CloseAsync(allowCommandsToComplete: false) with commands in flight races the read loop, still
+    // matching replies already in the buffer, against the failure path draining the sent-awaiting-reply
+    // queue one message at a time. Each reply the read loop matched after the drain had taken a message was
+    // handed to the message *after* the one it belonged to: GET key N completed, successfully, with key N-1's
+    // value. The only acceptable outcomes are the right value or a connection exception.
+    [Fact]
+    public async Task CloseWithCommandsInFlight_NeverCompletesACallWithAnotherCallsReply()
+    {
+        const int Rounds = 20, PerRound = 500;
+        SetExpectedAmbientFailureCount(-1); // every round tears a connection down with work in flight
+        await using var admin = Create();
+        var adminDb = admin.GetDatabase();
+        var prefix = Me();
+        long ok = 0, wrong = 0, failed = 0;
+
+        for (int round = 0; round < Rounds; round++)
+        {
+            var keys = Enumerable.Range(0, PerRound).Select(i => (RedisKey)$"{prefix}:{round}:{i}").ToArray();
+            await Task.WhenAll(keys.Select(k => adminDb.StringSetAsync(k, k.ToString(), TimeSpan.FromMinutes(1))));
+
+            var conn = Create(shared: false);
+            var db = conn.GetDatabase();
+            var calls = keys.Select(k => db.StringGetAsync(k)).ToArray();
+            await conn.CloseAsync(allowCommandsToComplete: false);
+
+            for (int i = 0; i < calls.Length; i++)
+            {
+                string? value;
+                try
+                {
+                    value = await calls[i];
+                }
+                catch (RedisConnectionException)
+                {
+                    failed++;
+                    continue;
+                }
+                catch (ObjectDisposedException)
+                {
+                    failed++;
+                    continue;
+                }
+
+                if (value == keys[i].ToString())
+                {
+                    ok++;
+                }
+                else if (wrong++ < 5)
+                {
+                    Log($"round {round}: GET {keys[i]} returned '{value}'");
+                }
+            }
+
+            await conn.DisposeAsync();
+            await adminDb.KeyDeleteAsync(keys);
+        }
+
+        Log($"ok={ok} wrong={wrong} failed={failed}");
+        Assert.Equal(0, wrong);
+    }
+}