#region allocator-custom-example
using System;
using AOT;
using System.Collections.Generic;
using System.Threading.Tasks;
using NUnit.Framework;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Burst;
using System.Threading;
// This is the example code used in
// Packages/com.unity.collections/Documentation~/allocator/allocator-custom.md
// Example custom allocator. The allocator is able to allocate memory from Allocator.Persistant,
// if successful, initialize the allocated memory with a user configured value and increment an
// allocation count. The allocator is able to deallocate the memory, if successful, decrement
// the allocation count.
// A custom allocator must implement AllocatorManager.IAllocator interface
[BurstCompile(CompileSynchronously = true)]
internal struct ExampleCustomAllocator : AllocatorManager.IAllocator
{
// A custom allocator must contain AllocatorManager.AllocatorHandle
AllocatorManager.AllocatorHandle m_handle;
// Implement the Function property required by IAllocator interface
public AllocatorManager.TryFunction Function => AllocatorFunction;
// Implement the Handle property required by IAllocator interface
public AllocatorManager.AllocatorHandle Handle { get { return m_handle; } set { m_handle = value; } }
// Implement the ToAllocator property required by IAllocator interface
public Allocator ToAllocator { get { return m_handle.ToAllocator; } }
// Implement the IsCustomAllocator property required by IAllocator interface
public bool IsCustomAllocator { get { return m_handle.IsCustomAllocator; } }
// Implement the IsAutoDispose property required by IAllocator interface
// Allocations made by this example allocator are not automatically disposed.
// This implementation can be skipped because the default implementation of
// this property is false.
public bool IsAutoDispose { get { return false; } }
// Implement the Dispose method required by IDisposable interface because
// AllocatorManager.IAllocator implements IDisposable
public void Dispose()
{
// Make sure no memory leaks
Assert.AreEqual(0, m_allocationCount);
m_handle.Dispose();
}
#region allocator-custom-try
// Value to initialize the allocated memory
byte m_initialValue;
// Allocation count
int m_allocationCount;
// Implement the Try method required by IAllocator interface
public unsafe int Try(ref AllocatorManager.Block block)
{
// Error status
int error = 0;
// Allocate
if (block.Range.Pointer == IntPtr.Zero)
{
// Allocate memory from Allocator.Persistant and restore the original allocator
AllocatorManager.AllocatorHandle tempAllocator = block.Range.Allocator;
block.Range.Allocator = Allocator.Persistent;
error = AllocatorManager.Try(ref block);
block.Range.Allocator = tempAllocator;
// return if error occurs
if (error != 0)
return error;
// if allocation succeeds, intialize the memory with the initial value and increment the allocation count
if (block.Range.Pointer != IntPtr.Zero)
{
UnsafeUtility.MemSet((void*)block.Range.Pointer, m_initialValue, block.Bytes);
Interlocked.Increment(ref m_allocationCount);
}
return 0;
}
// Deallocate
else
{
// Deallocate memory from Allocator.Persistant and restore the original allocator
AllocatorManager.AllocatorHandle tempAllocator = block.Range.Allocator;
block.Range.Allocator = Allocator.Persistent;
error = AllocatorManager.Try(ref block);
block.Range.Allocator = tempAllocator;
// return if error occurs
if (error != 0)
return error;
// if deallocation succeeds, decrement the allocation count
if (block.Range.Pointer == IntPtr.Zero)
{
Interlocked.Decrement(ref m_allocationCount);
}
return 0;
}
}
#endregion // allocator-custom-try
#region allocator-custom-allocator-function
// Implement the allocator function of delegate AllocatorManager.TryFunction that is
// required when register the allocator on the global allocator table
[BurstCompile(CompileSynchronously = true)]
[MonoPInvokeCallback(typeof(AllocatorManager.TryFunction))]
public static unsafe int AllocatorFunction(IntPtr customAllocatorPtr, ref AllocatorManager.Block block)
{
return ((ExampleCustomAllocator*)customAllocatorPtr)->Try(ref block);
}
#endregion // allocator-custom-allocator-function
// Property to get the initial value
public byte InitialValue => m_initialValue;
// Property to get the allocation count
public int AllocationCount => m_allocationCount;
// Initialize the allocator
public void Initialize(byte initialValue)
{
m_initialValue = initialValue;
m_allocationCount = 0;
}
}
#endregion // allocator-custom-example
#region allocator-custom-user-struct
// Example user structure that contains the custom allocator
internal struct ExampleCustomAllocatorStruct
{
// Use AllocatorHelper to help creating the example custom alloctor
AllocatorHelper<ExampleCustomAllocator> customAllocatorHelper;
// Custom allocator property for accessibility
public ref ExampleCustomAllocator customAllocator => ref customAllocatorHelper.Allocator;
// Create the example custom allocator
void CreateCustomAllocator(AllocatorManager.AllocatorHandle backgroundAllocator, byte initialValue)
{
// Allocate the custom allocator from backgroundAllocator and register the allocator
customAllocatorHelper = new AllocatorHelper<ExampleCustomAllocator>(backgroundAllocator);
// Set the initial value to initialize the memory
customAllocator.Initialize(initialValue);
}
#region allocator-custom-dispose
// Dispose the custom allocator
void DisposeCustomAllocator()
{
// Dispose the custom allocator
customAllocator.Dispose();
// Unregister the custom allocator and dispose it
customAllocatorHelper.Dispose();
}
#endregion // allocator-custom-dispose
// Constructor of user structure
public ExampleCustomAllocatorStruct(byte initialValue)
{
this = default;
CreateCustomAllocator(Allocator.Persistent, initialValue);
}
// Dispose the user structure
public void Dispose()
{
DisposeCustomAllocator();
}
#region allocator-custom-use
// Sample code to use the custom allocator to allocate containers
public void UseCustomAllocator(out NativeArray<int> nativeArray, out NativeList<int> nativeList)
{
// Use custom allocator to allocate a native array and check initial value.
nativeArray = CollectionHelper.CreateNativeArray<int, ExampleCustomAllocator>(100, ref customAllocator, NativeArrayOptions.UninitializedMemory);
Assert.AreEqual(customAllocator.InitialValue, (byte)nativeArray[0] & 0xFF);
nativeArray[0] = 0xFE;
// Use custom allocator to allocate a native list and check initial value.
nativeList = new NativeList<int>(customAllocator.Handle);
for (int i = 0; i < 50; i++)
{
nativeList.Add(i);
}
unsafe
{
// Use custom allocator to allocate a byte buffer.
byte* bytePtr = (byte*)AllocatorManager.Allocate(ref customAllocator, sizeof(byte), sizeof(byte), 10);
Assert.AreEqual(customAllocator.InitialValue, bytePtr[0]);
// Free the byte buffer.
AllocatorManager.Free(customAllocator.ToAllocator, bytePtr, 10);
}
}
#endregion // allocator-custom-use
// Get allocation count from the custom allocator
public int AllocationCount => customAllocator.AllocationCount;
public void UseCustomAllocatorHandle(out NativeArray<int> nativeArray, out NativeList<int> nativeList)
{
// Use custom allocator to allocate a native array and check initial value.
nativeArray = CollectionHelper.CreateNativeArray<int>(100, customAllocator.ToAllocator, NativeArrayOptions.UninitializedMemory);
Assert.AreEqual(customAllocator.InitialValue, (byte)nativeArray[0] & 0xFF);
nativeArray[0] = 0xFE;
// Use custom allocator to allocate a native list and check initial value.
nativeList = new NativeList<int>(customAllocator.Handle);
for (int i = 0; i < 50; i++)
{
nativeList.Add(i);
}
unsafe
{
// Use custom allocator to allocate a byte buffer.
byte* bytePtr = (byte*)AllocatorManager.Allocate(ref customAllocator, sizeof(byte), sizeof(byte), 10);
Assert.AreEqual(customAllocator.InitialValue, bytePtr[0]);
// Free the byte buffer.
AllocatorManager.Free(customAllocator.ToAllocator, bytePtr, 10);
}
}
}
internal class ExampleCustomAllocatorStructUsage
{
// Initial value for the custom allocator.
const int IntialValue = 0xAB;
// Test code.
[Test]
public void UseCustomAllocator_Works()
{
ExampleCustomAllocatorStruct exampleStruct = new ExampleCustomAllocatorStruct(IntialValue);
// Allocate native array and native list from the custom allocator
exampleStruct.UseCustomAllocator(out NativeArray<int> nativeArray, out NativeList<int> nativeList);
// Able to access the native array and native list
Assert.AreEqual(nativeArray[0], 0xFE);
Assert.AreEqual(nativeList[10], 10);
// Need to use CollectionHelper.DisposeNativeArray to dispose the native array from a custom allocator
CollectionHelper.Dispose(nativeArray) ;
// Dispose the native list
nativeList.Dispose();
#if ENABLE_UNITY_COLLECTIONS_CHECKS
// Object disposed exception throws because nativeArray is already disposed
Assert.Throws<ObjectDisposedException>(() =>
{
nativeArray[0] = 0xEF;
});
// Object disposed exception throws because nativeList is already disposed
Assert.Throws<ObjectDisposedException>(() =>
{
nativeList[10] = 0x10;
});
#endif
// Check allocation count after dispose the native array and native list
Assert.AreEqual(0, exampleStruct.AllocationCount);
// Dispose the user structure
exampleStruct.Dispose();
}
[Test]
public void UseCustomAllocatorHandle_Works()
{
ExampleCustomAllocatorStruct exampleStruct = new ExampleCustomAllocatorStruct(IntialValue);
// Allocate native array and native list from the custom allocator handle
exampleStruct.UseCustomAllocatorHandle(out NativeArray<int> nativeArray, out NativeList<int> nativeList);
// Able to access the native array and native list
Assert.AreEqual(nativeArray[0], 0xFE);
Assert.AreEqual(nativeList[10], 10);
// Need to use CollectionHelper.DisposeNativeArray to dispose the native array from a custom allocator
CollectionHelper.Dispose(nativeArray);
// Dispose the native list
nativeList.Dispose();
#if ENABLE_UNITY_COLLECTIONS_CHECKS
// Object disposed exception throws because nativeArray is already disposed
Assert.Throws<ObjectDisposedException>(() =>
{
nativeArray[0] = 0xEF;
});
// Object disposed exception throws because nativeList is already disposed
Assert.Throws<ObjectDisposedException>(() =>
{
nativeList[10] = 0x10;
});
#endif
// Check allocation count after dispose the native array and native list
Assert.AreEqual(0, exampleStruct.AllocationCount);
// Dispose the user structure
exampleStruct.Dispose();
}
[Test]
public void CustomAllocatorHandle_MultiThreadWorks()
{
ExampleCustomAllocatorStruct exampleStruct = new ExampleCustomAllocatorStruct(IntialValue);
var taskList = new List<Task>();
// create 128 native array with another threads
for (var i = 0; i < 128; i++)
{
var task = Task.Run(() =>
{
var nativeArray = CollectionHelper.CreateNativeArray<int, ExampleCustomAllocator>(100, ref exampleStruct.customAllocator,
NativeArrayOptions.UninitializedMemory);
CollectionHelper.Dispose(nativeArray);
});
taskList.Add(task);
}
Task.WaitAll(taskList.ToArray());
exampleStruct.Dispose();
}
}
#endregion // allocator-custom-user-struct