Nerfed/Nerfed.Runtime/JobSystem.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using System.Runtime.ExceptionServices;

namespace Nerfed.Runtime;

public enum JobDeadlineMode
{
    FrameCritical,
    Deferred,
}

public enum JobExecutionStatus
{
    Invalid,
    Queued,
    Running,
    Completed,
    Faulted,
    Canceled,
}

public readonly struct JobHandle : IEquatable<JobHandle>
{
    internal readonly int id;

    internal JobHandle(int id)
    {
        this.id = id;
    }

    public bool IsValid => id > 0;

    public bool Equals(JobHandle other)
    {
        return id == other.id;
    }

    public override bool Equals(object obj)
    {
        return obj is JobHandle other && Equals(other);
    }

    public override int GetHashCode()
    {
        return id;
    }

    public static bool operator ==(JobHandle left, JobHandle right)
    {
        return left.Equals(right);
    }

    public static bool operator !=(JobHandle left, JobHandle right)
    {
        return !left.Equals(right);
    }

    public override string ToString()
    {
        return IsValid ? $"JobHandle({id})" : "JobHandle(Invalid)";
    }
}

public sealed class JobSystem : IDisposable
{
    public static readonly JobSystem Default = new JobSystem();

    private Thread[] workers;
    private SemaphoreSlim startSignal;
    private CountdownEvent completionEvent;

    private Thread[] asyncWorkers;
    private SemaphoreSlim asyncSignal;
    private ConcurrentQueue<int> asyncQueue;
    private ConcurrentDictionary<int, AsyncJobState> asyncJobStates;
    private int nextAsyncJobId;

    private volatile bool running;

    // Shared per-dispatch state written by main thread before workers wake.
    private volatile Action<int> currentAction;
    private int workCount;
    private int nextIndex; // grabbed with Interlocked.Increment for work-stealing

    private sealed class AsyncJobState : IDisposable
    {
        public readonly Action Action;
        public readonly JobDeadlineMode Deadline;
        public readonly string Name;
        public readonly ManualResetEventSlim Completion = new(false);

        public ExceptionDispatchInfo CapturedException;
        public volatile JobExecutionStatus Status;
        public volatile bool CancellationRequested;

        public AsyncJobState(Action action, JobDeadlineMode deadline, string name)
        {
            Action = action;
            Deadline = deadline;
            Name = name;
            Status = JobExecutionStatus.Queued;
        }

        public bool IsTerminal =>
            Status == JobExecutionStatus.Completed ||
            Status == JobExecutionStatus.Faulted ||
            Status == JobExecutionStatus.Canceled;

        public void Dispose()
        {
            Completion.Dispose();
        }
    }

    public int WorkerCount => workers?.Length ?? 0;
    public int AsyncWorkerCount => asyncWorkers?.Length ?? 0;
    public bool IsInitialized => workers != null;

    public void Initialize(int threadCount = -1, int asyncThreadCount = -1)
    {
        if (IsInitialized)
        {
            throw new InvalidOperationException("JobSystem is already initialized. Call Shutdown first.");
        }

        threadCount = threadCount < 0
            ? Math.Max(1, Environment.ProcessorCount - 2)
            : Math.Max(1, threadCount);

        running = true;
        startSignal = new SemaphoreSlim(0);
        completionEvent = new CountdownEvent(1);
        asyncSignal = new SemaphoreSlim(0);
        workers = new Thread[threadCount];
        asyncQueue = new ConcurrentQueue<int>();
        asyncJobStates = new ConcurrentDictionary<int, AsyncJobState>();

        asyncThreadCount = asyncThreadCount < 0
            ? Math.Max(1, Environment.ProcessorCount >= 8 ? 2 : 1)
            : Math.Max(0, asyncThreadCount);

        asyncWorkers = new Thread[asyncThreadCount];

        for (int i = 0; i < threadCount; i++)
        {
            workers[i] = new Thread(WorkerLoop)
            {
                IsBackground = true,
                Name = $"Job-{i}",
            };
            workers[i].Start();
        }

        for (int i = 0; i < asyncThreadCount; i++)
        {
            asyncWorkers[i] = new Thread(AsyncWorkerLoop)
            {
                IsBackground = true,
                Name = $"AsyncJob-{i}",
            };
            asyncWorkers[i].Start();
        }
    }

