Update to support source jobs

Author: emelrad12

JobScheduler.Benchmarks/Benchmark.cs

@@ -210,7 +210,7 @@ private static long BenchVector(bool useVector)
         var parentJob = ParallelForJobCommon.GlobalScheduler.Schedule();
         for (var sindex = 0; sindex < loopCount; sindex++)
         {
-            var job = new ParallelJobProducer<VectorCalculationJob>(0, jobCount, data, 16, !useVector);
+            var job = new ParallelJobProducer<VectorCalculationJob>(0, jobCount, data, loopSize: 16, onlySingle: !useVector);
             job.CheckAndSplit();
             job.GetHandle().SetParent(parentJob);
             ParallelForJobCommon.GlobalScheduler.Flush(job.GetHandle());

JobScheduler.Test/JobSchedulerTests.cs

@@ -216,4 +216,14 @@ public void ManyMoreDependencyThanNodeLimitSize()
             jobScheduler.Wait(targetJob);
         });
     }
+
+    [Test]
+    public void TryToEditAJobThatIsReadyToExecute()
+    {
+        using var jobScheduler = new JobScheduler();
+        var newJob = jobScheduler.Schedule();
+        newJob.UnfinishedJobs--;
+        Throws<InvalidOperationException>(() => newJob.SetParent(new()));
+        Throws<InvalidOperationException>(() => newJob.SetDependsOn(new()));
+    }
 }

JobScheduler.Test/ParallelForTests.cs

@@ -102,6 +102,35 @@ public void NestedJobs()
                 Interlocked.Increment(ref triggers);
             });
         });
+        Assert.That(triggers, Is.EqualTo(100));
+    }
+
+    [Test]
+    public void UsingSourceJobButNotStarted()
+    {
+        var triggers = 0;
+        var sourceJob = ParallelForJobCommon.GlobalScheduler.Schedule();
+        var handle = ParallelForJob.Create(0, 100, i =>
+        {
+            Interlocked.Increment(ref triggers);
+        }, source: sourceJob);
+        ParallelForJobCommon.GlobalScheduler.Flush(handle);
+        Thread.Sleep(10);
         Assert.That(triggers, Is.EqualTo(0));
     }
+
+    [Test]
+    public void UsingSourceJobAndStarted()
+    {
+        var triggers = 0;
+        var sourceJob = ParallelForJobCommon.GlobalScheduler.Schedule();
+        var handle = ParallelForJob.Create(0, 100, i =>
+        {
+            Interlocked.Increment(ref triggers);
+        }, source: sourceJob);
+        ParallelForJobCommon.GlobalScheduler.Flush(handle);
+        ParallelForJobCommon.GlobalScheduler.Flush(sourceJob);
+        Thread.Sleep(10);
+        Assert.That(triggers, Is.EqualTo(100));
+    }
 }

JobScheduler/JobHandle.cs

@@ -155,6 +155,7 @@ public struct JobHandle
     internal ushort Generation;
 
     private readonly bool _isInitialized;
+
     /// <summary>
     /// Indicates whether this <see cref="JobHandle"/> is null.
     /// Default or Just JobHandle() is null, meaning it has not been initialized yet.
@@ -215,17 +216,39 @@ public ref int UnfinishedJobs
         get => ref Pool.UnfinishedJobs[Index];
     }
 
+    /// <summary>
+    /// Sets the job handle to be ready to execute.
+    /// The job still needs to be added to the worker queue to be executed.
+    /// It just means that before this point the job cannot be executed.
+    /// </summary>
+    // The readonly is not really true, but it silences the warning.
+    public readonly int SetReadyToExecute()
+    {
+        var result = Interlocked.Decrement(ref UnfinishedJobs);
+        if (result == 0)
+        {
+            throw new("JobHandle has no unfinished jobs, cannot set it ready to execute.");
+        }
+
+        return result;
+    }
+
     /// <summary>
     /// Sets a <see cref="JobHandle"/> that is to be the parent.
     /// If a null handle is passed it will be ignored.
     /// </summary>
     /// <param name="parent">The parent job handle</param>
     public void SetParent(JobHandle parent)
     {
+        if (IsJobEligibleForExecution())
+        {
+            throw new InvalidOperationException("Tried to set a parent on a job that is already eligible for execution. This is not allowed as it can lead to deadlocks or unexpected behavior.");
+        }
         if (parent.IsNull)
         {
-            return;
+            throw new ArgumentNullException(nameof(parent), "Cannot set a null parent on a job handle.");
         }
+
         Interlocked.Increment(ref parent.UnfinishedJobs);
         Parent = parent.Index;
     }
