| | | 1 | | // Copyright (c) ZeroC, Inc. |
| | | 2 | | |
| | | 3 | | using System.Buffers; |
| | | 4 | | using System.Collections.Concurrent; |
| | | 5 | | |
| | | 6 | | namespace ZeroC.Tests.Common; |
| | | 7 | | |
| | | 8 | | /// <summary>A memory pool with "poisoned" memory and small buffers.</summary> |
| | | 9 | | public class TestMemoryPool : MemoryPool<byte> |
| | | 10 | | { |
| | | 11 | | /// <inheritdoc/> |
| | 2222 | 12 | | public override int MaxBufferSize { get; } |
| | | 13 | | |
| | 51 | 14 | | private readonly ConcurrentStack<Memory<byte>> _stack = new(); |
| | | 15 | | |
| | | 16 | | /// <inheritdoc/> |
| | | 17 | | public override IMemoryOwner<byte> Rent(int minBufferSize = -1) |
| | 738 | 18 | | { |
| | 738 | 19 | | if (minBufferSize > MaxBufferSize) |
| | 0 | 20 | | { |
| | 0 | 21 | | throw new InvalidOperationException( |
| | 0 | 22 | | $"minBufferSize {minBufferSize} is greater than the pool's MaxBufferSize"); |
| | | 23 | | } |
| | | 24 | | |
| | 738 | 25 | | return new MemoryOwner(this); |
| | 738 | 26 | | } |
| | | 27 | | |
| | | 28 | | /// <inheritdoc/> |
| | | 29 | | protected override void Dispose(bool disposing) |
| | 51 | 30 | | { |
| | | 31 | | // no-op |
| | 51 | 32 | | } |
| | | 33 | | |
| | | 34 | | /// <summary>Construct a new test memory pool.</summary> |
| | | 35 | | /// <param name="maxBufferSize">The max buffer size.</param> |
| | 102 | 36 | | public TestMemoryPool(int maxBufferSize) => MaxBufferSize = maxBufferSize; |
| | | 37 | | |
| | | 38 | | private sealed class MemoryOwner : IMemoryOwner<byte> |
| | | 39 | | { |
| | 2196 | 40 | | public Memory<byte> Memory { get; } |
| | | 41 | | |
| | | 42 | | private TestMemoryPool? _pool; |
| | | 43 | | |
| | | 44 | | public void Dispose() |
| | 720 | 45 | | { |
| | 720 | 46 | | if (_pool is not null) |
| | 720 | 47 | | { |
| | 720 | 48 | | _pool._stack.Push(Memory); |
| | 720 | 49 | | _pool = null; |
| | 720 | 50 | | } |
| | 720 | 51 | | } |
| | | 52 | | |
| | 738 | 53 | | internal MemoryOwner(TestMemoryPool pool) |
| | 738 | 54 | | { |
| | 738 | 55 | | _pool = pool; |
| | | 56 | | |
| | 738 | 57 | | if (_pool._stack.TryPop(out Memory<byte> buffer)) |
| | 0 | 58 | | { |
| | 0 | 59 | | Memory = buffer; |
| | 0 | 60 | | } |
| | | 61 | | else |
| | 738 | 62 | | { |
| | 738 | 63 | | Memory = new byte[_pool.MaxBufferSize]; |
| | 738 | 64 | | } |
| | | 65 | | |
| | | 66 | | // "poison" memory |
| | 738 | 67 | | Memory.Span.Fill(0xAA); |
| | 738 | 68 | | } |
| | | 69 | | } |
| | | 70 | | } |