Updated the SegmentedPool to allow for zeroing of allocations and for returning a pointer to allocated memory to allow for use in unsafe structs (#2)
This commit was merged in pull request #2.
This commit is contained in:
516
UnmangedMMU/WorkspaceHeap.cs
Normal file
516
UnmangedMMU/WorkspaceHeap.cs
Normal file
@@ -0,0 +1,516 @@
|
||||
namespace UnmanagedMMU
|
||||
{
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Runtime.CompilerServices;
|
||||
using System.Runtime.InteropServices;
|
||||
using UnmanagedMMU.Allocators;
|
||||
using UnmanagedMMU.Handles;
|
||||
using UnmanagedMMU.Handles.Internal;
|
||||
|
||||
/// <summary>
|
||||
/// Provides an unmanaged heap for long-lived allocations with reuse.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// <see cref="WorkspaceHeap"/> minimizes calls to the underlying allocator by retaining
|
||||
/// freed blocks in size-segregated free lists and reusing them when possible.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// Allocation strategy:
|
||||
/// <list type="bullet">
|
||||
/// <item><description>
|
||||
/// <b>Small allocations</b> (≤ 1 KB) use fixed size buckets
|
||||
/// </description></item>
|
||||
/// <item><description>
|
||||
/// <b>Medium allocations</b> (≤ 256 KB) use best-fit reuse
|
||||
/// </description></item>
|
||||
/// <item><description>
|
||||
/// <b>Large allocations</b> (> 256 KB) use tolerance-based reuse
|
||||
/// </description></item>
|
||||
/// </list>
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
public unsafe sealed class WorkspaceHeap : IDisposable, IUnmanagedMemoryOwner
|
||||
{
|
||||
/// <summary>
|
||||
/// The maximum size, in bytes, for "small" allocations. Uses fixed-size buckets
|
||||
/// </summary>
|
||||
private const nuint _smallThreshold = 1024;
|
||||
|
||||
/// <summary>
|
||||
/// The maximum size, in bytes, for allocations considered "medium". Uses best-fit reuse
|
||||
/// </summary>
|
||||
private const nuint _mediumThreshold = 256 * 1024; // 256 KB
|
||||
|
||||
/// <summary>
|
||||
/// The maximum absolute number of bytes that may be wasted when reusing a large block for a "large" allocation.
|
||||
/// </summary>
|
||||
private const nuint _largeMaxWasteBytes = 256 * 1024; // 256 KB
|
||||
|
||||
/// <summary>
|
||||
/// The maximum allowed size ratio when reusing a large allocation block.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// For example, a value of <c>1.25</c> allows a block up to 25% larger than
|
||||
/// the requested size to be reused.
|
||||
/// </remarks>
|
||||
private const double _largeWasteRatioLimit = 1.25;
|
||||
|
||||
/// <summary>
|
||||
/// Predefined bucket sizes used for small allocation reuse.
|
||||
/// </summary>
|
||||
private static readonly nuint[] _sizeClasses =
|
||||
{
|
||||
32, 64, 96, 128, 160, 192, 224, 256,
|
||||
288, 320, 352, 384, 416, 448, 480, 512,
|
||||
544, 576, 608, 640, 672, 704, 736, 768,
|
||||
800, 832, 864, 896, 928, 960, 992, 1024
|
||||
};
|
||||
|
||||
/// <summary>
|
||||
/// Allocator interface used for all underlying unmanaged memory operations.
|
||||
/// </summary>
|
||||
private readonly IUnmanagedAllocator _allocator;
|
||||
|
||||
/// <summary>
|
||||
/// Internal lock, ensures thread safety while maintaining a simple interface
|
||||
/// </summary>
|
||||
private readonly Lock _lock = new();
|
||||
|
||||
/// <summary>
|
||||
/// Free lists for small allocations, keyed by the bucket size, in bytes.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// For a given bucket, the corresponding stack contains previously allocated blocks that are available to be used
|
||||
/// </remarks>
|
||||
private readonly Dictionary<nuint, Stack<IntPtr>> _smallFree = new();
|
||||
|
||||
/// <summary>
|
||||
/// Free lists for medium allocations keyed by the exactly allocated size, in bytes, and sorted for best-fit.
|
||||
/// </summary>
|
||||
private readonly SortedDictionary<nuint, Stack<IntPtr>> _mediumFree = new();
|
||||
|
||||
/// <summary>
|
||||
/// Free lists for large allocations keyed by the exactly allocated size, in bytes, sorted for tolerance-based reuse.
|
||||
/// </summary>
|
||||
private readonly SortedDictionary<nuint, Stack<IntPtr>> _largeFree = new();
|
||||
|
||||
/// <summary>
|
||||
/// Tracks the total bytes of memory reserved from the provided <see cref="IUnmanagedAllocator"/>.
|
||||
/// </summary>
|
||||
private nuint _totalReserved;
|
||||
|
||||
/// <summary>
|
||||
/// Total memory currently in use by active allocations, in bytes.
|
||||
/// </summary>
|
||||
private nuint _totalInUse;
|
||||
|
||||
/// <summary>
|
||||
/// Counts actual underlying OS allocations.
|
||||
/// </summary>
|
||||
private nuint _totalAllocations;
|
||||
|
||||
/// <summary>
|
||||
/// Indicates whether this <see cref="WorkspaceHeap"/> has been disposed.
|
||||
/// </summary>
|
||||
private volatile bool _disposed;
|
||||
|
||||
/// <summary>
|
||||
/// Internal header prepended to each allocation to track its size.
|
||||
/// </summary>
|
||||
[StructLayout(LayoutKind.Sequential)]
|
||||
private struct BlockHeader
|
||||
{
|
||||
/// <summary>
|
||||
/// Size of the allocation in bytes.
|
||||
/// </summary>
|
||||
public nuint Size;
|
||||
|
||||
/// <summary>
|
||||
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
|
||||
/// </summary>
|
||||
private readonly nuint _pad1;
|
||||
|
||||
/// <summary>
|
||||
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
|
||||
/// </summary>
|
||||
private readonly nuint _pad2;
|
||||
|
||||
/// <summary>
|
||||
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
|
||||
/// </summary>
|
||||
private readonly nuint _pad3;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Creates a new <see cref="WorkspaceHeap"/>.
|
||||
/// </summary>
|
||||
public WorkspaceHeap()
|
||||
: this(new DefaultUnmanagedAllocator())
|
||||
{
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Creates a new <see cref="WorkspaceHeap"/> using the specified <see cref="IUnmanagedAllocator"/>.
|
||||
/// </summary>
|
||||
/// <param name="allocator">Allocator implementing <see cref="IUnmanagedAllocator"/>.</param>
|
||||
internal WorkspaceHeap(IUnmanagedAllocator allocator)
|
||||
{
|
||||
_allocator = allocator;
|
||||
|
||||
// Initialize small-size buckets
|
||||
foreach (var size in _sizeClasses)
|
||||
_smallFree[size] = new Stack<IntPtr>();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets the total number of bytes currently allocated from the underlying allocator.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This includes both active allocations and freed blocks retained for reuse.
|
||||
/// </remarks>
|
||||
public nuint TotalReservedBytes
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_lock)
|
||||
return _totalReserved;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets the total number of bytes currently in use by active allocations.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This value decreases when memory is freed and increases when new allocations occur.
|
||||
/// </remarks>
|
||||
public nuint TotalUsedBytes
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_lock)
|
||||
{
|
||||
return _totalInUse;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets the total number of allocation operations performed by this heap.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This counts new underlying OS allocations, not reuse from free lists.
|
||||
/// Useful for performance diagnostics and testing reuse behavior.
|
||||
/// </remarks>
|
||||
public nuint TotalAllocationCount
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_lock)
|
||||
{
|
||||
return _totalAllocations;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Indicates whether the <see cref="WorkspaceHeap"/> has been disposed.
|
||||
/// </summary>
|
||||
public bool IsDisposed
|
||||
{
|
||||
get { return _disposed; }
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Determines the small size bucket for a requested allocation.
|
||||
/// </summary>
|
||||
/// <param name="size">The allocation size to get the bucket size for</param>
|
||||
/// <returns>The small size bucket for the requested allocation</returns>
|
||||
private static nuint GetSizeClass(nuint size)
|
||||
{
|
||||
foreach (nuint s in _sizeClasses)
|
||||
{
|
||||
if (size <= s)
|
||||
{
|
||||
return s;
|
||||
}
|
||||
}
|
||||
return size;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Allocates a new block from the underlying allocator including a header.
|
||||
/// </summary>
|
||||
/// <param name="payloadSize">Requested payload size in bytes.</param>
|
||||
/// <returns>
|
||||
/// Pointer to the allocated block (header included).
|
||||
/// </returns>
|
||||
private IntPtr AllocateNew(nuint payloadSize)
|
||||
{
|
||||
nuint total = payloadSize + (nuint)sizeof(BlockHeader);
|
||||
void* raw = _allocator.Alloc(total);
|
||||
_totalReserved += total;
|
||||
_totalAllocations++;
|
||||
|
||||
return (IntPtr)raw;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Allocates unmanaged memory from the workspace heap.
|
||||
/// </summary>
|
||||
/// <param name="count">Number of elements <typeparamref name="T"/> to allocate.</param>
|
||||
/// <param name="zero">If true, memory is zero-initialized.</param>
|
||||
/// <returns> An <see cref="IMemoryHandle{T}"/> to the allocated memory.</returns>
|
||||
/// <exception cref="ObjectDisposedException">Thrown if heap is disposed.</exception>
|
||||
/// <exception cref="ArgumentOutOfRangeException">Thrown if size is zero.</exception>
|
||||
public IMemoryHandle<T> Allocate<T>(int count, bool zero = false) where T : unmanaged
|
||||
{
|
||||
ArgumentOutOfRangeException.ThrowIfNegative(count);
|
||||
ArgumentOutOfRangeException.ThrowIfZero(count);
|
||||
|
||||
if ((nuint)count > nuint.MaxValue / (nuint)(sizeof(T)))
|
||||
{
|
||||
throw new OverflowException($"Requested allocation of {count} elements of type {typeof(T)} exceeds allowable maximum memory size.");
|
||||
}
|
||||
|
||||
nuint size = (nuint)count * (nuint)sizeof(T);
|
||||
|
||||
lock (_lock)
|
||||
{
|
||||
ThrowIfDisposed();
|
||||
void* ptr = null;
|
||||
if (size <= _smallThreshold)
|
||||
{
|
||||
ptr = AllocateSmall(size, zero);
|
||||
}
|
||||
else if (size <= _mediumThreshold)
|
||||
{
|
||||
ptr = AllocateMedium(size, zero);
|
||||
}
|
||||
else
|
||||
{
|
||||
ptr = AllocateLarge(size, zero);
|
||||
}
|
||||
|
||||
return new PersistentMemoryHandle<T>((T*)ptr, size, this);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Allocates a small-size block using bucketed free lists.</summary>
|
||||
private void* AllocateSmall(nuint size, bool zero)
|
||||
{
|
||||
nuint bucket = GetSizeClass(size);
|
||||
Stack<IntPtr> stack = _smallFree[bucket];
|
||||
|
||||
IntPtr block = stack.Count > 0
|
||||
? stack.Pop()
|
||||
: AllocateNew(bucket);
|
||||
|
||||
BlockHeader* header = (BlockHeader*)block;
|
||||
header->Size = bucket;
|
||||
|
||||
_totalInUse += bucket;
|
||||
|
||||
void* user = header + 1;
|
||||
|
||||
if (zero)
|
||||
Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
|
||||
|
||||
return user;
|
||||
}
|
||||
|
||||
/// <summary>Allocates a medium-size block using best-fit reuse.</summary>
|
||||
private void* AllocateMedium(nuint size, bool zero)
|
||||
{
|
||||
foreach (var kv in _mediumFree)
|
||||
{
|
||||
if (kv.Key >= size && kv.Value.Count > 0)
|
||||
{
|
||||
var block = kv.Value.Pop();
|
||||
var header = (BlockHeader*)block;
|
||||
_totalInUse += header->Size;
|
||||
|
||||
void* user = header + 1;
|
||||
if (zero)
|
||||
Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
|
||||
|
||||
return user;
|
||||
}
|
||||
}
|
||||
|
||||
var newBlock = AllocateNew(size);
|
||||
var newHeader = (BlockHeader*)newBlock;
|
||||
newHeader->Size = size;
|
||||
_totalInUse += size;
|
||||
|
||||
void* newUser = newHeader + 1;
|
||||
if (zero)
|
||||
{
|
||||
Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
|
||||
}
|
||||
|
||||
|
||||
return newUser;
|
||||
}
|
||||
|
||||
/// <summary>Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).</summary>
|
||||
private void* AllocateLarge(nuint size, bool zero)
|
||||
{
|
||||
foreach (var kv in _largeFree)
|
||||
{
|
||||
nuint blockSize = kv.Key;
|
||||
if (blockSize < size || kv.Value.Count == 0)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
nuint waste = blockSize - size;
|
||||
bool acceptable =
|
||||
waste <= _largeMaxWasteBytes ||
|
||||
((double)blockSize / size) <= _largeWasteRatioLimit;
|
||||
|
||||
if (!acceptable)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
var block = kv.Value.Pop();
|
||||
var header = (BlockHeader*)block;
|
||||
_totalInUse += header->Size;
|
||||
|
||||
void* user = header + 1;
|
||||
if (zero)
|
||||
{
|
||||
Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
|
||||
}
|
||||
|
||||
return user;
|
||||
}
|
||||
|
||||
var newBlock = AllocateNew(size);
|
||||
var newHeader = (BlockHeader*)newBlock;
|
||||
newHeader->Size = size;
|
||||
_totalInUse += size;
|
||||
|
||||
void* newUser = newHeader + 1;
|
||||
if (zero)
|
||||
{
|
||||
Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
|
||||
}
|
||||
|
||||
return newUser;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Frees a previously allocated block, returning it to the appropriate free list.
|
||||
/// </summary>
|
||||
void IUnmanagedMemoryOwner.Free(IOwnedHandle handle)
|
||||
{
|
||||
ThrowIfDisposed();
|
||||
if (handle.Pointer == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (handle.GetOwner() != this)
|
||||
{
|
||||
throw new InvalidOperationException(
|
||||
"Attempted to free a handle from a different allocator pool.");
|
||||
}
|
||||
|
||||
lock (_lock)
|
||||
{
|
||||
var header = ((BlockHeader*)handle.Pointer) - 1;
|
||||
nuint size = header->Size;
|
||||
_totalInUse -= size;
|
||||
|
||||
if (size <= _smallThreshold)
|
||||
_smallFree[size].Push((IntPtr)header);
|
||||
else if (size <= _mediumThreshold)
|
||||
{
|
||||
if (!_mediumFree.TryGetValue(size, out var stack))
|
||||
_mediumFree[size] = stack = new Stack<IntPtr>();
|
||||
stack.Push((IntPtr)header);
|
||||
}
|
||||
else
|
||||
{
|
||||
if (!_largeFree.TryGetValue(size, out var stack))
|
||||
_largeFree[size] = stack = new Stack<IntPtr>();
|
||||
stack.Push((IntPtr)header);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Releases all unused blocks back to the underlying allocator.
|
||||
/// </summary>
|
||||
public void Prune()
|
||||
{
|
||||
ThrowIfDisposed();
|
||||
lock (_lock)
|
||||
{
|
||||
PruneDictionary(_smallFree);
|
||||
PruneDictionary(_mediumFree);
|
||||
PruneDictionary(_largeFree);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Helper to free all blocks in a dictionary of free stacks.</summary>
|
||||
private void PruneDictionary(IDictionary<nuint, Stack<IntPtr>> dict)
|
||||
{
|
||||
foreach (var kv in dict)
|
||||
{
|
||||
var stack = kv.Value;
|
||||
while (stack.Count > 0)
|
||||
{
|
||||
var block = stack.Pop();
|
||||
var header = (BlockHeader*)block;
|
||||
_allocator.Free((void*)block);
|
||||
_totalReserved -= (header->Size + (nuint)sizeof(BlockHeader));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Releases all memory and marks the heap as disposed.
|
||||
/// </summary>
|
||||
public void Dispose()
|
||||
{
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
lock (_lock)
|
||||
{
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
if (_totalInUse > 0)
|
||||
{
|
||||
throw new InvalidOperationException(
|
||||
"Cannot dispose WorkspaceHeap while active allocations exist. " +
|
||||
"Dispose all handles returned from this heap before disposing the heap.");
|
||||
}
|
||||
Prune();
|
||||
// Reset stats
|
||||
_totalInUse = 0;
|
||||
_totalAllocations = 0;
|
||||
_disposed = true;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Throws an <see cref="ObjectDisposedException"/> if the <see cref="WorkspaceHeap"/> has already been disposed.
|
||||
/// </summary>
|
||||
/// <exception cref="ObjectDisposedException">
|
||||
/// Thrown when this instance is no longer valid for use.
|
||||
/// </exception>
|
||||
private void ThrowIfDisposed()
|
||||
{
|
||||
ObjectDisposedException.ThrowIf(_disposed, this);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user