Skip to content
This repository was archived by the owner on Jan 23, 2023. It is now read-only.

Commit d67ea51

Browse files
benaadamsahsonkhan
benaadams
authored and
ahsonkhan
committed
ArrayMemoryPoolBuffer should only implement IMemoryOwner<T> (#28690)
* ArrayMemoryPoolBuffer should only implement IMemoryOwner<T> * React pipelines to changes
1 parent b01e361 commit d67ea51

File tree

4 files changed

+70
-143
lines changed

4 files changed

+70
-143
lines changed

‎src/System.IO.Pipelines/tests/System.IO.Pipelines.Tests.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
<Import Project="$([MSBuild]::GetDirectoryNameOfFileAbove($(MSBuildThisFileDirectory), dir.props))\dir.props" />
44
<PropertyGroup>
55
<ProjectGuid>{9E984EB2-827E-4029-9647-FB5F8B67C553}</ProjectGuid>
6+
<AllowUnsafeBlocks>true</AllowUnsafeBlocks>
67
</PropertyGroup>
78
<PropertyGroup Condition="'$(Configuration)|$(Platform)' == 'netcoreapp-Debug|AnyCPU'" />
89
<PropertyGroup Condition="'$(Configuration)|$(Platform)' == 'netcoreapp-Release|AnyCPU'" />

‎src/System.IO.Pipelines/tests/TestMemoryPool.cs‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
using System.Buffers;
66
using System.Diagnostics;
7+
using System.Runtime.CompilerServices;
78
using System.Runtime.InteropServices;
89
using System.Threading;
910

@@ -18,7 +19,7 @@ public class TestMemoryPool: MemoryPool<byte>
1819
public override IMemoryOwner<byte> Rent(int minBufferSize = -1)
1920
{
2021
CheckDisposed();
21-
return new PooledMemory((MemoryManager<byte>)_pool.Rent(minBufferSize), this);
22+
return new PooledMemory(_pool.Rent(minBufferSize), this);
2223
}
2324

2425
protected override void Dispose(bool disposing)
@@ -38,7 +39,7 @@ internal void CheckDisposed()
3839

3940
private class PooledMemory : MemoryManager<byte>
4041
{
41-
private MemoryManager<byte> _manager;
42+
private IMemoryOwner<byte> _owner;
4243

4344
private readonly TestMemoryPool _pool;
4445

@@ -48,9 +49,9 @@ private class PooledMemory : MemoryManager<byte>
4849

4950
private string _leaser;
5051

51-
public PooledMemory(MemoryManager<byte> manager, TestMemoryPool pool)
52+
public PooledMemory(IMemoryOwner<byte> owner, TestMemoryPool pool)
5253
{
53-
_manager = manager;
54+
_owner = owner;
5455
_pool = pool;
5556
_leaser = Environment.StackTrace;
5657
_referenceCount = 1;
@@ -70,13 +71,36 @@ public override MemoryHandle Pin(int elementIndex = 0)
7071
{
7172
_pool.CheckDisposed();
7273
Interlocked.Increment(ref _referenceCount);
73-
return _manager.Pin(elementIndex);
74+
75+
if (!MemoryMarshal.TryGetArray(_owner.Memory, out ArraySegment<byte> segment))
76+
{
77+
throw new InvalidOperationException();
78+
}
79+
80+
unsafe
81+
{
82+
try
83+
{
84+
if ((uint)elementIndex > (uint)segment.Count)
85+
{
86+
throw new ArgumentOutOfRangeException(nameof(elementIndex));
87+
}
88+
89+
GCHandle handle = GCHandle.Alloc(segment.Array, GCHandleType.Pinned);
90+
91+
return new MemoryHandle(Unsafe.Add<byte>(((void*)handle.AddrOfPinnedObject()), elementIndex + segment.Offset), handle, this);
92+
}
93+
catch
94+
{
95+
Unpin();
96+
throw;
97+
}
98+
}
7499
}
75100

76101
public override void Unpin()
77102
{
78103
_pool.CheckDisposed();
79-
_manager.Unpin();
80104

81105
int newRefCount = Interlocked.Decrement(ref _referenceCount);
82106

@@ -92,22 +116,22 @@ public override void Unpin()
92116
protected override bool TryGetArray(out ArraySegment<byte> segment)
93117
{
94118
_pool.CheckDisposed();
95-
return MemoryMarshal.TryGetArray(_manager.Memory, out segment);
119+
return MemoryMarshal.TryGetArray(_owner.Memory, out segment);
96120
}
97121

98122
public override int Length
99123
{
100124
get
101125
{
102126
_pool.CheckDisposed();
103-
return _manager.Length;
127+
return _owner.Memory.Length;
104128
}
105129
}
106130

107131
public override Span<byte> GetSpan()
108132
{
109133
_pool.CheckDisposed();
110-
return _manager.GetSpan();
134+
return _owner.Memory.Span;
111135
}
112136
}
113137
}

‎src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs‎

Lines changed: 12 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
// See the LICENSE file in the project root for more information.
44

55
using System.Runtime.InteropServices;
6-
using System.Threading;
76

87
#if !netstandard
98
using Internal.Runtime.CompilerServices;
@@ -15,102 +14,36 @@ namespace System.Buffers
1514
{
1615
internal sealed partial class ArrayMemoryPool<T> : MemoryPool<T>
1716
{
18-
private sealed class ArrayMemoryPoolBuffer : MemoryManager<T>
17+
private sealed class ArrayMemoryPoolBuffer : IMemoryOwner<T>
1918
{
2019
private T[] _array;
21-
private int _refCount;
2220

2321
public ArrayMemoryPoolBuffer(int size)
2422
{
2523
_array = ArrayPool<T>.Shared.Rent(size);
26-
_refCount = 1;
2724
}
2825

29-
public sealed override int Length => _array.Length;
30-
31-
public bool IsDisposed => _array == null;
32-
33-
public bool IsRetained => Volatile.Read(ref _refCount) > 0;
34-
35-
public sealed override Span<T> GetSpan()
36-
{
37-
if (IsDisposed)
38-
ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
39-
40-
return _array;
41-
}
42-
43-
protected sealed override void Dispose(bool disposing)
26+
public Memory<T> Memory
4427
{
45-
if (_array != null)
28+
get
4629
{
47-
ArrayPool<T>.Shared.Return(_array);
48-
_array = null;
49-
}
50-
}
51-
52-
// TryGetArray is exposed as "protected internal". Normally, the rules of C# dictate we override it as "protected" because the base class is
53-
// in a different assembly. Except in the netstandard config where the base class is in the same assembly.
54-
protected
55-
#if netstandard
56-
internal
57-
#endif
58-
sealed override bool TryGetArray(out ArraySegment<T> segment)
59-
{
60-
if (IsDisposed)
61-
ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
62-
63-
segment = new ArraySegment<T>(_array);
64-
return true;
65-
}
66-
67-
public sealed override MemoryHandle Pin(int elementIndex = 0)
68-
{
69-
unsafe
70-
{
71-
while (true)
30+
T[] array = _array;
31+
if (array == null)
7232
{
73-
int currentCount = Volatile.Read(ref _refCount);
74-
if (currentCount <= 0)
75-
ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
76-
if (Interlocked.CompareExchange(ref _refCount, currentCount + 1, currentCount) == currentCount)
77-
break;
33+
ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
7834
}
7935

80-
try
81-
{
82-
if ((uint)elementIndex > (uint)_array.Length)
83-
{
84-
ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.elementIndex);
85-
}
86-
87-
GCHandle handle = GCHandle.Alloc(_array, GCHandleType.Pinned);
88-
89-
return new MemoryHandle(Unsafe.Add<T>(((void*)handle.AddrOfPinnedObject()), elementIndex), handle, this);
90-
}
91-
catch
92-
{
93-
Unpin();
94-
throw;
95-
}
36+
return new Memory<T>(array);
9637
}
9738
}
9839

99-
public sealed override void Unpin()
40+
public void Dispose()
10041
{
101-
while (true)
42+
T[] array = _array;
43+
if (array != null)
10244
{
103-
int currentCount = Volatile.Read(ref _refCount);
104-
if (currentCount <= 0)
105-
ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
106-
if (Interlocked.CompareExchange(ref _refCount, currentCount - 1, currentCount) == currentCount)
107-
{
108-
if (currentCount == 1)
109-
{
110-
Dispose(true);
111-
}
112-
break;
113-
}
45+
_array = null;
46+
ArrayPool<T>.Shared.Return(array);
11447
}
11548
}
11649
}

‎src/System.Memory/tests/MemoryPool/MemoryPool.cs‎

Lines changed: 24 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ public static void MemoryPoolSpan()
4848
Memory<int> memory = block.Memory;
4949
Span<int> sp = memory.Span;
5050
Assert.Equal(memory.Length, sp.Length);
51-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(memory, out MemoryManager<int> manager));
52-
using (MemoryHandle newMemoryHandle = manager.Pin())
51+
Assert.True(MemoryMarshal.TryGetArray(memory, out ArraySegment<int> arraySegment));
52+
using (MemoryHandle newMemoryHandle = memory.Pin())
5353
{
5454
unsafe
5555
{
@@ -60,6 +60,7 @@ public static void MemoryPoolSpan()
6060
}
6161
}
6262

63+
6364
[Theory]
6465
[InlineData(0)]
6566
[InlineData(3)]
@@ -72,8 +73,8 @@ public static void MemoryPoolPin(int elementIndex)
7273
Memory<int> memory = block.Memory;
7374
Span<int> sp = memory.Span;
7475
Assert.Equal(memory.Length, sp.Length);
75-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(memory, out MemoryManager<int> manager));
76-
using (MemoryHandle newMemoryHandle = manager.Pin(elementIndex: elementIndex))
76+
Assert.True(MemoryMarshal.TryGetArray(memory, out ArraySegment<int> segment));
77+
using (MemoryHandle newMemoryHandle = memory.Slice(elementIndex).Pin())
7778
{
7879
unsafe
7980
{
@@ -94,7 +95,7 @@ public static void MemoryPoolPinBadOffset(int elementIndex)
9495
Memory<int> memory = block.Memory;
9596
Span<int> sp = memory.Span;
9697
Assert.Equal(memory.Length, sp.Length);
97-
Assert.Throws<ArgumentOutOfRangeException>(() => ((MemoryManager<int>)block).Pin(elementIndex: elementIndex));
98+
Assert.Throws<ArgumentOutOfRangeException>(() => memory.Slice(elementIndex).Pin());
9899
}
99100

100101
[Fact]
@@ -107,8 +108,8 @@ public static void MemoryPoolPinOffsetAtEnd()
107108
Assert.Equal(memory.Length, sp.Length);
108109

109110
int elementIndex = memory.Length;
110-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(memory, out MemoryManager<int> manager));
111-
using (MemoryHandle newMemoryHandle = manager.Pin(elementIndex: elementIndex))
111+
Assert.True(MemoryMarshal.TryGetArray(memory, out ArraySegment<int> segment));
112+
using (MemoryHandle newMemoryHandle = memory.Slice(elementIndex).Pin())
112113
{
113114
unsafe
114115
{
@@ -128,8 +129,8 @@ public static void MemoryPoolPinBadOffsetTooLarge()
128129
Assert.Equal(memory.Length, sp.Length);
129130

130131
int elementIndex = memory.Length + 1;
131-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(memory, out MemoryManager<int> manager));
132-
Assert.Throws<ArgumentOutOfRangeException>(() => manager.Pin(elementIndex: elementIndex));
132+
Assert.True(MemoryMarshal.TryGetArray(memory, out ArraySegment<int> segment));
133+
Assert.Throws<ArgumentOutOfRangeException>(() => memory.Slice(elementIndex).Pin());
133134
}
134135

135136
[Fact]
@@ -150,13 +151,13 @@ public static void EachRentalIsUniqueUntilDisposed()
150151
IMemoryOwner<int> newBlock = pool.Rent(minBufferSize);
151152
Memory<int> memory = newBlock.Memory;
152153
Assert.True(memory.Length >= minBufferSize);
153-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(newBlock.Memory, out MemoryManager<int> newManager));
154+
Assert.True(MemoryMarshal.TryGetArray(newBlock.Memory, out ArraySegment<int> newArraySegment));
154155
foreach (IMemoryOwner<int> prior in priorBlocks)
155156
{
156-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(prior.Memory, out MemoryManager<int> priorManager));
157-
using (MemoryHandle priorMemoryHandle = priorManager.Pin())
157+
Assert.True(MemoryMarshal.TryGetArray(prior.Memory, out ArraySegment<int> priorArraySegment));
158+
using (MemoryHandle priorMemoryHandle = prior.Memory.Pin())
158159
{
159-
using (MemoryHandle newMemoryHandle = newManager.Pin())
160+
using (MemoryHandle newMemoryHandle = newBlock.Memory.Pin())
160161
{
161162
unsafe
162163
{
@@ -170,8 +171,7 @@ public static void EachRentalIsUniqueUntilDisposed()
170171

171172
foreach (IMemoryOwner<int> prior in priorBlocks)
172173
{
173-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(prior.Memory, out MemoryManager<int> priorManager));
174-
priorManager.Unpin();
174+
Assert.True(MemoryMarshal.TryGetArray(prior.Memory, out ArraySegment<int> priorArraySegment));
175175
prior.Dispose();
176176
}
177177
}
@@ -219,10 +219,10 @@ public static void MemoryPoolTryGetArray()
219219
Assert.Equal(memory.Length, arraySegment.Count);
220220
unsafe
221221
{
222-
Assert.True(MemoryMarshal.TryGetMemoryManager<int, MemoryManager<int>>(memory, out MemoryManager<int> manager));
223-
void* pSpan = Unsafe.AsPointer(ref MemoryMarshal.GetReference(manager.GetSpan()));
222+
Assert.True(MemoryMarshal.TryGetArray(memory, out arraySegment));
224223
fixed (int* pArray = arraySegment.Array)
225224
{
225+
void* pSpan = Unsafe.AsPointer(ref MemoryMarshal.GetReference(memory.Span));
226226
Assert.Equal((IntPtr)pSpan, (IntPtr)pArray);
227227
}
228228
}
@@ -233,47 +233,16 @@ public static void MemoryPoolTryGetArray()
233233
public static void ExtraDisposesAreIgnored()
234234
{
235235
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
236-
((MemoryManager<int>)block).Unpin();
237-
block.Dispose();
238-
block.Dispose();
239-
}
240-
241-
[Fact]
242-
public static void NoSpanAfterDispose()
243-
{
244-
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
245-
((MemoryManager<int>)block).Unpin();
246-
block.Dispose();
247-
Assert.Throws<ObjectDisposedException>(() => ((MemoryManager<int>)block).GetSpan().DontBox());
248-
}
249-
250-
[Fact]
251-
public static void NoPinAfterDispose()
252-
{
253-
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
254-
((MemoryManager<int>)block).Unpin();
255236
block.Dispose();
256-
Assert.Throws<ObjectDisposedException>(() => ((MemoryManager<int>)block).Pin());
257-
}
258-
259-
[Fact]
260-
public static void NoUnpin_AfterDispose()
261-
{
262-
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
263-
((MemoryManager<int>)block).Unpin();
264237
block.Dispose();
265-
Assert.Throws<ObjectDisposedException>(() => ((MemoryManager<int>)block).Unpin());
266238
}
267239

268-
[Fact]
269-
public static void NoTryGetArrayAfterDispose()
270-
{
271-
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
272-
Memory<int> memory = block.Memory;
273-
((MemoryManager<int>)block).Unpin();
274-
block.Dispose();
275-
Assert.Throws<ObjectDisposedException>(() => MemoryMarshal.TryGetArray(memory, out ArraySegment<int> arraySegment));
276-
}
240+
[Fact]
241+
public static void NoMemoryAfterDispose()
242+
{
243+
IMemoryOwner<int> block = MemoryPool<int>.Shared.Rent(42);
244+
block.Dispose();
245+
Assert.Throws<ObjectDisposedException>(() => block.Memory);
246+
}
277247
}
278248
}
279-

0 commit comments

Comments
 (0)