    public JobHandle Submit(Action action, JobDeadlineMode deadline = JobDeadlineMode.Deferred, string name = "")
    {
        ArgumentNullException.ThrowIfNull(action);

        if (!running)
        {
            throw new InvalidOperationException("JobSystem is not running. Call Initialize before submitting jobs.");
        }

        int id = Interlocked.Increment(ref nextAsyncJobId);
        var state = new AsyncJobState(action, deadline, name ?? string.Empty);

        if (!asyncJobStates.TryAdd(id, state))
        {
            state.Dispose();
            throw new InvalidOperationException($"Failed to register async job {id}.");
        }

        asyncQueue.Enqueue(id);
        asyncSignal.Release();
        return new JobHandle(id);
    }

    public JobExecutionStatus GetStatus(JobHandle handle)
    {
        if (!handle.IsValid)
        {
            return JobExecutionStatus.Invalid;
        }

        return asyncJobStates.TryGetValue(handle.id, out AsyncJobState state)
            ? state.Status
            : JobExecutionStatus.Invalid;
    }

    public bool IsCompleted(JobHandle handle)
    {
        return GetStatus(handle) is JobExecutionStatus.Completed or JobExecutionStatus.Faulted or JobExecutionStatus.Canceled;
    }

    public bool Cancel(JobHandle handle)
    {
        if (!handle.IsValid || !asyncJobStates.TryGetValue(handle.id, out AsyncJobState state))
        {
            return false;
        }

        state.CancellationRequested = true;
        return true;
    }

    public void Wait(JobHandle handle)
    {
        if (!handle.IsValid || !asyncJobStates.TryGetValue(handle.id, out AsyncJobState state))
        {
            throw new InvalidOperationException($"Unknown job handle {handle}.");
        }

        state.Completion.Wait();

        if (state.Status == JobExecutionStatus.Faulted)
        {
            state.CapturedException?.Throw();
        }
    }

    public bool TryForget(JobHandle handle)
    {
        if (!handle.IsValid || !asyncJobStates.TryGetValue(handle.id, out AsyncJobState state) || !state.IsTerminal)
        {
            return false;
        }

        if (asyncJobStates.TryRemove(handle.id, out AsyncJobState removedState))
        {
            removedState.Dispose();
            return true;
        }

        return false;
    }

    public void Dispatch(int count, Action<int> action)
    {
        if (count <= 0) return;

        if (!IsInitialized)
        {
            // Safe fallback: run inline if Initialize was never called.
            for (int i = 0; i < count; i++)
                action(i);
            return;
        }

        currentAction = action;
        workCount = count;
        Volatile.Write(ref nextIndex, 0);

        completionEvent.Reset(workers.Length);
        startSignal.Release(workers.Length);
        completionEvent.Wait();

        currentAction = null;
    }

    public void Shutdown()
    {
        if (!IsInitialized) return;

        running = false;
        startSignal.Release(workers.Length); // wake all workers so they can see running=false and exit

        foreach (Thread t in workers)
            t.Join();

        for (int i = 0; i < asyncWorkers.Length; i++)
        {
            asyncSignal.Release();
        }

        foreach (Thread t in asyncWorkers)
        {
            t.Join();
        }

        completionEvent.Dispose();
        startSignal.Dispose();
        asyncSignal.Dispose();

        foreach (KeyValuePair<int, AsyncJobState> pair in asyncJobStates)
        {
            pair.Value.Dispose();
        }

        asyncJobStates.Clear();

        asyncWorkers = null;
        asyncQueue = null;
        asyncJobStates = null;
        workers = null;
    }

    public void Dispose() => Shutdown();

    private void WorkerLoop()
    {
        while (true)
        {
            startSignal.Wait();
            if (!running) return;

            Action<int> action = currentAction;
            int total = workCount;

            while (true)
            {
                int index = Interlocked.Increment(ref nextIndex) - 1;
                if (index >= total) break;
                action(index);
            }

            completionEvent.Signal();
        }
    }

    private void AsyncWorkerLoop()
    {
        while (true)
        {
            asyncSignal.Wait();

            if (!running)
            {
                return;
            }

            if (!asyncQueue.TryDequeue(out int jobId))
            {
                continue;
            }

            if (!asyncJobStates.TryGetValue(jobId, out AsyncJobState state))
            {
                continue;
            }

            if (state.CancellationRequested)
            {
                state.Status = JobExecutionStatus.Canceled;
                state.Completion.Set();
                continue;
            }

            state.Status = JobExecutionStatus.Running;
            try
            {
                state.Action();
                state.Status = state.CancellationRequested
                    ? JobExecutionStatus.Canceled
                    : JobExecutionStatus.Completed;
            }
            catch (Exception ex)
            {
                state.CapturedException = ExceptionDispatchInfo.Capture(ex);
                state.Status = JobExecutionStatus.Faulted;
            }
            finally
            {
                state.Completion.Set();
            }
        }
    }
}