@@ -273,6 +296,10 @@ public bool IsHandleAlreadyRecycled()
     /// <param name="source">The handle that triggers the current handle</param>
     public void SetDependsOn(JobHandle source)
     {
+        if (IsJobEligibleForExecution())
+        {
+            throw new InvalidOperationException("Tried to set a dependency on a job that is already eligible for execution. This is not allowed as it can lead to deadlocks or unexpected behavior.");
+        }
         source.AddDependent(this);
     }
 
@@ -315,6 +342,10 @@ public void FlushAndWait()
     /// <param name="target">What to be added to be executed after</param>
     private void AddDependent(JobHandle target)
     {
+        if (IsJobEligibleForExecution())
+        {
+            throw new("Tried to add a dependent to a job that is already eligible for execution. This is not allowed as it can lead to deadlocks or unexpected behavior.");
+        }
         Dependents ??= Pool.DependenciesBackingArray[Index];
         if (!IsAnySpaceAvailable())
         {
@@ -384,4 +415,9 @@ private int GetFreeDependentsIndex()
 
         return -1;
     }
+
+    public bool IsJobEligibleForExecution()
+    {
+        return UnfinishedJobs == 1;
+    }
 }

JobScheduler/JobScheduler.cs

@@ -111,13 +111,6 @@ public float CalculateThreadUsage()
         return total;
     }
 
-    private void FlushDirect(JobHandle job)
-    {
-        while (!Worker.Enqueue(job))
-        {
-        }
-    }
-
     /// <summary>
     /// Tries to transfer <see cref="JobHandle"/> to the <see cref="Workers"/> so that it can be executed.
     /// If it has unfinished children it will not be flushed, and instead will be left for the children to flush.
@@ -126,18 +119,25 @@ private void FlushDirect(JobHandle job)
     public void Flush(JobHandle job)
     {
         // This is to prevent the job from being flushed multiple times.
-        var unfinishedJobs = Interlocked.Decrement(ref job.UnfinishedJobs);
-        if (unfinishedJobs == 0)
+        var unfinishedJobs = job.SetReadyToExecute();
+
+        if (unfinishedJobs != 1)
         {
-            throw new("JobHandle has no unfinished jobs, cannot flush it.");
+            return;
         }
 
-        if (unfinishedJobs != 1)
+        // It is pointless to flush a job that has no work to do.
+        // This is a job that triggers other jobs,
+        // so might as well do it now, to lower latency.
+        if (job.Job == null)
         {
+            Finish(job);
             return;
         }
 
-        FlushDirect(job);
+        while (!Worker.Enqueue(job))
+        {
+        }
     }
 
     /// <summary>

JobScheduler/Utils/IParallelJobProducer.cs

@@ -41,6 +41,7 @@ public struct ParallelJobProducer<T> : IJob where T : struct, IParallelJobProduc
     private readonly int _to;
     private readonly T _producer;
     private readonly JobHandle _selfHandle;
+    private readonly JobHandle _sourceHandle;
     private readonly bool _onlySingle;
     private readonly int _loopSize;
 
@@ -50,28 +51,38 @@ public struct ParallelJobProducer<T> : IJob where T : struct, IParallelJobProduc
     /// <param name="from">Where the loop starts</param>
     /// <param name="to">Maximum to loop to</param>
     /// <param name="producer">The job to call</param>
-    /// <param name="scheduler">The scheduler where the jobs should be put</param>
+    /// <param name="source"></param>
     /// <param name="loopSize">Size of the loop, useful for when you want to use vectorization</param>
     /// <param name="onlySingle">Makes sure you can ignore the end parameter in the IParallelJobProducer, and use start as index</param>
