Enhance JobScheduler and related classes

Added XML documentation comments to the `JobScheduler` class and its methods.
Added a new method `GetJobStatus` in the `JobScheduler` class for job status retrieval.
Added a new `CollectionHandle` struct for collection management.
Added a new test class `TestUnsafeSparseSet` with unit tests for `UnsafeSparseSet`.

Changed the `WorkerThread` class to improve job retrieval logic with a new `FindJob` method.
Changed the `DynamicArena` class by removing commented-out code to streamline memory management.

Removed commented-out code in the `WorkerThread` class for improved readability.
Removed the `ArenaAllocator` struct from `AllocationManager` to clean up unused code.
Removed the `ParallelWriter` struct from `UnsafeSparseSet`, indicating a shift in handling sparse sets.

Misaki.HighPerformance.Jobs/JobScheduler.cs

@@ -6,6 +6,15 @@ using System.Runtime.CompilerServices;
 
 namespace Misaki.HighPerformance.Jobs;
 
+/// <summary>
+/// Provides a mechanism for scheduling and executing jobs across multiple worker threads.
+/// </summary>
+/// <remarks>The <see cref="JobScheduler"/> class is designed to manage the execution of jobs, including support
+/// for dependencies, parallel execution, and thread-specific job assignment. It allows developers to schedule jobs that
+/// implement the <see cref="IJob"/> or <see cref="IJobParallelFor"/> interfaces, and it ensures efficient utilization
+/// of worker threads through job batching and work-stealing mechanisms.  This class is thread-safe and can be used in
+/// multi-threaded environments. However, it must be disposed when no longer needed to release resources and terminate
+/// worker threads.</remarks>
 public unsafe sealed class JobScheduler : IDisposable
 {
     private FreeList _jobDataAllocator;
@@ -23,6 +32,10 @@ public unsafe sealed class JobScheduler : IDisposable
 
     internal bool IsCancellationRequested => _cts.IsCancellationRequested;
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="JobScheduler"/> class with the specified number of worker threads.
+    /// </summary>
+    /// <param name="threadCount">The number of worker threads to create. If less than 1, at least one thread will be created.</param>
     public JobScheduler(int threadCount)
     {
         _jobDataAllocator = new(8);
@@ -333,7 +346,7 @@ public unsafe sealed class JobScheduler : IDisposable
     /// </summary>
     /// <param name="dependencies">A collection of <see cref="JobHandle"/> instances representing the dependencies to combine.</param>
     /// <returns>A <see cref="JobHandle"/> that represents the combined dependencies. The returned handle can be used to ensure
-    /// that all specified dependencies are completed before proceeding.</returns>
+    ///     that all specified dependencies are completed before proceeding.</returns>
     public JobHandle CombineDependencies(params ReadOnlySpan<JobHandle> dependencies)
     {
         var jobInfo = new JobInfo
@@ -350,6 +363,28 @@ public unsafe sealed class JobScheduler : IDisposable
         return CreateJobHandle(ref jobInfo, dependencies);
     }
 
+    /// <summary>
+    /// Retrieves the current status of a job identified by the specified handle.
+    /// </summary>
+    /// <param name="handle">The handle representing the job whose status is to be retrieved. The handle must be valid.</param>
+    /// <returns>The current status of the job as a <see cref="JobStatus"/> value.
+    ///     Returns <see cref="JobStatus.Invalid"/> if the handle is invalid or the job does not exist.</returns>
+    public JobStatus GetJobStatus(JobHandle handle)
+    {
+        if (!handle.IsValid)
+        {
+            return JobStatus.Invalid;
+        }
+
+        ref var jobInfo = ref _jobInfoPool.GetElementReferenceAt(handle._id, handle._generation, out var exist);
+        if (!exist)
+        {
+            return JobStatus.Invalid;
+        }
+
+        return (JobStatus)Volatile.Read(ref Unsafe.As<JobStatus, int>(ref jobInfo.status));
+    }
+
     /// <summary>
     /// Blocks the calling thread until the specified job is completed.
     /// </summary>

Misaki.HighPerformance.Jobs/WorkerThread.cs

@@ -29,23 +29,67 @@ internal class WorkerThread : IDisposable
 
     public void Start() => _thread.Start();
 
+    private JobHandle FindJob()
+    {
+        var handle = JobHandle.Invalid;
+        if (_localQueue.TryDequeue(out handle)
+            || _scheduler.TryStealJob(-1, out handle))
+        {
+            return handle;
+        }
+
+        while (true)
+        {
+            var randomIndex = _random.Next(0, _scheduler.WorkerCount);
+            if (_scheduler.TryStealJob(randomIndex, out handle))
+            {
+                return handle;
+            }
+        }
+    }
+
     private unsafe void WorkLoop()
     {
         while (!_scheduler.IsCancellationRequested)
         {
-            var handle = JobHandle.Invalid;
+            var spinner = new SpinWait();
-
+            for (var i = 0; i < 25; i++)
-            // Always try the local thread and main thread queue first.
-            if (!_localQueue.TryDequeue(out handle)
-                && !_scheduler.TryStealJob(-1, out handle))
             {
-                var randomIndex = _random.Next(0, _scheduler.WorkerCount);
+                spinner.SpinOnce(-1);
-                if (_scheduler.TryStealJob(randomIndex, out var tempHandle))
+
+                if (_scheduler.HasWork())
                 {
-                    handle = tempHandle;
+                    // Instead of goto, we still need to go through the WaitForWork to claim a release.
+                    // This causes lock and lots of branches inside the SemaphoreSlim, which lost 0.03ms.
+                    // goto DoWork;
+                    break;
                 }
             }
 
+            try
+            {
+                _scheduler.WaitForWork();
+            }
+            catch (OperationCanceledException)
+            {
+                continue;
+            }
+
+            //var handle = JobHandle.Invalid;
+
+            //// Always try the local thread and main thread queue first.
+            //if (!_localQueue.TryDequeue(out handle)
+            //    && !_scheduler.TryStealJob(-1, out handle))
+            //{
+            //    var randomIndex = _random.Next(0, _scheduler.WorkerCount);
+            //    if (_scheduler.TryStealJob(randomIndex, out var tempHandle))
+            //    {
+            //        handle = tempHandle;
+            //    }
+            //}
+
+            //DoWork:
+            var handle = FindJob();
             ref var jobInfo = ref _scheduler.GetJobInfoReference(handle, out var exist);
 
             if (exist)
@@ -59,31 +103,6 @@ internal class WorkerThread : IDisposable
                     _scheduler.MarkJobComplete(handle);
                 }
             }
-            else
-            {
-                var spinner = new SpinWait();
-                for (var i = 0; i < 25; i++)
-                {
-                    spinner.SpinOnce(-1);
-
-                    if (_scheduler.HasWork())
-                    {
-                        goto FoundWork;
-                    }
-                }
-
-                try
-                {
-                    _scheduler.WaitForWork();
-                }
-                catch (OperationCanceledException)
-                {
-                    continue;
-                }
-            }
-
-        FoundWork:
-            ;
         }
     }
 

