Skip to content

Commit

Permalink
Use native allocation instead of ArrayPool when possible
Browse files Browse the repository at this point in the history
  • Loading branch information
sakno committed Jan 13, 2024
1 parent 815c03c commit abde7c5
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 8 deletions.
41 changes: 33 additions & 8 deletions src/DotNext/Buffers/SpanOwner.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@

namespace DotNext.Buffers;

using Intrinsics = Runtime.Intrinsics;

/// <summary>
/// Represents the memory obtained from the pool or allocated
/// on the stack or heap.
Expand Down Expand Up @@ -47,10 +49,7 @@ public ref struct SpanOwner<T>
/// <param name="span">The span that references the memory to rent.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public SpanOwner(Span<T> span)
{
memory = span;
owner = null;
}
=> memory = span;

/// <summary>
/// Rents the memory referenced by the span.
Expand Down Expand Up @@ -96,18 +95,37 @@ public SpanOwner(MemoryPool<T> pool)
}

/// <summary>
/// Rents the memory from <see cref="ArrayPool{T}.Shared"/>.
/// Rents the memory from <see cref="ArrayPool{T}.Shared"/>, if <typeparamref name="T"/>
/// contains at least one field of reference type; or use <see cref="NativeMemory"/>.
/// </summary>
/// <param name="minBufferSize">The minimum size of the memory to rent.</param>
/// <param name="exactSize"><see langword="true"/> to return the buffer of <paramref name="minBufferSize"/> length; otherwise, the returned buffer is at least of <paramref name="minBufferSize"/>.</param>
/// <exception cref="ArgumentOutOfRangeException"><paramref name="minBufferSize"/> is less than or equal to zero.</exception>
public SpanOwner(int minBufferSize, bool exactSize = true)
{
var owner = ArrayPool<T>.Shared.Rent(minBufferSize);
memory = exactSize ? new(owner, 0, minBufferSize) : new(owner);
this.owner = owner;
if (UseNativeAllocation)
{
ArgumentOutOfRangeException.ThrowIfNegativeOrZero(minBufferSize);

unsafe
{
var ptr = NativeMemory.Alloc((uint)minBufferSize, (uint)Unsafe.SizeOf<T>());
memory = new(ptr, minBufferSize);
}

owner = Sentinel.Instance;
}
else
{
var owner = ArrayPool<T>.Shared.Rent(minBufferSize);
memory = exactSize ? new(owner, 0, minBufferSize) : new(owner);
this.owner = owner;
}
}

private static bool UseNativeAllocation
=> !LibrarySettings.DisableNativeAllocation && !RuntimeHelpers.IsReferenceOrContainsReferences<T>() && Intrinsics.AlignOf<T>() <= nuint.Size;

/// <summary>
/// Gets the rented memory.
/// </summary>
Expand Down Expand Up @@ -163,6 +181,13 @@ public void Dispose()
{
ArrayPool<T>.Shared.Return(array, RuntimeHelpers.IsReferenceOrContainsReferences<T>());
}
else if (ReferenceEquals(owner, Sentinel.Instance))
{
unsafe
{
NativeMemory.Free(Unsafe.AsPointer(ref MemoryMarshal.GetReference(memory)));
}
}
else
{
Unsafe.As<IDisposable>(owner)?.Dispose();
Expand Down
14 changes: 14 additions & 0 deletions src/DotNext/LibrarySettings.cs
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,18 @@ internal static bool DisableRandomStringInternalBufferCleanup
return result;
}
}

internal static bool DisableNativeAllocation
{
get
{
const string switchName = "DotNext.Buffers.DisableNativeAllocation";
const bool defaultValue = false;

if (!AppContext.TryGetSwitch(switchName, out var result))
result = defaultValue;

return result;
}
}
}

0 comments on commit abde7c5

Please sign in to comment.