Use shared MsQuicBuffers struct

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

@@ -39,35 +39,94 @@ private void FreeNativeMemory()
         NativeMemory.Free(buffers);
     }
 
+    private void Reserve(int count)
+    {
+        if (_handles.Length < count)
+        {
+            _handles = new MemoryHandle[count];
+            FreeNativeMemory();
+            _buffers = (QUIC_BUFFER*)NativeMemory.Alloc((nuint)count, (nuint)sizeof(QUIC_BUFFER));
+        }
+    }
+
+    private void SetBuffer(int index, ReadOnlyMemory<byte> buffer)
+    {
+        MemoryHandle handle = buffer.Pin();
+
+        _handles[index] = handle;
+        _buffers[index].Buffer = (byte*)handle.Pointer;
+        _buffers[index].Length = (uint)buffer.Length;
+    }
+
     /// <summary>
-    /// The method initializes QUIC_BUFFER* with data from inputs, converted via toBuffer.
+    /// Initializes QUIC_BUFFER* with data from inputs, converted via toBuffer.
     /// Note that the struct either needs to be freshly created via new or previously cleaned up with Reset.
     /// </summary>
     /// <param name="inputs">Inputs to get their byte array, pin them and pepare them to be passed to MsQuic as QUIC_BUFFER*.</param>
     /// <param name="toBuffer">Method extracting byte array from the inputs, e.g. applicationProtocol.Protocol.</param>
     /// <typeparam name="T">The type of the inputs.</typeparam>
     public void Initialize<T>(IList<T> inputs, Func<T, ReadOnlyMemory<byte>> toBuffer)
     {
-        if (_handles.Length < inputs.Count)
+        Reserve(inputs.Count);
+        _count = inputs.Count;
+
+        for (int i = 0; i < inputs.Count; ++i)
         {
-            _handles = new MemoryHandle[inputs.Count];
+            ReadOnlyMemory<byte> buffer = toBuffer(inputs[i]);
+            SetBuffer(i, buffer);
         }
-        if (_count < inputs.Count)
+    }
+
+    /// <summary>
+    /// Initializes QUIC_BUFFER* with the provided buffer.
+    /// Note that the struct either needs to be freshly created via new or previously cleaned up with Reset.
+    /// </summary>
+    /// <param name="buffer">Buffer to be passed to MsQuic as QUIC_BUFFER*.</param>
+    public void Initialize(ReadOnlyMemory<byte> buffer)
+    {
+        Reserve(1);
+        _count = 1;
+        SetBuffer(0, buffer);
+    }
+
+    /// <summary>
+    /// Initializes QUIC_BUFFER* with the provided buffers.
+    /// Note that the struct either needs to be freshly created via new or previously cleaned up with Reset.
+    /// </summary>
+    /// <param name="buffers">Buffers to be passed to MsQuic as QUIC_BUFFER*.</param>
+    public void Initialize(ReadOnlySequence<byte> buffers)
+    {
+        int count = 0;
+        foreach (ReadOnlyMemory<byte> _ in buffers)
         {
-            FreeNativeMemory();
-            _buffers = (QUIC_BUFFER*)NativeMemory.Alloc((nuint)sizeof(QUIC_BUFFER), (nuint)inputs.Count);
+            ++count;
         }
 
-        _count = inputs.Count;
+        Reserve(count);
+        _count = count;
+        count = 0;
 
-        for (int i = 0; i < inputs.Count; ++i)
+        foreach (ReadOnlyMemory<byte> buffer in buffers)
         {
-            ReadOnlyMemory<byte> buffer = toBuffer(inputs[i]);
-            MemoryHandle handle = buffer.Pin();
+            SetBuffer(count, buffer);
+            ++count;
+        }
+    }
 
-            _handles[i] = handle;
-            _buffers[i].Buffer = (byte*)handle.Pointer;
-            _buffers[i].Length = (uint)buffer.Length;
+    /// <summary>
+    /// Initializes QUIC_BUFFER* with the provided buffers.
+    /// Note that the struct either needs to be freshly created via new or previously cleaned up with Reset.
+    /// </summary>
+    /// <param name="buffers">Buffers to be passed to MsQuic as QUIC_BUFFER*.</param>
+    public void Initialize(ReadOnlyMemory<ReadOnlyMemory<byte>> buffers)
+    {
+        int count = buffers.Length;
+        Reserve(count);
+        _count = count;
+
+        for (int i = 0; i < buffers.Length; i++)
+        {
+            SetBuffer(i, buffers.Span[i]);
         }
     }
 

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicStream.cs

@@ -58,11 +58,7 @@ private sealed class State
             public SendState SendState;
             public long SendErrorCode = -1;
 
-            // Buffers to hold during a call to send.
-            public MemoryHandle[] BufferArrays = new MemoryHandle[1];
-            public IntPtr SendQuicBuffers;
-            public int SendBufferMaxCount;
-            public int SendBufferCount;
+            public MsQuicBuffers SendBuffers;
 
             // Resettable completions to be used for multiple calls to send.
             public readonly ResettableCompletionSource<int> SendResettableCompletionSource = new ResettableCompletionSource<int>();
