Custom Native Container [Part 2]: Deallocate On Job Completion

Introduction

In the previous article in this series, Custom Native Container [Part 1]: The Basics, we looked into how we can create a bare basic custom native container for usage with the job system. In this article we will extend our NativeIntArray container to add support for usage with .WithDeallocateOnJobCompletion and [DeallocateOnJobCompletion].

The result of the previous article can be found here.
The final result of this article can be found here.

1) Enable Support

To enable support for deallocation on job completion we must add the [NativeContainerSupportsDeallocateOnJobCompletion] attribute to our container struct. We will also use this opportunity to have a quick reminder of what our container roughly looked like: a simple array of integers.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Threading;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;

// Enable support for ".WithDeallocateOnJobCompletion" and "[DeallocateOnJobCompletion]".
[NativeContainerSupportsDeallocateOnJobCompletion]
[NativeContainer] 
[StructLayout(LayoutKind.Sequential)] 
public unsafe struct NativeIntArray : IDisposable
{
    [NativeDisableUnsafePtrRestriction] internal void* m_Buffer;
    internal int m_Length;

#if ENABLE_UNITY_COLLECTIONS_CHECKS
    internal AtomicSafetyHandle m_Safety;
    [NativeSetClassTypeToNullOnSchedule] internal DisposeSentinel m_DisposeSentinel;
#endif

    internal Allocator m_AllocatorLabel;

    public NativeIntArray(int length, Allocator allocator, NativeArrayOptions options = NativeArrayOptions.ClearMemory) { /* More Code */ }

    static void Allocate(int length, Allocator allocator, out NativeIntArray array)
    {
        long size = UnsafeUtility.SizeOf<int>() * (long)length;

		/* More Code */

        array = default(NativeIntArray);
        // Allocate memory for our buffer.
        array.m_Buffer = UnsafeUtility.Malloc(size, UnsafeUtility.AlignOf<int>(), allocator);
        array.m_Length = length;
        array.m_AllocatorLabel = allocator;

		/* More Code */
    }

	/* More Code */

2) JobHandle Dispose

Next we will add a new dispose function to our container struct. This dispose function will not deallocate our container immediately, but will instead return a job handle that can be scheduled later. This is how deallocate on job completion works, by scheduling another job to do the cleanup once our job is completed.

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
    /*
	 * ... More Code ...
	 */

	public unsafe JobHandle Dispose(JobHandle inputDeps)
    {
#if ENABLE_UNITY_COLLECTIONS_CHECKS
        // DisposeSentinel needs to be cleared on the main thread.
        DisposeSentinel.Clear(ref m_DisposeSentinel);
#endif

        // Create a job to dispose of our container and pass a copy of our pointer to it.
        NativeCustomArrayDisposeJob disposeJob = new NativeCustomArrayDisposeJob()
        {
            Data = new NativeCustomArrayDispose() 
            { 
                m_Buffer = m_Buffer,
                m_AllocatorLabel = m_AllocatorLabel
            }
        };
        JobHandle result = disposeJob.Schedule(inputDeps);

#if ENABLE_UNITY_COLLECTIONS_CHECKS
        AtomicSafetyHandle.Release(m_Safety);
#endif
        
        m_Buffer = null;
        m_Length = 0;

        return result;
    }

    /*
	 * ... More Code ...
	 */

3) NativeCustomArrayDisposeJob And NativeCustomArrayDispose

As you may have noticed, inside our Dispose function we make use of two new structs. These need to be defined outside out container struct. NativeCustomArrayDispose is used to hold a copy of our container pointer and NativeCustomArrayDisposeJob will call Dispose. Whenever the job gets scheduled, it will internally make sure that no other job is reading or writing to our container.

150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
/*
* ... More Code ...
*/

[NativeContainer]
internal unsafe struct NativeCustomArrayDispose
{
    // Relax the pointer safety so jobs can schedule with this struct.
    [NativeDisableUnsafePtrRestriction] internal void* m_Buffer;
    internal Allocator m_AllocatorLabel;

    public void Dispose()
    {
        // Free the allocated memory
        UnsafeUtility.Free(m_Buffer, m_AllocatorLabel);
    }
}

[BurstCompile]
internal struct NativeCustomArrayDisposeJob : IJob
{
    internal NativeCustomArrayDispose Data;

    public void Execute()
    {
        Data.Dispose();
    }
}

Usage

And that is it! We now have added support for parallel jobs to our NativeIntArray. An example of this is shown below.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.Jobs;
using Unity.Mathematics;

public class NativeIntArraySystem : SystemBase
{
	protected override void OnUpdate()
    {
        NativeIntArray myArray = new NativeIntArray(100, Allocator.TempJob);
		Job.WithName("NativeIntArrayJob")
            .WithDeallocateOnJobCompletion(myArray)
            .WithCode(() =>
            {
                for (int i = 0; i < myArray.Length; i++)
                    myArray.Increment(i);
            }).Run();
	}
}

Conclusion

This article showed how to add support for .WithDeallocateOnJobCompletion and [DeallocateOnJobCompletion]. In the next part(s) we will continue to add features to our NativeIntArray by supporting parallel jobs.

Custom Native Container [Part 1]: The Basics
Custom Native Container [Part 2]: Deallocate On Job Completion
Custom Native Container [Part 3]: Parallel Job Using Min Max
Custom Native Container [Part 4]: Parallel Job Using ParallelWriter
Custom Native Container [Part 5]: ParallelFor Using ParallelWriter With Thread Index