-    public ParallelJobProducer(int from, int to, T producer, int loopSize = 16, bool onlySingle = false)
+    /// <param name="scheduler">The scheduler where the jobs should be put</param>
+    public ParallelJobProducer(int from, int to, T producer, JobHandle source = default, int loopSize = 16, bool onlySingle = false)
     {
+        _sourceHandle = source;
         _from = from;
         _to = to;
         _producer = producer;
         _onlySingle = onlySingle;
         _loopSize = loopSize;
         _selfHandle = ParallelForJobCommon.GlobalScheduler!.Schedule(this);
+        if (!source.IsNull)
+        {
+            _selfHandle.SetDependsOn(source);
+        }
     }
 
     //Only used to spawn sub-jobs
-    private ParallelJobProducer(T producer, JobHandle parent, int start, int end, int loopSize, bool onlySingle)
+    private ParallelJobProducer(T producer, JobHandle parent, int start, int end, int loopSize, bool onlySingle, JobHandle source)
     {
         _producer = producer;
         _from = start;
         _to = end;
         _loopSize = loopSize;
         _onlySingle = onlySingle;
         _selfHandle = ParallelForJobCommon.GlobalScheduler.Schedule(this, parent);
+        if (!source.IsNull)
+        {
+            _selfHandle.SetDependsOn(source);
+        }
     }
 
     /// <summary>
@@ -108,7 +119,9 @@ public void Execute()
     /// <returns>The amount.</returns>
     private int CalculateChildrenToSplitInto()
     {
-        const int ChildrenToSplitInto = 128; //This should be equal to the number of threads(or that times 2/3?) but for now it's just a constant
+        // If you change this also change the bulk queue segment size.
+        // This should be equal to the number of threads(or that times 2/3?) but for now it's just a constant
+        const int ChildrenToSplitInto = 128;
         var range = _to - _from;
         return range < ChildrenToSplitInto ? range : ChildrenToSplitInto;
     }
@@ -134,11 +147,20 @@ public bool CheckAndSplit()
     /// <exception cref="Exception">An exception occuring when an invalid range was passed.</exception>
     private void Split()
     {
+        if (_selfHandle.IsJobEligibleForExecution())
+        {
+            throw new InvalidOperationException("Tried to split a job that is already eligible for execution. You should split before flushing the job.");
+        }
+
         var childrenToSplitInto = CalculateChildrenToSplitInto();
-        BulkQueue<JobHandle>.Segment segment;
-        while (!Worker.GetEmptySegment(out segment))
+        BulkQueue<JobHandle>.Segment segment = default;
+        if (_sourceHandle.IsNull)
         {
+            while (!Worker.GetEmptySegment(out segment))
+            {
+            }
         }
+
         _selfHandle.Job = null;
         for (var i = 0; i < childrenToSplitInto; i++)
         {
@@ -149,15 +171,23 @@ private void Split()
                 throw new($"Invalid range from {start} to {end}");
             }
 
-            var pjob = new ParallelJobProducer<T>(_producer, _selfHandle, start, end, _loopSize, _onlySingle);
-            if (!segment.Enqueue(pjob._selfHandle))
+            var pjob = new ParallelJobProducer<T>(_producer, _selfHandle, start, end, _loopSize, _onlySingle, _sourceHandle);
+            if (_sourceHandle.IsNull)
             {
-                throw new InvalidOperationException("Failed to enqueue job in segment. Probably segment size was changed without changing the code using it.");
+                if (!segment.Enqueue(pjob._selfHandle))
+                {
+                    throw new InvalidOperationException("Failed to enqueue job in segment. Probably segment size was changed without changing the code using it.");
+                }
             }
+
             //Simulate flushing but without using the scheduler
-            pjob._selfHandle.UnfinishedJobs--;
+            pjob._selfHandle.SetReadyToExecute();
+        }
+
+        if (_sourceHandle.IsNull)
+        {
+            Worker.EnqueueFullSegment(segment);
         }
-        Worker.EnqueueFullSegment(segment);
     }
 
     /// <summary>

JobScheduler/Utils/LambdaJob.cs

@@ -20,10 +20,15 @@ public void Execute()
     /// <param name="action">The function to run</param>
     /// <param name="parent">The parent handle</param>
     /// <returns></returns>