Misaki.HighPerformance.LowLevel/Buffer/AllocationManager.cs

@@ -51,7 +51,6 @@ public readonly unsafe struct AllocationInfo
 /// </summary>
 public static unsafe class AllocationManager
 {
-
     private unsafe struct ArenaAllocator : IAllocator, IDisposable
     {
         private DynamicArena _arena;

Misaki.HighPerformance.LowLevel/Buffer/Arena .cs

@@ -49,15 +49,6 @@ public unsafe struct Arena : IDisposable
             throw new ObjectDisposedException(nameof(DynamicArena));
         }
 
-        //var offset = _offset + alignment - 1 & ~(alignment - 1);
-        //if (offset + size > _size)
-        //{
-        //    return null;
-        //}
-
-        //_offset = offset + size;
-        //var ptr = _buffer + offset;
-
         nuint currentOffset, newOffset, alignedOffset;
 
         do

Misaki.HighPerformance.LowLevel/Buffer/DynamicArena.cs

@@ -90,22 +90,6 @@ public unsafe struct DynamicArena : IDisposable
             // Release the spinlock
             Interlocked.Exchange(ref _nodeCreationLock, 0);
         }
-
-        //var newNode = (ArenaNode*)Malloc(SizeOf<ArenaNode>());
-        //try
-        //{
-        //    newNode->arena = new Arena(size);
-        //    newNode->next = null;
-
-        //    _current->next = newNode;
-        //    _current = newNode;
-        //    return true;
-        //}
-        //catch
-        //{
-        //    Free(newNode);
-        //    return false;
-        //}
     }
 
     /// <summary>
@@ -162,9 +146,6 @@ public unsafe struct DynamicArena : IDisposable
         _current = _root;
     }
 
