Remove TryGetPointer from RioTcpConnection

src/System.IO.Pipelines.Networking.Windows.RIO/Internal/RioThread.cs

@@ -1,12 +1,12 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO.Pipelines.Networking.Windows.RIO.Internal.Winsock;
 using System.Runtime.InteropServices;
 using System.Threading;
-using System.IO.Pipelines.Networking.Windows.RIO.Internal.Winsock;
 
 namespace System.IO.Pipelines.Networking.Windows.RIO.Internal
 {
@@ -92,7 +92,7 @@ public void RemoveConnection(long key)
             }
         }
 
-        public IntPtr GetBufferId(IntPtr address, out long startAddress)
+        private IntPtr GetBufferId(IntPtr address, out long startAddress)
         {
             var id = IntPtr.Zero;
             startAddress = 0;
@@ -116,6 +116,25 @@ public IntPtr GetBufferId(IntPtr address, out long startAddress)
             return id;
         }
 
+        public unsafe RioBufferSegment GetSegmentFromMemory(Buffer<byte> memory)
+        {
+            // It's ok to unpin the handle here because the memory is from the pool
+            // we created, which is already pinned.
+            var pin = memory.Pin();
+            var spanPtr = (IntPtr)pin.PinnedPointer;
+            pin.Free();
+
+            long startAddress;
+            long spanAddress = spanPtr.ToInt64();
+            var bufferId = GetBufferId(spanPtr, out startAddress);
+
+            checked
+            {
+                var offset = (uint)(spanAddress - startAddress);
+                return new RioBufferSegment(bufferId, offset, (uint)memory.Length);
+            }
+        }
+
         private void OnSlabAllocated(MemoryPoolSlab slab)
         {
             lock (_bufferIdMappings)

src/System.IO.Pipelines.Networking.Windows.RIO/RioTcpConnection.cs

@@ -1,12 +1,11 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
 using System.Buffers;
-using System.Threading;
-using System.Threading.Tasks;
 using System.IO.Pipelines.Networking.Windows.RIO.Internal;
 using System.IO.Pipelines.Networking.Windows.RIO.Internal.Winsock;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace System.IO.Pipelines.Networking.Windows.RIO
 {
@@ -64,7 +63,7 @@ internal RioTcpConnection(IntPtr socket, long connectionId, IntPtr requestQueue,
         private void ProcessReceives()
         {
             _buffer = _input.Writer.Alloc(2048);
-            var receiveBufferSeg = GetSegmentFromMemory(_buffer.Buffer);
+            var receiveBufferSeg = _rioThread.GetSegmentFromMemory(_buffer.Buffer);
 
             if (!_rio.RioReceive(_requestQueue, ref receiveBufferSeg, 1, RioReceiveFlags.None, 0))
             {
@@ -120,7 +119,7 @@ private Task SendAsync(Buffer<byte> memory, bool endOfMessage)
 
             var flushSends = endOfMessage || MaxOutstandingSendsReached;
 
-            Send(GetSegmentFromMemory(memory), flushSends);
+            Send(_rioThread.GetSegmentFromMemory(memory), flushSends);
 
             if (flushSends && !endOfMessage)
             {
@@ -136,7 +135,7 @@ private async Task SendAsyncAwaited(Buffer<byte> memory, bool endOfMessage)
 
             var flushSends = endOfMessage || MaxOutstandingSendsReached;
 
-            Send(GetSegmentFromMemory(memory), flushSends);
+            Send(_rioThread.GetSegmentFromMemory(memory), flushSends);
 
             if (flushSends && !endOfMessage)
             {
@@ -219,25 +218,6 @@ public void ReceiveEndComplete()
             _buffer.FlushAsync();
         }
 
-        private unsafe RioBufferSegment GetSegmentFromMemory(Buffer<byte> memory)
-        {
-            void* pointer;
-            if (!memory.TryGetPointer(out pointer))
-            {
-                throw new InvalidOperationException("Memory needs to be pinned");
-            }
-            var spanPtr = (IntPtr)pointer;
-            long startAddress;
-            long spanAddress = spanPtr.ToInt64();
-            var bufferId = _rioThread.GetBufferId(spanPtr, out startAddress);
-
-            checked
-            {
-                var offset = (uint)(spanAddress - startAddress);
-                return new RioBufferSegment(bufferId, offset, (uint)memory.Length);
-            }
-        }
-
         private static void ThrowError(ErrorType type)
         {
             var errorNo = RioImports.WSAGetLastError();