Derive SocketAwaitable from SocketAsyncEventArgs (#2592)

* SocketSenderReceiverBase as baseclass for SocketSender and SocketReceiver
* Made SocketAwaitable.GetResult inlineable
* SocketAwaitable renamed to SocketAwaitableEventArgs and inherited from SocketAsyncEventArgs

Thus saving an extra object (`SocketAwaitable`) and two int-fields, so less allocations.

src/Kestrel.Transport.Sockets/Internal/SocketAwaitableEventArgs.cs

@@ -11,22 +11,20 @@ using System.Threading.Tasks;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 {
-    public class SocketAwaitable : ICriticalNotifyCompletion
+    public class SocketAwaitableEventArgs : SocketAsyncEventArgs, ICriticalNotifyCompletion
     {
         private static readonly Action _callbackCompleted = () => { };
 
         private readonly PipeScheduler _ioScheduler;
 
         private Action _callback;
-        private int _bytesTransferred;
-        private SocketError _error;
 
-        public SocketAwaitable(PipeScheduler ioScheduler)
+        public SocketAwaitableEventArgs(PipeScheduler ioScheduler)
         {
             _ioScheduler = ioScheduler;
         }
 
-        public SocketAwaitable GetAwaiter() => this;
+        public SocketAwaitableEventArgs GetAwaiter() => this;
         public bool IsCompleted => ReferenceEquals(_callback, _callbackCompleted);
 
         public int GetResult()
@@ -35,12 +33,17 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 
             _callback = null;
 
-            if (_error != SocketError.Success)
+            if (SocketError != SocketError.Success)
             {
-                throw new SocketException((int)_error);
+                ThrowSocketException(SocketError);
             }
 
-            return _bytesTransferred;
+            return BytesTransferred;
+
+            void ThrowSocketException(SocketError e)
+            {
+                throw new SocketException((int)e);
+            }
         }
 
         public void OnCompleted(Action continuation)
@@ -57,10 +60,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
             OnCompleted(continuation);
         }
 
-        public void Complete(int bytesTransferred, SocketError socketError)
+        public void Complete()
+        {
+            OnCompleted(this);
+        }
+
+        protected override void OnCompleted(SocketAsyncEventArgs _)
         {
-            _error = socketError;
-            _bytesTransferred = bytesTransferred;
             var continuation = Interlocked.Exchange(ref _callback, _callbackCompleted);
 
             if (continuation != null)

src/Kestrel.Transport.Sockets/Internal/SocketReceiver.cs

@@ -7,42 +7,29 @@ using System.Net.Sockets;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 {
-    public class SocketReceiver : IDisposable
+    public sealed class SocketReceiver : SocketSenderReceiverBase
     {
-        private readonly Socket _socket;
+        public SocketReceiver(Socket socket, PipeScheduler scheduler) : base(socket, scheduler)
-        private readonly SocketAsyncEventArgs _eventArgs = new SocketAsyncEventArgs();
-        private readonly SocketAwaitable _awaitable;
-
-        public SocketReceiver(Socket socket, PipeScheduler scheduler)
         {
-            _socket = socket;
-            _awaitable = new SocketAwaitable(scheduler);
-            _eventArgs.UserToken = _awaitable;
-            _eventArgs.Completed += (_, e) => ((SocketAwaitable)e.UserToken).Complete(e.BytesTransferred, e.SocketError);
         }
 
-        public SocketAwaitable ReceiveAsync(Memory<byte> buffer)
+        public SocketAwaitableEventArgs ReceiveAsync(Memory<byte> buffer)
         {
 #if NETCOREAPP2_1
-            _eventArgs.SetBuffer(buffer);
+            _awaitableEventArgs.SetBuffer(buffer);
 #elif NETSTANDARD2_0
             var segment = buffer.GetArray();
 
-            _eventArgs.SetBuffer(segment.Array, segment.Offset, segment.Count);
+            _awaitableEventArgs.SetBuffer(segment.Array, segment.Offset, segment.Count);
 #else
 #error TFMs need to be updated
 #endif
-            if (!_socket.ReceiveAsync(_eventArgs))
+            if (!_socket.ReceiveAsync(_awaitableEventArgs))
             {
-                _awaitable.Complete(_eventArgs.BytesTransferred, _eventArgs.SocketError);
+                _awaitableEventArgs.Complete();
             }
 
-            return _awaitable;
+            return _awaitableEventArgs;
-        }
-
-        public void Dispose()
-        {
-            _eventArgs.Dispose();
         }
     }
 }

src/Kestrel.Transport.Sockets/Internal/SocketSender.cs

@@ -11,23 +11,15 @@ using System.Runtime.InteropServices;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 {
-    public class SocketSender : IDisposable
+    public sealed class SocketSender : SocketSenderReceiverBase
     {
-        private readonly Socket _socket;
-        private readonly SocketAsyncEventArgs _eventArgs = new SocketAsyncEventArgs();
-        private readonly SocketAwaitable _awaitable;
-
         private List<ArraySegment<byte>> _bufferList;
 
-        public SocketSender(Socket socket, PipeScheduler scheduler)
+        public SocketSender(Socket socket, PipeScheduler scheduler) : base(socket, scheduler)
         {
-            _socket = socket;
-            _awaitable = new SocketAwaitable(scheduler);
-            _eventArgs.UserToken = _awaitable;
-            _eventArgs.Completed += (_, e) => ((SocketAwaitable)e.UserToken).Complete(e.BytesTransferred, e.SocketError);
         }
 
-        public SocketAwaitable SendAsync(ReadOnlySequence<byte> buffers)
+        public SocketAwaitableEventArgs SendAsync(ReadOnlySequence<byte> buffers)
         {
             if (buffers.IsSingleSegment)
             {
@@ -35,49 +27,49 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
             }
 
 #if NETCOREAPP2_1
-            if (!_eventArgs.MemoryBuffer.Equals(Memory<byte>.Empty))
+            if (!_awaitableEventArgs.MemoryBuffer.Equals(Memory<byte>.Empty))
 #elif NETSTANDARD2_0
-            if (_eventArgs.Buffer != null)
+            if (_awaitableEventArgs.Buffer != null)
 #else
 #error TFMs need to be updated
 #endif
             {
-                _eventArgs.SetBuffer(null, 0, 0);
+                _awaitableEventArgs.SetBuffer(null, 0, 0);
             }
 
-            _eventArgs.BufferList = GetBufferList(buffers);
+            _awaitableEventArgs.BufferList = GetBufferList(buffers);
 
-            if (!_socket.SendAsync(_eventArgs))
+            if (!_socket.SendAsync(_awaitableEventArgs))
             {
-                _awaitable.Complete(_eventArgs.BytesTransferred, _eventArgs.SocketError);
+                _awaitableEventArgs.Complete();
             }
 
-            return _awaitable;
+            return _awaitableEventArgs;
         }
 
-        private SocketAwaitable SendAsync(ReadOnlyMemory<byte> memory)
+        private SocketAwaitableEventArgs SendAsync(ReadOnlyMemory<byte> memory)
         {
             // The BufferList getter is much less expensive then the setter.
-            if (_eventArgs.BufferList != null)
+            if (_awaitableEventArgs.BufferList != null)
             {
-                _eventArgs.BufferList = null;
+                _awaitableEventArgs.BufferList = null;
             }
 
 #if NETCOREAPP2_1
-            _eventArgs.SetBuffer(MemoryMarshal.AsMemory(memory));
+            _awaitableEventArgs.SetBuffer(MemoryMarshal.AsMemory(memory));
 #elif NETSTANDARD2_0
             var segment = memory.GetArray();
 
-            _eventArgs.SetBuffer(segment.Array, segment.Offset, segment.Count);
+            _awaitableEventArgs.SetBuffer(segment.Array, segment.Offset, segment.Count);
 #else
 #error TFMs need to be updated
 #endif
-            if (!_socket.SendAsync(_eventArgs))
+            if (!_socket.SendAsync(_awaitableEventArgs))
             {
-                _awaitable.Complete(_eventArgs.BytesTransferred, _eventArgs.SocketError);
+                _awaitableEventArgs.Complete();
             }
 
-            return _awaitable;
+            return _awaitableEventArgs;
         }
 
         private List<ArraySegment<byte>> GetBufferList(ReadOnlySequence<byte> buffer)
@@ -102,10 +94,5 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 
             return _bufferList;
         }
-
-        public void Dispose()
-        {
-            _eventArgs.Dispose();
-        }
     }
 }

src/Kestrel.Transport.Sockets/Internal/SocketSenderReceiverBase.cs

@@ -0,0 +1,23 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.IO.Pipelines;
+using System.Net.Sockets;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
+{
+    public abstract class SocketSenderReceiverBase : IDisposable
+    {
+        protected readonly Socket _socket;
+        protected readonly SocketAwaitableEventArgs _awaitableEventArgs;
+
+        protected SocketSenderReceiverBase(Socket socket, PipeScheduler scheduler)
+        {
+            _socket = socket;
+            _awaitableEventArgs = new SocketAwaitableEventArgs(scheduler);
+        }
+
+        public void Dispose() => _awaitableEventArgs.Dispose();
+    }
+}