@@ -85,15 +81,19 @@ private sealed class State
             // Set once stream have been shutdown.
             public readonly TaskCompletionSource ShutdownCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
 
+            public State()
+            {
+                SendBuffers = new MsQuicBuffers();
+            }
+
             public void Cleanup()
             {
                 if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"{Handle} releasing handles.");
 
                 ShutdownState = ShutdownState.Finished;
                 CleanupSendState(this);
                 Handle?.Dispose();
-                Marshal.FreeHGlobal(SendQuicBuffers);
-                SendQuicBuffers = IntPtr.Zero;
+                SendBuffers.Dispose();
                 if (StateGCHandle.IsAllocated) StateGCHandle.Free();
                 ConnectionState?.RemoveStream(null);
             }
@@ -323,10 +323,10 @@ private unsafe ValueTask WriteAsyncInternal<TAdapter>(TAdapter adapter, QUIC_SEN
             adapter.SetupSendState(_state);
 
             Debug.Assert(!Monitor.IsEntered(_state), "!Monitor.IsEntered(_state)");
-            uint status = MsQuicApi.Api.StreamSendDelegate(
-                _state.Handle,
-                (QuicBuffer*)_state.SendQuicBuffers,
-                (uint)_state.SendBufferCount,
+            int status = MsQuicApi.Api.ApiTable->StreamSend(
+                _state.Handle.QuicHandle,
+                _state.SendBuffers.Buffers,
+                (uint)_state.SendBuffers.Count,
                 flags,
                 IntPtr.Zero);
 
@@ -1416,12 +1416,7 @@ private static void CleanupSendState(State state)
             lock (state)
             {
                 Debug.Assert(state.SendState != SendState.Pending);
-                Debug.Assert(state.SendBufferCount <= state.BufferArrays.Length);
-
-                for (int i = 0; i < state.SendBufferCount; i++)
-                {
-                    state.BufferArrays[i].Dispose();
-                }
+                state.SendBuffers.Reset();
             }
         }
 
@@ -1713,27 +1708,6 @@ private interface ISendBufferAdapter
             void SetupSendState(State state);
 
             bool IsEmpty { get; }
-
-            static unsafe void Reserve(State state, int count)
-            {
-                if (state.SendBufferMaxCount < count)
-                {
-                    NativeMemory.Free((void*)state.SendQuicBuffers);
-                    state.SendQuicBuffers = IntPtr.Zero;
-                    state.SendQuicBuffers = (IntPtr)NativeMemory.Alloc((nuint)count, (nuint)sizeof(QuicBuffer));
-                    state.SendBufferMaxCount = count;
-                    state.BufferArrays = new MemoryHandle[count];
-                }
-            }
-
-            static unsafe void SetBuffer(State state, int index, ReadOnlyMemory<byte> buffer)
-            {
-                MemoryHandle handle = buffer.Pin();
-                QuicBuffer* quicBuffer = (QuicBuffer*)state.SendQuicBuffers + index;
-                quicBuffer->Length = (uint)buffer.Length;
-                quicBuffer->Buffer = (byte*)handle.Pointer;
-                state.BufferArrays[index] = handle;
-            }
         }
 
         private struct WriteMemoryAdapter : ISendBufferAdapter
@@ -1744,11 +1718,9 @@ private struct WriteMemoryAdapter : ISendBufferAdapter
 
             public bool IsEmpty => _buffer.IsEmpty;
 
-            public unsafe void SetupSendState(State state)
+            public void SetupSendState(State state)
             {
-                ISendBufferAdapter.Reserve(state, 1);
-                ISendBufferAdapter.SetBuffer(state, 0, _buffer);
-                state.SendBufferCount = 1;
+                state.SendBuffers.Initialize(_buffer);
             }
         }
 
@@ -1760,24 +1732,9 @@ private struct WriteSequenceAdapter : ISendBufferAdapter
 
             public bool IsEmpty => _buffers.IsEmpty;
 
-            public unsafe void SetupSendState(State state)
+            public void SetupSendState(State state)
             {
-                int count = 0;
-
-                foreach (ReadOnlyMemory<byte> _ in _buffers)
-                {
-                    ++count;
-                }
-
-                ISendBufferAdapter.Reserve(state, count);
-                state.SendBufferCount = count;
-                count = 0;
-
-                foreach (ReadOnlyMemory<byte> buffer in _buffers)
-                {
-                    ISendBufferAdapter.SetBuffer(state, count, buffer);
-                    ++count;
-                }
+                state.SendBuffers.Initialize(_buffers);
             }
         }
 
@@ -1789,18 +1746,9 @@ private struct WriteMemoryOfMemoryAdapter : ISendBufferAdapter
 
             public bool IsEmpty => _buffers.IsEmpty;
 
-            public unsafe void SetupSendState(State state)
+            public void SetupSendState(State state)
             {
-                int count = _buffers.Length;
-                ISendBufferAdapter.Reserve(state, count);
-                state.SendBufferCount = count;
-                count = 0;
-
-                foreach (ReadOnlyMemory<byte> buffer in _buffers.Span)
-                {
-                    ISendBufferAdapter.SetBuffer(state, count, buffer);
-                    ++count;
-                }
+                state.SendBuffers.Initialize(_buffers);
             }
         }
     }