-    public static JobHandle Create(Action action, JobHandle parent)
+    public static JobHandle Create(Action action, JobHandle source = default)
     {
         var job = new JobProducer(action);
-        var handle = ParallelForJobCommon.GlobalScheduler.Schedule(job, parent);
+        var handle = ParallelForJobCommon.GlobalScheduler!.Schedule(job);
+        if (!source.IsNull)
+        {
+            handle.SetDependsOn(source);
+        }
+
         return handle;
     }
 
@@ -32,7 +37,7 @@ public static JobHandle Create(Action action, JobHandle parent)
     /// </summary>
     public static void CreateFlushAndWait(Action action)
     {
-        var handle = Create(action, JobHandle.Null);
+        var handle = Create(action);
         handle.FlushAndWait();
     }
 }

JobScheduler/Utils/ParallelForEachJob.cs

@@ -22,19 +22,23 @@ public void RunSingle(int index)
     /// <summary>
     /// Creates a ParallelJobProducer that executes the provided action on every item of the given collection.
     /// </summary>
-    /// <param name="source">The data to run the action on</param>
+    /// <param name="data">The data to run the scheduler over</param>
     /// <param name="action">What to run</param>
-    /// <param name="scheduler">The scheduler to use</param>
+    /// <param name="source">The job that will trigger this one</param>
     /// <param name="parent"></param>
     /// <returns>A struct that wraps the handle</returns>
-    public static JobHandle Create<T>(IEnumerable<T> source, Action<T> action, JobHandle parent = default)
+    public static JobHandle Create<T>(IEnumerable<T> data, Action<T> action, JobHandle parent = default, JobHandle source = default)
     {
-        var partitioner = Partitioner.Create(source);
+        var partitioner = Partitioner.Create(data);
         // We use Environment.ProcessorCount * 4 to account for uneven distribution of work and to ensure that we have enough partitions to keep all cores busy.
         // The overhead of jobs is very low hence we can afford to have more partitions than cores.
         var partitions = partitioner.GetPartitions(Environment.ProcessorCount * 4);
-        var producer = new ParallelJobProducer<PartitionedJobProducer<T>>(0, partitions.Count, new(partitions, action));
-        producer.GetHandle().SetParent(parent);
+        var producer = new ParallelJobProducer<PartitionedJobProducer<T>>(0, partitions.Count, new(partitions, action), source);
+        if (!parent.IsNull)
+        {
+            producer.GetHandle().SetParent(parent);
+        }
+
         producer.CheckAndSplit();
         return producer.GetHandle();
     }

JobScheduler/Utils/ParallelForJob.cs

@@ -19,12 +19,16 @@ public void RunSingle(int index)
     /// <param name="start">Where to start executing</param>
     /// <param name="count">Executes to this</param>
     /// <param name="action">What to run</param>
-    /// <param name="scheduler">The scheduler to use</param>
+    /// <param name="target">The job to trigger after finishing</param>
+    /// <param name="source">The job that will trigger this one</param>
     /// <returns>A struct that wraps the handle</returns>
-    public static JobHandle Create(int start, int count, Action<int> action, JobHandle parent = default)
+    public static JobHandle Create(int start, int count, Action<int> action, JobHandle target = default, JobHandle source = default)
     {
-        var producer = new ParallelJobProducer<JobProducer>(start, count, new(action));
-        producer.GetHandle().SetParent(parent);
+        var producer = new ParallelJobProducer<JobProducer>(start, count, new(action), source);
+        if (!target.IsNull)
+        {
+            producer.GetHandle().SetParent(target);
+        }
         producer.CheckAndSplit();
         return producer.GetHandle();
     }

JobScheduler/Worker.cs

@@ -21,6 +21,8 @@ internal class Worker
     private readonly WorkStealingDeque<JobHandle> _queue;
     private readonly JobScheduler _jobScheduler;
     private volatile CancellationTokenSource _cancellationToken;
+    // 128 for now to match the IParallelJobProducer size.
+    // if you change this then change also the split IParallelJobProducer size, to not be larger than this value.
     private static readonly BulkQueue<JobHandle> _bulkQueue = new(1000, 128);
     public bool IsCurrentlyWorking { get; private set; } = false;