Define a custom allocator

You can use a custom allocator for specific memory allocation needs. To create a custom allocator, it must contain an allocator handle of type AllocatorManager.AllocatorHandle and implement the interface, AllocatorManager.IAllocator. After you create a custom allocator, you need to register it in a global allocator table in AllocatorManager.

Add the AllocatorManager.AllocatorHandle type to the custom allocator

A custom allocator must contain an allocator handle of type AllocatorManager.AllocatorHandle. An allocator handle includes the following:

Version: A 2 byte unsigned version number. Only the lower 15 bits are valid.
Index: A 2 byte unsigned index of the global allocator table obtained during registration.
Method to add a safety handle to the list of child safety handles of the allocator handle.
Method to add a child allocator to the list of child allocators of the allocator handle.
A rewind method to invalidate and unregister all the child allocators, invalidate all the child safety handles of the allocator handle, and increment the allocator handle’ Version and OfficialVersion.

Implement AllocatorManager.IAllocator interface

To define a custom allocator, you must implement the AllocatorManager.IAllocator interface.

Because AllocatorManager.IAllocator implements IDisposable, your custom allocator must implement the Dispose method.

The following is an example of how to set up the IAllocator interface and its required properties except the Try and AllocatorFunction method:

// 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();
    }
}

The Try method tells a custom allocator how to allocate or deallocate memory. The following is an example of the Try method where a custom allocator allocates memory from Allocator.Persistent, initializes the allocated memory with a user configured value, and increments an allocation count. The custom allocator also decrements the allocation count when deallocating the allocated memory.

byte m_initialValue;

int m_allocationCount;

public unsafe int Try(ref AllocatorManager.Block block)
{
    int error = 0;

    if (block.Range.Pointer == IntPtr.Zero)
    {
        // Allocate memory from Allocator.Persistent 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;

        if (error != 0)
            return error;

        // If allocation succeeds, initialize 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;
    }
    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;

        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;
    }
}

Example method AllocatorFunction below shows an allocator function of the custom allocator.

// Implement the allocator function of delegate AllocatorManager.TryFunction that is
// required when register the allocator on the global allocator table
[BurstCompile]
[MonoPInvokeCallback(typeof(AllocatorManager.TryFunction))]
public static unsafe int AllocatorFunction(IntPtr customAllocatorPtr, ref AllocatorManager.Block block)
{
    return ((ExampleCustomAllocator*)customAllocatorPtr)->Try(ref block);
}

Global allocator table

The global allocator table in AllocatorManager stores all the necessary information for custom allocators to work. When you instantiate a custom allocator, you must register the allocator in the global allocator table. The table stores the following information:

A pointer to the custom allocator instance.
A pointer to the allocator function of the custom allocator instance.
The current official version of the custom allocator instance, stored in the lower 15 bits of a 2 byte unsigned integer value.
A list of child safety handles of native containers that are created using the custom allocator instance.
A list of child allocators that are allocated using the custom allocator instance.
A bit flag indicating whether the custom allocator is able to dispose individual allocations.

Custom allocator example

The following is an example of a custom allocator that has an AllocatorManager.AllocatorHandle and initializes the allocated memory with a user-configured value and increments the allocation count. It also uses AllocatorManager.TryFunction to register the allocator on the global allocator table:

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;

// Example custom allocator. The allocator can allocate memory from Allocator.Persistant, and
// if successful, initializes the allocated memory with a user configured value and increments an
// allocation count. The allocator can also deallocate memory and, if successful, decrements
// the allocation count.
[BurstCompile(CompileSynchronously = true)]
internal struct ExampleCustomAllocator : AllocatorManager.IAllocator
{
    AllocatorManager.AllocatorHandle m_handle;

    public AllocatorManager.TryFunction Function => AllocatorFunction;

    public AllocatorManager.AllocatorHandle Handle { get { return m_handle; } set { m_handle = value; } }

    public Allocator ToAllocator { get { return m_handle.ToAllocator; } }

    public bool IsCustomAllocator { get { return m_handle.IsCustomAllocator; } }

    public bool IsAutoDispose { get { return false; } }

    public void Dispose()
    {
        // Make sure no memory leaks
        Assert.AreEqual(0, m_allocationCount);

        m_handle.Dispose();
    }

    #region allocator-custom-try
    byte m_initialValue;

    int m_allocationCount;

    public unsafe int Try(ref AllocatorManager.Block block)
    {
        int error = 0;

        if (block.Range.Pointer == IntPtr.Zero)
        {
            // Allocate memory from Allocator.Persistent 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;

            if (error != 0)
                return error;

            // If allocation succeeds, initialize 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;
        }
        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;

            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;
        }
    }

Additional resources

Create and use a custom allocator

Rewindable allocator

Create and use a custom allocator