pool async state in SslStream (#69418)

* pool async state in SslStream * fix tests * remove GetAwaiter from test
dotnet · May 26, 2022 · 008f128 · 008f128
1 parent 4e9ac49
commit 008f128
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/src/libraries/System.Net.Security/src/System/Net/Security/SslStream.IO.cs b/src/libraries/System.Net.Security/src/System/Net/Security/SslStream.IO.cs
@@ -2,9 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers;
-using System.ComponentModel;
 using System.Diagnostics;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Runtime.ExceptionServices;
 using System.Security.Authentication;
 using System.Security.Cryptography.X509Certificates;
@@ -675,7 +675,7 @@ private bool HaveFullTlsFrame(out int frameSize)
             return _buffer.EncryptedLength >= frameSize;
         }
 
-
+        [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
         private async ValueTask<int> EnsureFullTlsFrameAsync<TIOAdapter>(CancellationToken cancellationToken)
             where TIOAdapter : IReadWriteAdapter
         {
@@ -760,6 +760,7 @@ private SecurityStatusPal DecryptData(int frameSize)
             return status;
         }
 
+        [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
         private async ValueTask<int> ReadAsyncInternal<TIOAdapter>(Memory<byte> buffer, CancellationToken cancellationToken)
             where TIOAdapter : IReadWriteAdapter
         {

diff --git a/src/libraries/System.Net.Security/tests/FunctionalTests/SslStreamNetworkStreamTest.cs b/src/libraries/System.Net.Security/tests/FunctionalTests/SslStreamNetworkStreamTest.cs
@@ -961,19 +961,20 @@ public async Task SslStream_RandomSizeWrites_OK(int bufferSize, int readBufferSi
 
                 await TestConfiguration.WhenAllOrAnyFailedWithTimeout(t1, t2);
 
-                Task writer = Task.Run(() =>
+                Task writer = Task.Run(async () =>
                 {
                     Memory<byte> data = new Memory<byte>(dataToCopy);
                     while (data.Length > 0)
                     {
                         int writeLength = Math.Min(data.Length, writeBufferSize);
                         if (useAsync)
                         {
-                            server.WriteAsync(data.Slice(0, writeLength)).GetAwaiter().GetResult();
+                            await server.WriteAsync(data.Slice(0, writeLength));
                         }
                         else
                         {
                             server.Write(data.Span.Slice(0, writeLength));
+                            await Task.CompletedTask;
                         }
 
                         data = data.Slice(Math.Min(writeBufferSize, data.Length));
@@ -982,7 +983,7 @@ public async Task SslStream_RandomSizeWrites_OK(int bufferSize, int readBufferSi
                     server.ShutdownAsync().GetAwaiter().GetResult();
                 });
 
-                Task reader = Task.Run(() =>
+                Task reader = Task.Run(async () =>
                 {
                     Memory<byte> readBuffer = new Memory<byte>(dataReceived);
                     int totalLength = 0;
@@ -992,11 +993,12 @@ public async Task SslStream_RandomSizeWrites_OK(int bufferSize, int readBufferSi
                     {
                         if (useAsync)
                         {
-                            readLength = client.ReadAsync(readBuffer.Slice(totalLength, readBufferSize)).GetAwaiter().GetResult();
+                            readLength = await client.ReadAsync(readBuffer.Slice(totalLength, readBufferSize));
                         }
                         else
                         {
                             readLength = client.Read(readBuffer.Span.Slice(totalLength, readBufferSize));
+                            await Task.CompletedTask;
                         }
 
                         if (readLength == 0)