aspnetcore/test/Kestrel.Core.Tests/DiagnosticMemoryPoolTests.cs

// 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.Runtime.InteropServices;
using System.Threading.Tasks;
using Xunit;

namespace Microsoft.Extensions.Internal.Test
{
    public class DiagnosticMemoryPoolTests: MemoryPoolTests
    {
        protected override MemoryPool<byte> CreatePool() => new DiagnosticMemoryPool(new SlabMemoryPool());

        [Fact]
        public void DoubleDisposeThrows()
        {
            var memoryPool = CreatePool();
            memoryPool.Dispose();
            var exception = Assert.Throws<InvalidOperationException>(() => memoryPool.Dispose());
            Assert.Equal("Object is being disposed twice", exception.Message);
        }

        [Fact]
        public void DisposeWithActiveBlocksThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            ExpectDisposeException(memoryPool);

            var exception = Assert.Throws<InvalidOperationException>(() => block.Dispose());
            Assert.Equal("Block is being returned to disposed pool", exception.Message);
        }

        [Fact]
        public void DoubleBlockDisposeThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            block.Dispose();
            var exception = Assert.Throws<InvalidOperationException>(() => block.Dispose());
            Assert.Equal("Block is being disposed twice", exception.Message);

            ExpectDisposeAggregateException(memoryPool, exception);
        }

        [Fact]
        public void GetMemoryOfDisposedPoolThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();

            ExpectDisposeException(memoryPool);

            var exception = Assert.Throws<InvalidOperationException>(() => block.Memory);
            Assert.Equal("Block is backed by disposed slab", exception.Message);
        }

        [Fact]
        public void GetMemoryPinOfDisposedPoolThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            ExpectDisposeException(memoryPool);

            var exception = Assert.Throws<InvalidOperationException>(() => memory.Pin());
            Assert.Equal("Block is backed by disposed slab", exception.Message);
        }

        [Fact]
        public void GetMemorySpanOfDisposedPoolThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            ExpectDisposeException(memoryPool);

            var threw = false;
            try
            {
                _ = memory.Span;
            }
            catch (InvalidOperationException ode)
            {
                threw = true;
                Assert.Equal("Block is backed by disposed slab", ode.Message);
            }
            Assert.True(threw);
        }

        [Fact]
        public void GetMemoryTryGetArrayOfDisposedPoolThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            ExpectDisposeException(memoryPool);

            var exception = Assert.Throws<InvalidOperationException>(() => MemoryMarshal.TryGetArray<byte>(memory, out _));
            Assert.Equal("Block is backed by disposed slab", exception.Message);
        }

        [Fact]
        public void GetMemoryOfDisposedThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();

            block.Dispose();

            var exception = Assert.Throws<ObjectDisposedException>(() => block.Memory);
            Assert.Equal($"Cannot access a disposed object.{Environment.NewLine}Object name: 'MemoryPoolBlock'.", exception.Message);

            ExpectDisposeAggregateException(memoryPool, exception);
        }

        [Fact]
        public void GetMemoryPinOfDisposedThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            block.Dispose();

            var exception = Assert.Throws<ObjectDisposedException>(() => memory.Pin());
            Assert.Equal($"Cannot access a disposed object.{Environment.NewLine}Object name: 'MemoryPoolBlock'.", exception.Message);

            ExpectDisposeAggregateException(memoryPool, exception);
        }

        [Fact]
        public void GetMemorySpanOfDisposedThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            block.Dispose();

            Exception exception = null;
            try
            {
                _ = memory.Span;
            }
            catch (ObjectDisposedException ode)
            {
                exception = ode;
                Assert.Equal($"Cannot access a disposed object.{Environment.NewLine}Object name: 'MemoryPoolBlock'.", ode.Message);
            }
            Assert.NotNull(exception);

            ExpectDisposeAggregateException(memoryPool, exception);
        }

        [Fact]
        public void GetMemoryTryGetArrayOfDisposedThrows()
        {
            var memoryPool = CreatePool();
            var block = memoryPool.Rent();
            var memory = block.Memory;

            block.Dispose();

            var exception = Assert.Throws<ObjectDisposedException>(() => MemoryMarshal.TryGetArray<byte>(memory, out _));
            Assert.Equal($"Cannot access a disposed object.{Environment.NewLine}Object name: 'MemoryPoolBlock'.", exception.Message);

            ExpectDisposeAggregateException(memoryPool, exception);
        }

        [Fact]
        public async Task DoesNotThrowWithLateReturns()
        {
            var memoryPool = new DiagnosticMemoryPool(new SlabMemoryPool(), allowLateReturn: true);
            var block = memoryPool.Rent();
            memoryPool.Dispose();
            block.Dispose();
            await memoryPool.WhenAllBlocksReturnedAsync(TimeSpan.FromSeconds(5));
        }

        [Fact]
        public async Task ThrowsOnAccessToLateBlocks()
        {
            var memoryPool = new DiagnosticMemoryPool(new SlabMemoryPool(), allowLateReturn: true);
            var block = memoryPool.Rent();
            memoryPool.Dispose();

            var exception = Assert.Throws<InvalidOperationException>(() => block.Memory);
            Assert.Equal("Block is backed by disposed slab", exception.Message);

            block.Dispose();
            var aggregateException = await Assert.ThrowsAsync<AggregateException>(async () => await memoryPool.WhenAllBlocksReturnedAsync(TimeSpan.FromSeconds(5)));

            Assert.Equal(new Exception [] { exception }, aggregateException.InnerExceptions);
        }

        [Fact]
        public void ExceptionsContainStackTraceWhenEnabled()
        {
            var memoryPool = new DiagnosticMemoryPool(new SlabMemoryPool(), rentTracking: true);
            var block = memoryPool.Rent();

            ExpectDisposeException(memoryPool);

            var exception = Assert.Throws<InvalidOperationException>(() => block.Memory);
            Assert.Contains("Block is backed by disposed slab", exception.Message);
            Assert.Contains("ExceptionsContainStackTraceWhenEnabled", exception.Message);
        }

        private static void ExpectDisposeException(MemoryPool<byte> memoryPool)
        {
            var exception = Assert.Throws<InvalidOperationException>(() => memoryPool.Dispose());
            Assert.Contains("Memory pool with active blocks is being disposed, 0 of 1 returned", exception.Message);
        }

        private static void ExpectDisposeAggregateException(MemoryPool<byte> memoryPool, params Exception[] inner)
        {
            var exception = Assert.Throws<AggregateException>(() => memoryPool.Dispose());

            Assert.Equal(inner, exception.InnerExceptions);
        }
    }
}