implement cancellation support for SendFileAsync and DisconnectAsync (#…

…53062) implement cancellation support for SendFileAsync and DisconnectAsync, and rework some internal async logic to support this and reduce code duplication
dotnet · May 27, 2021 · bb38de7 · bb38de7
1 parent 18f5931
commit bb38de7
Show file tree

Hide file tree

Showing 7 changed files with 199 additions and 158 deletions.
diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.Tasks.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.Tasks.cs
@@ -1098,8 +1098,7 @@ public ValueTask SendPacketsAsync(Socket socket, CancellationToken cancellationT
             {
                 Debug.Assert(Volatile.Read(ref _continuation) == null, "Expected null continuation to indicate reserved for use");
 
-                // TODO: Support cancellation by passing cancellationToken down through SendPacketsAsync, etc.
-                if (socket.SendPacketsAsync(this))
+                if (socket.SendPacketsAsync(this, cancellationToken))
                 {
                     _cancellationToken = cancellationToken;
                     return new ValueTask(this, _token);
@@ -1379,7 +1378,10 @@ SocketReceiveMessageFromResult IValueTaskSource<SocketReceiveMessageFromResult>.
 
             private void ThrowException(SocketError error, CancellationToken cancellationToken)
             {
-                if (error == SocketError.OperationAborted)
+                // Most operations will report OperationAborted when canceled.
+                // On Windows, SendFileAsync will report ConnectionAborted.
+                // There's a race here anyway, so there's no harm in also checking for ConnectionAborted in all cases.
+                if (error == SocketError.OperationAborted || error == SocketError.ConnectionAborted)
                 {
                     cancellationToken.ThrowIfCancellationRequested();
                 }

diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.cs
@@ -3060,7 +3060,9 @@ private bool SendAsync(SocketAsyncEventArgs e, CancellationToken cancellationTok
             return socketError == SocketError.IOPending;
         }
 
-        public bool SendPacketsAsync(SocketAsyncEventArgs e)
+        public bool SendPacketsAsync(SocketAsyncEventArgs e) => SendPacketsAsync(e, default(CancellationToken));
+
+        private bool SendPacketsAsync(SocketAsyncEventArgs e, CancellationToken cancellationToken)
         {
             ThrowIfDisposed();
 
@@ -3082,7 +3084,7 @@ public bool SendPacketsAsync(SocketAsyncEventArgs e)
             SocketError socketError;
             try
             {
-                socketError = e.DoOperationSendPackets(this, _handle);
+                socketError = e.DoOperationSendPackets(this, _handle, cancellationToken);
             }
             catch (Exception)
             {

diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncContext.Unix.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncContext.Unix.cs
@@ -2100,7 +2100,7 @@ public SocketError SendFile(SafeFileHandle fileHandle, long offset, long count,
             return operation.ErrorCode;
         }
 
-        public SocketError SendFileAsync(SafeFileHandle fileHandle, long offset, long count, out long bytesSent, Action<long, SocketError> callback)
+        public SocketError SendFileAsync(SafeFileHandle fileHandle, long offset, long count, out long bytesSent, Action<long, SocketError> callback, CancellationToken cancellationToken = default)
         {
             SetHandleNonBlocking();
 
@@ -2122,7 +2122,7 @@ public SocketError SendFileAsync(SafeFileHandle fileHandle, long offset, long co
                 BytesTransferred = bytesSent
             };
 
-            if (!_sendQueue.StartAsyncOperation(this, operation, observedSequenceNumber))
+            if (!_sendQueue.StartAsyncOperation(this, operation, observedSequenceNumber, cancellationToken))
             {
                 bytesSent = operation.BytesTransferred;
                 return operation.ErrorCode;

diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.Unix.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.Unix.cs
@@ -244,7 +244,7 @@ internal unsafe SocketError DoOperationSend(SafeSocketHandle handle, Cancellatio
             return errorCode;
         }
 
-        internal SocketError DoOperationSendPackets(Socket socket, SafeSocketHandle handle)
+        internal SocketError DoOperationSendPackets(Socket socket, SafeSocketHandle handle, CancellationToken cancellationToken)
         {
             Debug.Assert(_sendPacketsElements != null);
             SendPacketsElement[] elements = (SendPacketsElement[])_sendPacketsElements.Clone();
@@ -288,7 +288,7 @@ internal SocketError DoOperationSendPackets(Socket socket, SafeSocketHandle hand
                 throw;
             }
 
-            SocketPal.SendPacketsAsync(socket, SendPacketsFlags, elements, files, (bytesTransferred, error) =>
+            SocketPal.SendPacketsAsync(socket, SendPacketsFlags, elements, files, cancellationToken, (bytesTransferred, error) =>
             {
                 if (error == SocketError.Success)
                 {