Allow NetworkStream.WriteAsync calls to be canceled asynchronously (d…

…otnet/corefx#40839)

In 3.0 we added the necessary support for this, and for Socket.WriteAsync calls we're appropriately passing the CancellationToken through, but we missed doing so for NetworkStream.WriteAsync, which takes a slightly different code path through AwaitableSocketAsyncEventArgs.

Commit migrated from dotnet/corefx@0a5b928

src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.Tasks.cs

@@ -376,7 +376,7 @@ internal ValueTask SendAsyncForNetworkStream(ReadOnlyMemory<byte> buffer, Socket
                 saea.SetBuffer(MemoryMarshal.AsMemory(buffer));
                 saea.SocketFlags = socketFlags;
                 saea.WrapExceptionsInIOExceptions = true;
-                return saea.SendAsyncForNetworkStream(this);
+                return saea.SendAsyncForNetworkStream(this, cancellationToken);
             }
             else
             {
@@ -929,12 +929,13 @@ public ValueTask<int> SendAsync(Socket socket, CancellationToken cancellationTok
                     new ValueTask<int>(Task.FromException<int>(CreateException(error)));
             }
 
-            public ValueTask SendAsyncForNetworkStream(Socket socket)
+            public ValueTask SendAsyncForNetworkStream(Socket socket, CancellationToken cancellationToken)
             {
                 Debug.Assert(Volatile.Read(ref _continuation) == null, $"Expected null continuation to indicate reserved for use");
 
-                if (socket.SendAsync(this))
+                if (socket.SendAsync(this, cancellationToken))
                 {
+                    _cancellationToken = cancellationToken;
                     return new ValueTask(this, _token);
                 }
 

src/libraries/System.Net.Sockets/tests/FunctionalTests/NetworkStreamTest.netcoreapp.cs

@@ -395,6 +395,37 @@ await RunWithConnectedNetworkStreamsAsync(async (server, client) =>
             });
         }
 
+        [Fact]
+        public async Task WriteAsync_CancelPendingWrite_SucceedsOrThrowsOperationCanceled()
+        {
+            await RunWithConnectedNetworkStreamsAsync(async (server, client) =>
+            {
+                await Assert.ThrowsAnyAsync<OperationCanceledException>(() => client.WriteAsync(new byte[1], 0, 1, new CancellationToken(true)));
+                await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => { await client.WriteAsync(new Memory<byte>(new byte[1]), new CancellationToken(true)); });
+
+                byte[] hugeBuffer = new byte[100_000_000];
+                Exception e;
+
+                var cts = new CancellationTokenSource();
+                Task t = client.WriteAsync(hugeBuffer, 0, hugeBuffer.Length, cts.Token);
+                cts.Cancel();
+                e = await Record.ExceptionAsync(async () => await t);
+                if (e != null)
+                {
+                    Assert.IsAssignableFrom<OperationCanceledException>(e);
+                }
+
+                cts = new CancellationTokenSource();
+                ValueTask vt = client.WriteAsync(new Memory<byte>(hugeBuffer), cts.Token);
+                cts.Cancel();
+                e = await Record.ExceptionAsync(async () => await vt);
+                if (e != null)
+                {
+                    Assert.IsAssignableFrom<OperationCanceledException>(e);
+                }
+            });
+        }
+
         private sealed class CustomSynchronizationContext : SynchronizationContext
         {
             public override void Post(SendOrPostCallback d, object state)