diff --git a/PowerKit.Tests/ArrayPoolExtensionsTests.cs b/PowerKit.Tests/ArrayPoolExtensionsTests.cs new file mode 100644 index 0000000..fba0a11 --- /dev/null +++ b/PowerKit.Tests/ArrayPoolExtensionsTests.cs @@ -0,0 +1,38 @@ +using System; +using System.Buffers; +using FluentAssertions; +using PowerKit.Extensions; +using Xunit; + +namespace PowerKit.Tests; + +public class ArrayPoolExtensionsTests +{ + [Fact] + public void RentOwner_Test() + { + // Arrange + var pool = ArrayPool.Shared; + + // Act + using var owner = pool.RentOwner(16); + + // Assert + owner.Memory.Length.Should().Be(16); + } + + [Fact] + public void RentOwner_Dispose_Test() + { + // Arrange + var pool = ArrayPool.Shared; + var owner = pool.RentOwner(16); + + // Act + owner.Dispose(); + var act = () => owner.Memory; + + // Assert + act.Should().Throw(); + } +} diff --git a/PowerKit/Extensions/ArrayPoolExtensions.cs b/PowerKit/Extensions/ArrayPoolExtensions.cs new file mode 100644 index 0000000..8ba2c9d --- /dev/null +++ b/PowerKit/Extensions/ArrayPoolExtensions.cs @@ -0,0 +1,44 @@ +using System; +using System.Buffers; +using System.Threading; + +namespace PowerKit.Extensions; + +internal static class ArrayPoolExtensions +{ + extension(ArrayPool pool) + { + /// + /// Rents a buffer of at least elements from the pool + /// and wraps it in an that returns the buffer to the pool + /// when disposed. + /// + public IMemoryOwner RentOwner(int minimumLength = 1) => + new ArrayPoolMemoryOwner(pool, pool.Rent(minimumLength), minimumLength); + } +} + +file sealed class ArrayPoolMemoryOwner(ArrayPool pool, T[] buffer, int minimumLength) + : IMemoryOwner +{ + private int _disposed; + + public Memory Memory + { + get + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + return buffer.AsMemory(0, minimumLength); + } + } + + public void Dispose() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + { + return; + } + + pool.Return(buffer); + } +}