If the stream passed to Http1Connection does not support gathered writes, add a WriteBufferingStream on top of it.

Author: scalablecory

NetworkToolkit/Connections/Connection.cs

@@ -61,9 +61,7 @@ public async ValueTask DisposeAsync(CancellationToken cancellationToken)
 
             _stream = null;
 
-            await (stream is ICancellableAsyncDisposable cancellable ?
-                cancellable.DisposeAsync(cancellationToken) :
-                stream.DisposeAsync()).ConfigureAwait(false);
+            await stream.DisposeAsync(cancellationToken).ConfigureAwait(false);
         }
 
         /// <summary>

NetworkToolkit/Connections/FilteringConnection.cs

@@ -36,11 +36,17 @@ public FilteringConnection(Connection baseConnection, Stream stream) : base(stre
         }
 
         /// <inheritdoc/>
-        protected override ValueTask DisposeAsyncCore(CancellationToken cancellationToken) =>
-            BaseConnection.DisposeAsync(cancellationToken);
+        protected override async ValueTask DisposeAsyncCore(CancellationToken cancellationToken)
+        {
+            await Stream.DisposeAsync(cancellationToken).ConfigureAwait(false);
+            await BaseConnection.DisposeAsync(cancellationToken).ConfigureAwait(false);
+        }
 
         /// <inheritdoc/>
-        public override ValueTask CompleteWritesAsync(CancellationToken cancellationToken) =>
-            BaseConnection.CompleteWritesAsync(cancellationToken);
+        public override async ValueTask CompleteWritesAsync(CancellationToken cancellationToken)
+        {
+            await Stream.FlushAsync(cancellationToken).ConfigureAwait(false);
+            await BaseConnection.CompleteWritesAsync(cancellationToken).ConfigureAwait(false);
+        }
     }
 }

NetworkToolkit/Connections/MemoryConnection.cs