-    /// <summary>
-    /// Disposes all arenas and frees associated memory.
-    /// </summary>
     public void Dispose()
     {
         if (_root == null)

Misaki.HighPerformance.LowLevel/Buffer/FreeList.cs

@@ -6,25 +6,6 @@ namespace Misaki.HighPerformance.LowLevel.Buffer;
 /// <summary>
 /// A lock-free, thread-safe variable-size allocator that manages memory blocks of different sizes.
 /// Optimized for high-performance scenarios with frequent allocations and deallocations.
-/// 
-/// Example usage:
-/// <code>
-/// // Create a free list with multiple size buckets
-/// var freeList = new FreeList();
-/// 
-/// // Allocate a 70-byte block
-/// var block = freeList.Allocate(70);
-/// if (block.IsValid)
-/// {
-///     // Use the memory block...
-///     
-///     // Free the block when done
-///     freeList.Free(block);
-/// }
-/// 
-/// // Dispose when finished
-/// freeList.Dispose();
-/// </code>
 /// </summary>
 [StructLayout(LayoutKind.Explicit, Size = 256)] // Cache line aligned to prevent false sharing
 public unsafe struct FreeList : IDisposable
@@ -476,10 +457,6 @@ public unsafe struct FreeList : IDisposable
         }
     }
 
-    /// <summary>
-    /// Disposes the free list and frees all allocated memory.
-    /// Note: This method is NOT thread-safe by design as requested.
-    /// </summary>
     public void Dispose()
     {
         if (Interlocked.CompareExchange(ref _disposed, 1, 0) == 0)

Misaki.HighPerformance.LowLevel/Collections/CollectionHandle.cs

@@ -0,0 +1,47 @@
+namespace Misaki.HighPerformance.LowLevel.Collections;
+
+public readonly struct CollectionHandle
+{
+    public readonly int id;
+    public readonly int generation;
+
+    public static CollectionHandle Invalid => new(-1, -1);
+
+    public bool IsValid => this != Invalid;
+
+    internal CollectionHandle(int id, int generation)
+    {
+        this.id = id;
+        this.generation = generation;
+    }
+
+    public bool Equals(CollectionHandle other)
+    {
+        return id == other.id && generation == other.generation;
+    }
+
+    public override bool Equals(object? obj)
+    {
+        return obj is CollectionHandle handle && Equals(handle);
+    }
+
+    public override int GetHashCode()
+    {
+        return HashCode.Combine(id, generation);
+    }
+
+    public override string ToString()
+    {
+        return IsValid ? $"CollectionHandle({id}, {generation})" : "CollectionHandle(Invalid)";
+    }
+
+    public static bool operator ==(CollectionHandle left, CollectionHandle right)
+    {
+        return left.Equals(right);
+    }
+
+    public static bool operator !=(CollectionHandle left, CollectionHandle right)
+    {
+        return !(left == right);
+    }
+}

Misaki.HighPerformance.LowLevel/Collections/UnsafeSparseSet.cs

@@ -66,122 +66,6 @@ public unsafe struct UnsafeSparseSet<T> : IUnsafeCollection<T>
         }
     }
 
