Skip to content

Commit 1dfc638

Browse files
pakrympakrym
authored
Cleanup memory pool code add some tests (#344)
1 parent e49ece1 commit 1dfc638

File tree

11 files changed

+803
-218
lines changed

11 files changed

+803
-218
lines changed

‎shared/Microsoft.Extensions.Buffers.MemoryPool.Sources/DiagnosticMemoryPool.cs‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
3+
4+
using System.Collections.Generic;
5+
using System.Linq;
6+
using System.Threading.Tasks;
7+
8+
namespace System.Buffers
9+
{
10+
/// <summary>
11+
/// Used to allocate and distribute re-usable blocks of memory.
12+
/// </summary>
13+
internal class DiagnosticMemoryPool : MemoryPool<byte>
14+
{
15+
private readonly MemoryPool<byte> _pool;
16+
17+
private readonly bool _allowLateReturn;
18+
19+
private readonly bool _rentTracking;
20+
21+
private readonly object _syncObj;
22+
23+
private readonly HashSet<DiagnosticPoolBlock> _blocks;
24+
25+
private readonly List<Exception> _blockAccessExceptions;
26+
27+
private readonly TaskCompletionSource<object> _allBlocksRetuned;
28+
29+
private int _totalBlocks;
30+
31+
/// <summary>
32+
/// This default value passed in to Rent to use the default value for the pool.
33+
/// </summary>
34+
private const int AnySize = -1;
35+
36+
public DiagnosticMemoryPool(MemoryPool<byte> pool, bool allowLateReturn = false, bool rentTracking = false)
37+
{
38+
_pool = pool;
39+
_allowLateReturn = allowLateReturn;
40+
_rentTracking = rentTracking;
41+
_blocks = new HashSet<DiagnosticPoolBlock>();
42+
_syncObj = new object();
43+
_allBlocksRetuned = new TaskCompletionSource<object>();
44+
_blockAccessExceptions = new List<Exception>();
45+
}
46+
47+
public bool IsDisposed { get; private set; }
48+
49+
public override IMemoryOwner<byte> Rent(int size = AnySize)
50+
{
51+
lock (_syncObj)
52+
{
53+
if (IsDisposed)
54+
{
55+
MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPool);
56+
}
57+
58+
var diagnosticPoolBlock = new DiagnosticPoolBlock(this, _pool.Rent(size));
59+
if (_rentTracking)
60+
{
61+
diagnosticPoolBlock.Track();
62+
}
63+
_totalBlocks++;
64+
_blocks.Add(diagnosticPoolBlock);
65+
return diagnosticPoolBlock;
66+
}
67+
}
68+
69+
public override int MaxBufferSize => _pool.MaxBufferSize;
70+
71+
internal void Return(DiagnosticPoolBlock block)
72+
{
73+
bool returnedAllBlocks;
74+
lock (_syncObj)
75+
{
76+
_blocks.Remove(block);
77+
returnedAllBlocks = _blocks.Count == 0;
78+
}
79+
80+
if (IsDisposed)
81+
{
82+
if (!_allowLateReturn)
83+
{
84+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockReturnedToDisposedPool(block);
85+
}
86+
87+
if (returnedAllBlocks)
88+
{
89+
SetAllBlocksReturned();
90+
}
91+
}
92+
93+
}
94+
95+
internal void ReportException(Exception exception)
96+
{
97+
lock (_syncObj)
98+
{
99+
_blockAccessExceptions.Add(exception);
100+
}
101+
}
102+
103+
protected override void Dispose(bool disposing)
104+
{
105+
if (IsDisposed)
106+
{
107+
MemoryPoolThrowHelper.ThrowInvalidOperationException_DoubleDispose();
108+
}
109+
110+
bool allBlocksReturned = false;
111+
try
112+
{
113+
lock (_syncObj)
114+
{
115+
IsDisposed = true;
116+
allBlocksReturned = _blocks.Count == 0;
117+
if (!allBlocksReturned && !_allowLateReturn)
118+
{
119+
MemoryPoolThrowHelper.ThrowInvalidOperationException_DisposingPoolWithActiveBlocks(_totalBlocks - _blocks.Count, _totalBlocks, _blocks.ToArray());
120+
}
121+
122+
if (_blockAccessExceptions.Any())
123+
{
124+
throw CreateAccessExceptions();
125+
}
126+
}
127+
}
128+
finally
129+
{
130+
if (allBlocksReturned)
131+
{
132+
SetAllBlocksReturned();
133+
}
134+
}
135+
}
136+
137+
private void SetAllBlocksReturned()
138+
{
139+
if (_blockAccessExceptions.Any())
140+
{
141+
_allBlocksRetuned.SetException(CreateAccessExceptions());
142+
}
143+
else
144+
{
145+
_allBlocksRetuned.SetResult(null);
146+
}
147+
}
148+
149+
private AggregateException CreateAccessExceptions()
150+
{
151+
return new AggregateException("Exceptions occurred while accessing blocks", _blockAccessExceptions.ToArray());
152+
}
153+
154+
public async Task WhenAllBlocksReturnedAsync(TimeSpan timeout)
155+
{
156+
var task = await Task.WhenAny(_allBlocksRetuned.Task, Task.Delay(timeout));
157+
if (task != _allBlocksRetuned.Task)
158+
{
159+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlocksWereNotReturnedInTime(_totalBlocks - _blocks.Count, _totalBlocks, _blocks.ToArray());
160+
}
161+
162+
await task;
163+
}
164+
}
165+
}