@@ -56,7 +56,7 @@ public override ValueTask CompleteWritesAsync(CancellationToken cancellationToke
 
         private sealed class MemoryConnectionStream : Stream, IGatheringStream
         {
-            PipeReader _reader;
+            PipeReader? _reader;
             PipeWriter? _writer;
 
             public override bool CanRead => true;
@@ -77,12 +77,12 @@ public MemoryConnectionStream(PipeReader reader, PipeWriter writer)
 
             protected override void Dispose(bool disposing)
             {
-                if (disposing && _writer != null)
+                if (disposing && _reader is PipeReader reader)
                 {
-                    _writer.Complete();
-                    _reader.Complete();
-                    _writer = null!;
-                    _reader = null!;
+                    _writer?.Complete();
+                    reader.Complete();
+                    _writer = null;
+                    _reader = null;
                 }
             }
 
@@ -106,13 +106,13 @@ public override int Read(byte[] buffer, int offset, int count)
 
             public override int Read(Span<byte> buffer)
             {
-                if (_reader == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_reader is not PipeReader reader) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
 
                 try
                 {
-                    return FinishRead(buffer, Tools.BlockForResult(_reader.ReadAsync()));
+                    return FinishRead(reader, buffer, Tools.BlockForResult(_reader.ReadAsync()), CancellationToken.None);
                 }
-                catch (Exception ex)
+                catch (Exception ex) when (ex is not OperationCanceledException)
                 {
                     throw new IOException(ex.Message, ex);
                 }
@@ -129,24 +129,25 @@ public override Task<int> ReadAsync(byte[] buffer, int offset, int count, Cancel
 
             public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
             {
-                if (_reader == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_reader is not PipeReader reader) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
 
                 try
                 {
-                    ReadResult result = await _reader.ReadAsync(cancellationToken).ConfigureAwait(false);
-                    return FinishRead(buffer.Span, result);
+                    ReadResult result = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
+                    return FinishRead(reader, buffer.Span, result, cancellationToken);
                 }
-                catch (Exception ex)
+                catch (Exception ex) when(ex is not OperationCanceledException)
                 {
                     throw new IOException(ex.Message, ex);
                 }
             }
 
-            private int FinishRead(Span<byte> buffer, in ReadResult result)
+            private static int FinishRead(PipeReader reader, Span<byte> buffer, in ReadResult result, CancellationToken cancellationToken)
             {
                 if (result.IsCanceled)
                 {
-                    throw new SocketException((int)SocketError.OperationAborted);
+                    cancellationToken.ThrowIfCancellationRequested();
+                    throw new OperationCanceledException();
                 }
 
                 ReadOnlySequence<byte> sequence = result.Buffer;
@@ -171,7 +172,7 @@ private int FinishRead(Span<byte> buffer, in ReadResult result)
                 }
                 finally
                 {
-                    _reader.AdvanceTo(consumed);
+                    reader.AdvanceTo(consumed);
                 }
             }
 
@@ -180,7 +181,8 @@ public override void Write(byte[] buffer, int offset, int count) =>
 
             public override void Write(ReadOnlySpan<byte> buffer)
             {
-                if (_writer == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_reader == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_writer == null) throw new InvalidOperationException($"{nameof(MemoryConnectionStream)} cannot be written to after writes have been completed.");
 
                 try
                 {
@@ -191,10 +193,10 @@ public override void Write(ReadOnlySpan<byte> buffer)
 
                     if (res.IsCanceled)
                     {
-                        throw new SocketException((int)SocketError.OperationAborted);
+                        throw new OperationCanceledException();
                     }
                 }
-                catch (Exception ex)
+                catch (Exception ex) when (ex is not OperationCanceledException)
                 {
                     throw new IOException(ex.Message, ex);
                 }
@@ -205,48 +207,58 @@ public override Task WriteAsync(byte[] buffer, int offset, int count, Cancellati
 
             public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
             {
-                if (_writer == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_reader is null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_writer is not PipeWriter writer) throw new InvalidOperationException($"{nameof(MemoryConnectionStream)} cannot be written to after writes have been completed.");
 
                 try
                 {
-                    FlushResult res = await _writer.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
+                    FlushResult res = await writer.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
 
                     if (res.IsCanceled)
                     {
-                        throw new SocketException((int)SocketError.OperationAborted);
+                        cancellationToken.ThrowIfCancellationRequested();
+                        throw new OperationCanceledException();
                     }
                 }
-                catch (Exception ex)
+                catch (Exception ex) when (ex is not OperationCanceledException)
                 {
                     throw new IOException(ex.Message, ex);
                 }
             }
 
             public async ValueTask WriteAsync(IReadOnlyList<ReadOnlyMemory<byte>> buffers, CancellationToken cancellationToken = default)
             {
-                if (_writer == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_reader == null) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                if (_writer is not PipeWriter writer) throw new InvalidOperationException($"{nameof(MemoryConnectionStream)} cannot be written to after writes have been completed.");
 
                 try
                 {
                     foreach (ReadOnlyMemory<byte> buffer in buffers)
                     {
-                        buffer.Span.CopyTo(_writer.GetSpan(buffer.Length));
-                        _writer.Advance(buffer.Length);
+                        buffer.Span.CopyTo(writer.GetSpan(buffer.Length));
+                        writer.Advance(buffer.Length);
                     }
 
-                    FlushResult res = await _writer.FlushAsync(cancellationToken).ConfigureAwait(false);
+                    FlushResult res = await writer.FlushAsync(cancellationToken).ConfigureAwait(false);
 
                     if (res.IsCanceled)
                     {
-                        throw new SocketException((int)SocketError.OperationAborted);
+                        cancellationToken.ThrowIfCancellationRequested();
+                        throw new OperationCanceledException();
                     }
                 }
-                catch (Exception ex)
+                catch (Exception ex) when (ex is not OperationCanceledException)
                 {
                     throw new IOException(ex.Message, ex);
                 }
             }
 
+            public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state) =>
+                TaskToApm.Begin(WriteAsync(buffer, offset, count), callback, state);
+
+            public override void EndWrite(IAsyncResult asyncResult) =>
+                TaskToApm.End(asyncResult);
+
             public override void Flush()
             {
             }
@@ -265,6 +277,22 @@ public override void SetLength(long value)
             {
                 throw new NotImplementedException();
             }
+
+            public override void CopyTo(Stream destination, int bufferSize) =>
+                CopyToAsync(destination, bufferSize, CancellationToken.None).GetAwaiter().GetResult();
+
+            public override async Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken)
+            {
+                if (_reader is not PipeReader reader) throw new ObjectDisposedException(nameof(MemoryConnectionStream));
+                try
+                {
+                    await reader.CopyToAsync(destination, cancellationToken).ConfigureAwait(false);
+                }
+                catch (Exception ex) when (ex is not OperationCanceledException)
+                {
+                    throw new IOException(ex.Message, ex);
+                }
+            }
         }
     }
 }