-    public struct ParallelWriter
-    {
-        private UnsafeSparseSet<T>* _sparseSet;
-
-        internal ParallelWriter(UnsafeSparseSet<T>* sparseSet)
-        {
-            _sparseSet = sparseSet;
-        }
-
-        /// <summary>
-        /// Adds a value to the sparse set without resizing the internal arrays.
-        /// </summary>
-        /// <param name="value">The value to add to the sparse set.</param>
-        /// <returns> Returns the sparse index assigned to the value. -1 if the sparse index is out of bounds.</returns>
-        public int AddNoResize(T value)
-        {
-            int sparseIndex;
-
-            if (_sparseSet->_freeCount > 0)
-            {
-                var index = Interlocked.Decrement(ref _sparseSet->_freeCount);
-                sparseIndex = _sparseSet->_freeList[index];
-            }
-            else
-            {
-                sparseIndex = Interlocked.Increment(ref _sparseSet->_nextId) - 1;
-            }
-
-            if (sparseIndex >= _sparseSet->_sparse.Count)
-            {
-                return -1;
-            }
-
-            var count = Interlocked.Increment(ref _sparseSet->_count) - 1;
-
-            _sparseSet->_dense[count] = value;
-            _sparseSet->_sparse[sparseIndex] = count;
-            _sparseSet->_reverse[count] = sparseIndex;
-
-
-            return sparseIndex;
-        }
-
-        /// <summary>
-        /// Attempts to add a value at the specified sparse index without resizing the underlying collection.
-        /// </summary>
-        /// <param name="sparseIndex">The index in the sparse array where the value should be added. Must be within the valid range of the sparse
-        /// array.</param>
-        /// <param name="value">The value to add to the collection.</param>
-        /// <returns><see langword="true"/> if the value was successfully added at the specified index; otherwise, <see
-        /// langword="false"/>.</returns>
-        public bool AddAtNoResize(int sparseIndex, T value)
-        {
-            if (sparseIndex < 0 || sparseIndex >= _sparseSet->_sparse.Count)
-            {
-                return false;
-            }
-
-            if (_sparseSet->Contains(sparseIndex))
-            {
-                return false;
-            }
-
-            if (_sparseSet->_count >= _sparseSet->_dense.Count)
-            {
-                return false;
-            }
-
-            var count = Interlocked.Increment(ref _sparseSet->_count) - 1;
-
-            _sparseSet->_dense[count] = value;
-            _sparseSet->_sparse[sparseIndex] = count;
-            _sparseSet->_reverse[count] = sparseIndex;
-
-            return true;
-        }
-
-        /// <summary>
-        /// Removes a value at the specified sparse index without resizing the internal arrays.
-        /// </summary>
-        /// <param name="sparseIndex">The sparse index of the value to remove.</param>
-        /// <returns>Returns <see langword="true"/> if the value was successfully removed; otherwise, <see langword="false"/>.</returns>
-        public bool RemoveNoResize(int sparseIndex)
-        {
-            if (!_sparseSet->Contains(sparseIndex))
-            {
-                return false;
-            }
-
-            var denseIndex = _sparseSet->_sparse[sparseIndex];
-            var lastIndex = _sparseSet->_count - 1;
-
-            if (denseIndex != lastIndex)
-            {
-                var lastValue = _sparseSet->_dense[lastIndex];
-                var lastSparseIndex = _sparseSet->_reverse[lastIndex];
-                _sparseSet->_dense[denseIndex] = lastValue;
-                _sparseSet->_reverse[denseIndex] = lastSparseIndex;
-                _sparseSet->_sparse[lastSparseIndex] = denseIndex;
-            }
-
-            _sparseSet->_sparse[sparseIndex] = -1;
-            if (_sparseSet->_freeCount >= _sparseSet->_freeList.Count)
-            {
-                return false;
-            }
-
-            _sparseSet->_freeList[_sparseSet->_freeCount] = sparseIndex;
-
-            Interlocked.Increment(ref _sparseSet->_freeCount);
-            Interlocked.Decrement(ref _sparseSet->_count);
-
-            return true;
-        }
-    }
-
     private UnsafeArray<T> _dense;
     private UnsafeArray<int> _sparse;
     private UnsafeArray<int> _reverse; // Maps dense index to sparse index
@@ -194,12 +78,6 @@ public unsafe struct UnsafeSparseSet<T> : IUnsafeCollection<T>
     public readonly int Capacity => _dense.Count;
     public readonly bool IsCreated => _dense.IsCreated && _sparse.IsCreated && _reverse.IsCreated && _freeList.IsCreated;
 
-    public readonly ref T this[int index]
-    {
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        get => ref _dense[index];
-    }
-
     public IEnumerator<T> GetEnumerator() => new Enumerator((UnsafeSparseSet<T>*)UnsafeUtilities.AddressOf(ref this));
     IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 

Misaki.HighPerformance.Test/UnitTest/Collections/TestUnsafeSparseSet.cs

@@ -0,0 +1,46 @@
+using Misaki.HighPerformance.LowLevel.Buffer;
+using Misaki.HighPerformance.LowLevel.Collections;
+
+namespace Misaki.HighPerformance.Test.UnitTest.Collections;
+
+[TestClass]
+public class TestUnsafeSparseSet
+{
+    private UnsafeSparseSet<int> _sparseSet;
+
+    [TestInitialize]
+    public void Initialize()
+    {
+        _sparseSet = new UnsafeSparseSet<int>(16, Allocator.Persistent);
+    }
+
+    [TestMethod]
+    public void Add()
+    {
+        var id = _sparseSet.Add(10);
+        Assert.IsTrue(_sparseSet.Contains(id));
+    }
+
+    [TestMethod]
+    public void Remove()
+    {
+        var id = _sparseSet.Add(20);
+        Assert.IsTrue(_sparseSet.Contains(id));
+
+        _sparseSet.Remove(id);
+        Assert.IsFalse(_sparseSet.Contains(id));
+    }
+
+    [TestMethod]
+    public void IndexReuse()
+    {
+        var id = _sparseSet.Add(20);
+        Assert.IsTrue(_sparseSet.Contains(id));
+
+        _sparseSet.Remove(id);
+        Assert.IsFalse(_sparseSet.Contains(id));
+
+        var newId = _sparseSet.Add(30);
+        Assert.AreEqual(id, newId);
+    }
+}