tests: init output and allow time to complete

Author: mgravell

src/StackExchange.Redis/PhysicalConnection.Write.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Buffers;
+using System.Diagnostics;
 using System.IO;
 using System.IO.Pipelines;
 using System.Threading;
@@ -20,28 +21,37 @@ public IBufferWriter<byte> Output
         }
     }
 
-    private void CreateOutputPipe()
+    private Task _writeComplete = Task.CompletedTask;
+
+    private void InitOutput(Stream? stream)
     {
-        if (_ioStream is not { } stream) return;
+        if (stream is null) return;
+        _ioStream = stream;
         var pipe = new Pipe();
         _output = pipe.Writer;
-        _ = Task.Run(() => CopyOutputAsync(this, pipe.Reader, stream));
+        _writeComplete = Task.Run(() => CopyOutputAsync(this, pipe.Reader), OutputCancel);
     }
 
     internal bool HasOutputPipe => _output is not null;
 
-    private static async Task CopyOutputAsync(PhysicalConnection connection, PipeReader from, Stream to)
+    internal Task CompleteOutputAsync(Exception? exception = null)
+    {
+        _output?.Complete(exception);
+        return _writeComplete;
+    }
+
+    private static async Task CopyOutputAsync(PhysicalConnection connection, PipeReader from)
     {
         try
         {
             bool pendingFlush = false;
-            while (true)
+            while (connection._ioStream is { } stream)
             {
                 if (!from.TryRead(out var read))
                 {
                     if (pendingFlush)
                     {
-                        await to.FlushAsync(connection.OutputCancel).ConfigureAwait(false);
+                        await stream.FlushAsync(connection.OutputCancel).ConfigureAwait(false);
                         pendingFlush = false;
                     }
                     read = await from.ReadAsync(connection.OutputCancel).ConfigureAwait(false);
@@ -55,7 +65,7 @@ private static async Task CopyOutputAsync(PhysicalConnection connection, PipeRea
                     {
                         pendingFlush = true;
                         connection.totalBytesSent += segment.Length;
-                        await to.WriteAsync(buffer.First, connection.OutputCancel).ConfigureAwait(false);
+                        await stream.WriteAsync(buffer.First, connection.OutputCancel).ConfigureAwait(false);
                     }
                 }
                 else
@@ -66,7 +76,7 @@ private static async Task CopyOutputAsync(PhysicalConnection connection, PipeRea
                         {
                             pendingFlush = true;
                             connection.totalBytesSent += segment.Length;
-                            await to.WriteAsync(segment, connection.OutputCancel).ConfigureAwait(false);
+                            await stream.WriteAsync(segment, connection.OutputCancel).ConfigureAwait(false);
                         }
                     }
                 }

src/StackExchange.Redis/PhysicalConnection.cs

@@ -110,7 +110,7 @@ public PhysicalConnection(
             _protocol = protocol;
             _bridge = new WeakReference(null);
             _physicalName = name;
-            _ioStream = ioStream;
+            InitOutput(ioStream);
             OnCreateEcho();
         }
         public PhysicalConnection(PhysicalBridge bridge)
@@ -1144,8 +1144,7 @@ static Stream DemandSocketStream(Socket? socket)
                 stream ??= DemandSocketStream(socket);
                 OnWrapForLogging(ref stream, _physicalName);
 
-                _ioStream = stream;
-                CreateOutputPipe();
+                InitOutput(stream);
 
                 log?.LogInformationConnected(bridge.Name);
 

tests/StackExchange.Redis.Tests/RoundTripUnitTests/TestConnection.cs

@@ -49,6 +49,7 @@ internal static async Task<T> ExecuteAsync<T>(
         message.SetSource(box, processor);
         conn.WriteOutbound(message, commandMap, channelPrefix);
         Assert.Equal(TaskStatus.WaitingForActivation, tcs.Task.Status); // should be pending, since we haven't responded yet
+        await conn.CompleteOutputAsync();
 
         // check the request
         conn.AssertOutbound(requestResp);
@@ -59,6 +60,8 @@ internal static async Task<T> ExecuteAsync<T>(
         return await tcs.Task;
     }
 
+    private Task CompleteOutputAsync(Exception? exception = null) => _physical.CompleteOutputAsync(exception);
+
     public TestConnection(
         bool startReading = true,
         ConnectionType connectionType = ConnectionType.Interactive,