Fix ManagedWebSocket cancellation race condition (#27407)

When a cancelable cancellation token is passed to Receive/SendAsync, the code registers with the token to abort the web socket.  The resulting exception is then translated into an OperationCanceledException if _abort is true.  But this logic doesn't count on the possibility that the underlying Read/WriteAsync on the transport stream could actually throw a cancellation exception.  If it does, there's a race condition as to whether the web socket's abort flag will be set in time; if it is, we wrap the cancellation exception in another cancellation exception, and if it's not, we treat it as a non-cancellation exception.  This commit fixes that, by explicitly allowing OperationCanceledExceptions to proceed uninterrupted.  This should fix the ReceiveAsync_AfterCancellationDoReceiveAsync_ThrowsWebSocketException test, which recently started failing sporadically, I believe because we improved the cancellation support in SocketsHttpHandler, on top of which ClientWebSocket now sits.

Author: stephentoub

src/Common/src/System/Net/WebSockets/ManagedWebSocket.cs

@@ -415,8 +415,9 @@ private Task SendFrameLockAcquiredNonCancelableAsync(MessageOpcode opcode, bool
             }
             catch (Exception exc)
             {
-                return Task.FromException(_state == WebSocketState.Aborted ?
-                    CreateOperationCanceledException(exc) :
+                return Task.FromException(
+                    exc is OperationCanceledException ? exc :
+                    _state == WebSocketState.Aborted ? CreateOperationCanceledException(exc) :
                     new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc));
             }
             finally
@@ -437,7 +438,7 @@ private Task SendFrameLockAcquiredNonCancelableAsync(MessageOpcode opcode, bool
                 thisRef.ReleaseSendBuffer();
 
                 try { t.GetAwaiter().GetResult(); }
-                catch (Exception exc)
+                catch (Exception exc) when (!(exc is OperationCanceledException))
                 {
                     throw thisRef._state == WebSocketState.Aborted ?
                         CreateOperationCanceledException(exc) :
@@ -457,7 +458,7 @@ private async Task SendFrameFallbackAsync(MessageOpcode opcode, bool endOfMessag
                     await _stream.WriteAsync(_sendBuffer, 0, sendBytes, cancellationToken).ConfigureAwait(false);
                 }
             }
-            catch (Exception exc)
+            catch (Exception exc) when (!(exc is OperationCanceledException))
             {
                 throw _state == WebSocketState.Aborted ?
                     CreateOperationCanceledException(exc, cancellationToken) :
@@ -739,7 +740,7 @@ private async ValueTask<TWebSocketReceiveResult> ReceiveAsyncPrivate<TWebSocketR
                         null, null);
                 }
             }
-            catch (Exception exc)
+            catch (Exception exc) when (!(exc is OperationCanceledException))
             {
                 if (_state == WebSocketState.Aborted)
                 {