NetworkToolkit/Connections/SocketConnectionFactory.cs

@@ -219,10 +219,15 @@ private sealed class SocketConnection : Connection
 
             private Socket Socket => ((NetworkStream)Stream).Socket;
 
-            public SocketConnection(Socket socket) : base(new GatheringNetworkStream(socket))
+            public SocketConnection(Socket socket) : base(CreateStream(socket))
             {
             }
 
+            private static NetworkStream CreateStream(Socket socket) =>
+                GatheringNetworkStream.IsSupported
+                    ? new GatheringNetworkStream(socket)
+                    : new NetworkStream(socket);
+
             protected override ValueTask DisposeAsyncCore(CancellationToken cancellationToken)
                 => default;
 

NetworkToolkit/Connections/SslConnectionFactory.cs

@@ -118,12 +118,6 @@ public SslConnection(Connection baseConnection, SslStream stream) : base(baseCon
             {
             }
 
-            protected override async ValueTask DisposeAsyncCore(CancellationToken cancellationToken)
-            {
-                await Stream.DisposeAsync().ConfigureAwait(false);
-                await base.DisposeAsyncCore(cancellationToken).ConfigureAwait(false);
-            }
-
             public override bool TryGetProperty(Type type, out object? value)
             {
                 if (type == typeof(SslStream))

NetworkToolkit/Connections/WriteBufferingConnectionFactory.cs

@@ -0,0 +1,49 @@
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace NetworkToolkit.Connections
+{
+    /// <summary>
+    /// A connection factory that adds write buffering to underlying connections.
+    /// </summary>
+    public sealed class WriteBufferingConnectionFactory : FilteringConnectionFactory
+    {
+        /// <summary>
+        /// Instantiates a new <see cref="WriteBufferingConnectionFactory"/>.
+        /// </summary>
+        /// <param name="baseFactory">The underlying factory that will have write buffering added to its connections.</param>
+        public WriteBufferingConnectionFactory(ConnectionFactory baseFactory) : base(baseFactory)
+        {
+        }
+
+        /// <inheritdoc/>
+        public override async ValueTask<Connection> ConnectAsync(EndPoint endPoint, IConnectionProperties? options = null, CancellationToken cancellationToken = default)
+        {
+            Connection con = await BaseFactory.ConnectAsync(endPoint, options, cancellationToken).ConfigureAwait(false);
+            return new FilteringConnection(con, new WriteBufferingStream(con.Stream));
+        }
+
+        /// <inheritdoc/>
+        public override async ValueTask<ConnectionListener> ListenAsync(EndPoint? endPoint = null, IConnectionProperties? options = null, CancellationToken cancellationToken = default)
+        {
+            ConnectionListener listener = await BaseFactory.ListenAsync(endPoint, options, cancellationToken).ConfigureAwait(false);
+            return new WriteBufferingConnectionListener(listener);
+        }
+
+        private sealed class WriteBufferingConnectionListener : FilteringConnectionListener
+        {
+            public WriteBufferingConnectionListener(ConnectionListener baseListener) : base(baseListener)
+            {
+            }
+
+            public override async ValueTask<Connection?> AcceptConnectionAsync(IConnectionProperties? options = null, CancellationToken cancellationToken = default)
+            {
+                Connection? con = await BaseListener.AcceptConnectionAsync(options, cancellationToken).ConfigureAwait(false);
+                if (con == null) return con;
+
+                return new FilteringConnection(con, new WriteBufferingStream(con.Stream));
+            }
+        }
+    }
+}

NetworkToolkit/GatheringNetworkStream.cs

@@ -1,8 +1,10 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Net.Sockets;
+using System.Reflection;
 using System.Runtime.ExceptionServices;
 using System.Runtime.InteropServices;
 using System.Threading;
@@ -13,9 +15,25 @@ namespace NetworkToolkit
     internal sealed class GatheringNetworkStream : NetworkStream, IGatheringStream
     {
         private EventArgs? _gatheredEventArgs;
+        private static Func<Socket, SocketAsyncEventArgs, CancellationToken, bool>? s_sendAsyncWithCancellation;
+
+        public static bool IsSupported => s_sendAsyncWithCancellation != null;
+
+        static GatheringNetworkStream()
+        {
+            MethodInfo? sendAsync = typeof(Socket).GetMethod("SendAsync", BindingFlags.NonPublic | BindingFlags.Instance, binder: null, new[] { typeof(SocketAsyncEventArgs), typeof(CancellationToken) }, modifiers: null);
+
+            if (sendAsync != null)
+            {
+                s_sendAsyncWithCancellation =
+                    (Func<Socket, SocketAsyncEventArgs, CancellationToken, bool>)
+                    Delegate.CreateDelegate(typeof(Func<Socket, SocketAsyncEventArgs, CancellationToken, bool>), firstArgument: null, sendAsync);
+            }
+        }
 
         public GatheringNetworkStream(Socket socket) : base(socket, ownsSocket: true)
         {
+            Debug.Assert(IsSupported);
         }
 
         protected override void Dispose(bool disposing)
@@ -35,6 +53,12 @@ public ValueTask WriteAsync(IReadOnlyList<ReadOnlyMemory<byte>> buffers, Cancell
                 return WriteAsync(buffers[0], cancellationToken);
             }
 
+            if (cancellationToken.IsCancellationRequested)
+            {
+                // There is no SocketAsyncEventArgs call that is cancellable...
+                return ValueTask.FromException(ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceledException(cancellationToken)));
+            }
+
             _gatheredEventArgs ??= new EventArgs();
             return _gatheredEventArgs.WriteAsync(Socket, buffers, cancellationToken);
         }
@@ -44,7 +68,7 @@ private sealed class EventArgs : SocketTaskEventArgs<int>
             private List<ArraySegment<byte>>? _gatheredSegments;
             private List<byte[]>? _pooledArrays;
 
-            public ValueTask WriteAsync(Socket socket, IReadOnlyList<ReadOnlyMemory<byte>> buffers, CancellationToken cancellationToken = default)
+            public ValueTask WriteAsync(Socket socket, IReadOnlyList<ReadOnlyMemory<byte>> buffers, CancellationToken cancellationToken)
             {
                 int bufferCount = buffers.Count;
 
@@ -71,7 +95,7 @@ public ValueTask WriteAsync(Socket socket, IReadOnlyList<ReadOnlyMemory<byte>> b
 
                 BufferList = _gatheredSegments;
                 Reset();
-                if (!socket.SendAsync(this))
+                if (!s_sendAsyncWithCancellation!(socket, this, cancellationToken))
                 {
                     OnCompleted();
                 }
@@ -104,7 +128,8 @@ public void OnCompleted()
                 }
                 else
                 {
-                    SetException(ExceptionDispatchInfo.SetCurrentStackTrace(new IOException($"{nameof(WriteAsync)} failed. See InnerException for more details.", new SocketException((int)SocketError))));
+                    var ex = new SocketException((int)SocketError);
+                    SetException(ExceptionDispatchInfo.SetCurrentStackTrace(new IOException(ex.Message, ex)));
                 }
             }
         }