‎shared/Microsoft.Extensions.Buffers.MemoryPool.Sources/DiagnosticPoolBlock.cs‎

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
3+
4+
using System.Threading;
5+
using System.Diagnostics;
6+
using System.Runtime.InteropServices;
7+
8+
namespace System.Buffers
9+
{
10+
/// <summary>
11+
/// Block tracking object used by the byte buffer memory pool. A slab is a large allocation which is divided into smaller blocks. The
12+
/// individual blocks are then treated as independent array segments.
13+
/// </summary>
14+
internal sealed class DiagnosticPoolBlock : MemoryManager<byte>
15+
{
16+
/// <summary>
17+
/// Back-reference to the memory pool which this block was allocated from. It may only be returned to this pool.
18+
/// </summary>
19+
private readonly DiagnosticMemoryPool _pool;
20+
21+
private readonly IMemoryOwner<byte> _memoryOwner;
22+
private MemoryHandle? _memoryHandle;
23+
private Memory<byte> _memory;
24+
25+
private readonly object _syncObj = new object();
26+
private bool _isDisposed;
27+
private int _pinCount;
28+
29+
30+
/// <summary>
31+
/// This object cannot be instantiated outside of the static Create method
32+
/// </summary>
33+
internal DiagnosticPoolBlock(DiagnosticMemoryPool pool, IMemoryOwner<byte> memoryOwner)
34+
{
35+
_pool = pool;
36+
_memoryOwner = memoryOwner;
37+
_memory = memoryOwner.Memory;
38+
}
39+
40+
public override Memory<byte> Memory
41+
{
42+
get
43+
{
44+
try
45+
{
46+
lock (_syncObj)
47+
{
48+
if (_isDisposed)
49+
{
50+
MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
51+
}
52+
53+
if (_pool.IsDisposed)
54+
{
55+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
56+
}
57+
58+
return CreateMemory(_memory.Length);
59+
}
60+
}
61+
catch (Exception exception)
62+
{
63+
_pool.ReportException(exception);
64+
throw;
65+
}
66+
}
67+
}
68+
69+
protected override void Dispose(bool disposing)
70+
{
71+
try
72+
{
73+
lock (_syncObj)
74+
{
75+
if (Volatile.Read(ref _pinCount) > 0)
76+
{
77+
MemoryPoolThrowHelper.ThrowInvalidOperationException_ReturningPinnedBlock(this);
78+
}
79+
80+
if (_isDisposed)
81+
{
82+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockDoubleDispose(this);
83+
}
84+
85+
_memoryOwner.Dispose();
86+
87+
_pool.Return(this);
88+
89+
_isDisposed = true;
90+
}
91+
}
92+
catch (Exception exception)
93+
{
94+
_pool.ReportException(exception);
95+
throw;
96+
}
97+
}
98+
99+
public override Span<byte> GetSpan()
100+
{
101+
try
102+
{
103+
lock (_syncObj)
104+
{
105+
if (_isDisposed)
106+
{
107+
MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
108+
}
109+
110+
if (_pool.IsDisposed)
111+
{
112+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
113+
}
114+
115+
return _memory.Span;
116+
}
117+
}
118+
catch (Exception exception)
119+
{
120+
_pool.ReportException(exception);
121+
throw;
122+
}
123+
}
124+
125+
public override MemoryHandle Pin(int byteOffset = 0)
126+
{
127+
try
128+
{
129+
lock (_syncObj)
130+
{
131+
if (_isDisposed)
132+
{
133+
MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
134+
}
135+
136+
if (_pool.IsDisposed)
137+
{
138+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
139+
}
140+
141+
if (byteOffset < 0 || byteOffset > _memory.Length)
142+
{
143+
MemoryPoolThrowHelper.ThrowArgumentOutOfRangeException(_memory.Length, byteOffset);
144+
}
145+
146+
_pinCount++;
147+
148+
_memoryHandle = _memoryHandle ?? _memory.Pin();
149+
150+
unsafe
151+
{
152+
return new MemoryHandle(((IntPtr)_memoryHandle.Value.Pointer + byteOffset).ToPointer(), default, this);
153+
}
154+
}
155+
}
156+
catch (Exception exception)
157+
{
158+
_pool.ReportException(exception);
159+
throw;
160+
}
161+
}
162+
163+
protected override bool TryGetArray(out ArraySegment<byte> segment)
164+
{
165+
try
166+
{
167+
lock (_syncObj)
168+
{
169+
if (_isDisposed)
170+
{
171+
MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
172+
}
173+
174+
if (_pool.IsDisposed)
175+
{
176+
MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
177+
}
178+
179+
return MemoryMarshal.TryGetArray(_memory, out segment);
180+
}
181+
}
182+
catch (Exception exception)
183+
{
184+
_pool.ReportException(exception);
185+
throw;
186+
}
187+
}
188+
189+
public override void Unpin()
190+
{
191+
try
192+
{
193+
lock (_syncObj)
194+
{
195+
if (_pinCount == 0)
196+
{
197+
MemoryPoolThrowHelper.ThrowInvalidOperationException_PinCountZero(this);
198+
}
199+
200+
_pinCount--;
201+
202+
if (_pinCount == 0)
203+
{
204+
Debug.Assert(_memoryHandle.HasValue);
205+
_memoryHandle.Value.Dispose();
206+
_memoryHandle = null;
207+
}
208+
}
209+
}
210+
catch (Exception exception)
211+
{
212+
_pool.ReportException(exception);
213+
throw;
214+
}
215+
}
216+
217+
public StackTrace Leaser { get; set; }
218+
219+
public void Track()
220+
{
221+
Leaser = new StackTrace(false);
222+
}
223+
}
224+
}

0 commit comments

Comments
 (0)