Apply PR feedback from dotnet#74611

src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/MsQuicBuffers.cs

@@ -1,11 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Diagnostics;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
 using Microsoft.Quic;
-using System.Threading;
 
 namespace System.Net.Quic;
 
@@ -20,27 +18,14 @@ internal unsafe struct MsQuicBuffers : IDisposable
     private QUIC_BUFFER* _buffers;
     // Number of QUIC_BUFFER instance currently allocated in _buffers, so that we can reuse the memory instead of reallocating.
     private int _count;
-    private bool _initialized;
-    private bool _disposed;
 
     public MsQuicBuffers()
     {
         _buffers = null;
         _count = 0;
-        _initialized=false;
-        _disposed = false;
-    }
-
-    public QUIC_BUFFER* Buffers
-    {
-        get
-        {
-            ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
-            Debug.Assert(_initialized);
-            return _buffers;
-        }
     }
 
+    public QUIC_BUFFER* Buffers => _buffers;
     public int Count => _count;
 
     private void FreeNativeMemory()
@@ -63,7 +48,6 @@ private void Reserve(int count)
 
     private void SetBuffer(int index, ReadOnlyMemory<byte> buffer)
     {
-        Debug.Assert(_buffers[index].Buffer == null);
         _buffers[index].Buffer = (byte*)NativeMemory.Alloc((nuint)buffer.Length, (nuint)sizeof(byte));
         _buffers[index].Length = (uint)buffer.Length;
         buffer.Span.CopyTo(_buffers[index].Span);
@@ -78,14 +62,11 @@ private void SetBuffer(int index, ReadOnlyMemory<byte> buffer)
     /// <typeparam name="T">The type of the inputs.</typeparam>
     public void Initialize<T>(IList<T> inputs, Func<T, ReadOnlyMemory<byte>> toBuffer)
     {
-        ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
-        Debug.Assert(!_initialized);
         Reserve(inputs.Count);
         for (int i = 0; i < inputs.Count; ++i)
         {
             SetBuffer(i, toBuffer(inputs[i]));
         }
-        _initialized = true;
     }
 
     /// <summary>
@@ -95,21 +76,15 @@ public void Initialize<T>(IList<T> inputs, Func<T, ReadOnlyMemory<byte>> toBuffe
     /// <param name="buffer">Buffer to be passed to MsQuic as QUIC_BUFFER*.</param>
     public void Initialize(ReadOnlyMemory<byte> buffer)
     {
-        ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
-        Debug.Assert(!_initialized);
         Reserve(1);
         SetBuffer(0, buffer);
-        _initialized = true;
     }
 
     /// <summary>
     /// Release the native memory of individual buffers and allows reuse of this struct.
     /// </summary>
-    public void Reset() => Reset(disposing: false);
-
-    private void Reset(bool disposing)
+    public void Reset()
     {
-        ObjectDisposedException.ThrowIf(_disposed && !disposing, typeof(MsQuicBuffers));
         for (int i = 0; i < _count; ++i)
         {
             if (_buffers[i].Buffer is null)
@@ -121,16 +96,14 @@ private void Reset(bool disposing)
             NativeMemory.Free(buffer);
             _buffers[i].Length = 0;
         }
-        _initialized = false;
     }
 
     /// <summary>
     /// Releases all the native memory.
     /// </summary>
     public void Dispose()
     {
-        _disposed = true;
-        Reset(disposing: true);
+        Reset();
         FreeNativeMemory();
     }
 }

src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs

@@ -363,10 +363,21 @@ public ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, bool completeWrites, Ca
 
         lock (_sendBuffersLock)
         {
-            ObjectDisposedException.ThrowIf(_disposed == 1, this);
-            _sendBuffers.Initialize(buffer);
+            ObjectDisposedException.ThrowIf(_disposed == 1, this); // TODO: valueTask is left unobserved
             unsafe
             {
+                if (_sendBuffers.Count > 0 && _sendBuffers.Buffers[0].Buffer != null)
+                {
+                    // _sendBuffers are not reset, meaning SendComplete for the previous WriteAsync call didn't arrive yet.
+                    // In case of cancellation, the task from _sendTcs is finished before the aborting. It is technically possible for subsequent
+                    // WriteAsync to grab the next task from _sendTcs and start executing before SendComplete event occurs for the previous (canceled) write.
+                    // This is not an "invalid nested call", because the previous task has finished. Best guess is to mimic OperationAborted as it will be from Abort
+                    // that would execute soon enough, if not already. Not final, because Abort should be the one to set final exception.
+                    _sendTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_writing_aborted), final: false);
+                    return valueTask;
+                }
+
+                _sendBuffers.Initialize(buffer);
                 int status = MsQuicApi.Api.StreamSend(
                     _handle,
                     _sendBuffers.Buffers,
@@ -515,6 +526,8 @@ private unsafe int HandleEventSendComplete(ref SEND_COMPLETE data)
             NetEventSource.Info(this, $"{this} Received event SEND_COMPLETE with {nameof(data.Canceled)}={data.Canceled}");
         }
 
+		// In case of cancellation, the task from _sendTcs is finished before the aborting. It is technically possible for subsequent WriteAsync to grab the next task
+        // from _sendTcs and start executing before SendComplete event occurs for the previous (canceled) write
         lock (_sendBuffersLock)
         {
             _sendBuffers.Reset();