Optimize deterministic scheduler performance

Clockwork.slnx

@@ -1,4 +1,7 @@
 <Solution>
+  <Folder Name="/benchmarks/">
+    <Project Path="benchmarks/Clockwork.Benchmarks/Clockwork.Benchmarks.csproj" />
+  </Folder>
   <Folder Name="/src/">
     <Project Path="src/Clockwork/Clockwork.csproj" />
     <Project Path="src/Clockwork.Instrumentation/Clockwork.Instrumentation.csproj" />

Directory.Packages.props

@@ -3,11 +3,12 @@
     <ManagePackageVersionsCentrally>true</ManagePackageVersionsCentrally>
   </PropertyGroup>
   <ItemGroup>
+    <PackageVersion Include="BenchmarkDotNet" Version="0.15.8" />
     <PackageVersion Include="Microsoft.Build.Framework" Version="17.14.28" />
     <PackageVersion Include="Microsoft.Build.Utilities.Core" Version="17.14.28" />
     <PackageVersion Include="Microsoft.CodeAnalysis.CSharp" Version="4.14.0" />
     <PackageVersion Include="Microsoft.Extensions.Logging.Abstractions" Version="10.0.1" />
     <PackageVersion Include="Mono.Cecil" Version="0.11.6" />
     <PackageVersion Include="xunit.v3.mtp-v2" Version="3.2.1" />
   </ItemGroup>
-</Project>
+</Project>

benchmarks/Clockwork.Benchmarks/Clockwork.Benchmarks.csproj

@@ -0,0 +1,18 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\Clockwork\Clockwork.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <PackageReference Include="BenchmarkDotNet" />
+  </ItemGroup>
+
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <TargetFramework>net10.0</TargetFramework>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <Nullable>enable</Nullable>
+  </PropertyGroup>
+
+</Project>

benchmarks/Clockwork.Benchmarks/DeterministicSchedulerBenchmarks.cs

@@ -0,0 +1,79 @@
+using BenchmarkDotNet.Attributes;
+using Clockwork.Runtime.Execution;
+using Clockwork.Runtime.Scheduling;
+
+namespace Clockwork.Benchmarks;
+
+[MemoryDiagnoser]
+public class DeterministicSchedulerBenchmarks
+{
+    private const int OperationCount = 4;
+    private const int StepsPerOperation = 32;
+    private const int SchedulingPointCount = OperationCount * StepsPerOperation;
+    private int _initialCompletedSteps = 1;
+
+    [Benchmark(Baseline = true, OperationsPerInvoke = SchedulingPointCount)]
+    public int Direct()
+    {
+        var workload = new Workload(scheduler: null, _initialCompletedSteps);
+        for (var operation = 0; operation < OperationCount; operation++)
+        {
+            workload.Run();
+        }
+
+        return workload.CompletedSteps;
+    }
+
+    [Benchmark(OperationsPerInvoke = SchedulingPointCount)]
+    public int DeterministicScheduler() => RunScheduler(_initialCompletedSteps);
+
+    public static int RunTrace(int iterationCount)
+    {
+        var completed = 0;
+        for (var iteration = 0; iteration < iterationCount; iteration++)
+        {
+            completed += RunScheduler(initialCompletedSteps: 0);
+        }
+
+        return completed;
+    }
+
+    private static int RunScheduler(int initialCompletedSteps)
+    {
+        using var scheduler = new SimulationScheduler(
+            new SimulationRuntimeIdentity(Guid.Empty, Seed: 1, Description: "benchmark"));
+        var workload = new Workload(scheduler, initialCompletedSteps);
+        Action body = workload.Run;
+
+        for (var operation = 0; operation < OperationCount; operation++)
+        {
+            scheduler.Schedule("benchmark", body);
+        }
+
+        int dispatched = scheduler.Drain();
+        if (dispatched != SchedulingPointCount)
+        {
+            throw new InvalidOperationException(
+                $"Expected {SchedulingPointCount} dispatches but observed {dispatched}.");
+        }
+
+        return workload.CompletedSteps;
+    }
+
+    private sealed class Workload(SimulationScheduler? scheduler, int initialCompletedSteps)
+    {
+        public int CompletedSteps { get; private set; } = initialCompletedSteps;
+
+        public void Run()
+        {
+            for (var step = 0; step < StepsPerOperation; step++)
+            {
+                CompletedSteps++;
+                if (step + 1 < StepsPerOperation)
+                {
+                    scheduler?.Yield();
+                }
+            }
+        }
+    }
+}

benchmarks/Clockwork.Benchmarks/Program.cs

@@ -0,0 +1,12 @@
+using BenchmarkDotNet.Running;
+using Clockwork.Benchmarks;
+
+if (args is ["--trace", var iterationCount]
+    && int.TryParse(iterationCount, out var iterations)
+    && iterations > 0)
+{
+    Console.WriteLine(DeterministicSchedulerBenchmarks.RunTrace(iterations));
+    return;
+}
+
+BenchmarkSwitcher.FromAssembly(typeof(Program).Assembly).Run(args);

src/Clockwork/Runtime/Scheduling/SimulationScheduler.cs

@@ -57,10 +57,16 @@ public sealed class SimulationScheduler : IDisposable
     [ThreadStatic]
     private static SimulationOperation? t_currentOperation;
 
+    private static readonly ContextCallback s_runBodyInCapturedContext = static state =>
+    {
+        var operation = (SimulationOperation)state!;
+        operation.Scheduler.RunBodyScoped(operation);
+    };
+
     private readonly object _gate = new();
     private readonly object _transitionPublicationGate = new();
-    private readonly SortedDictionary<SimulationOperationId, SimulationOperation> _operations = new();
-    private readonly SortedDictionary<SimulationOperationId, SimulationOperation> _activeOperations = new();
+    private readonly List<SimulationOperation> _operations = [];
+    private readonly List<SimulationOperation> _activeOperations = [];
     private readonly SortedDictionary<SimulationResourceId, SimulationResource> _resources = new();
     private readonly HashSet<string> _suspendedNodes = new(StringComparer.Ordinal);
     private readonly SemaphoreSlim _handback = new(0, 1);
@@ -486,8 +492,8 @@ private SimulationOperation RegisterCore(
                     body,
                     capturedContext,
                     priority);
-                _operations.Add(id, operation);
-                _activeOperations.Add(id, operation);
+                _operations.Add(operation);
+                _activeOperations.Add(operation);
                 _raceTracker.RegisterOperation(operation, parent);
             }
 
@@ -1201,7 +1207,7 @@ internal bool HasPendingWork(SimulationNodeIdentity node)
         ArgumentNullException.ThrowIfNull(node);
         lock (_gate)
         {
-            foreach (var operation in _activeOperations.Values)
+            foreach (var operation in _activeOperations)
             {
                 if (operation.Node == node)
                 {
@@ -1234,7 +1240,7 @@ internal void RemovePendingWork(SimulationNodeIdentity node)
                 _readinessWaits.RemoveAt(index);
             }
 
-            operations = _activeOperations.Values
+            operations = _activeOperations
                 .Where(operation => operation.Node == node)
                 .ToArray();
         }
@@ -1305,7 +1311,7 @@ private int ElapseTimers(IReadOnlyList<ISimulationTimerEntry> due)
 
     private bool HasRunnableUnderLock()
     {
-        foreach (var operation in _activeOperations.Values)
+        foreach (var operation in _activeOperations)
         {
             if (operation.State == SimulationOperationState.Runnable && IsEligibleUnderLock(operation))
             {
@@ -1382,7 +1388,7 @@ public void Cancel(SimulationOperation operation)
                 operation.RequestTermination();
                 registration = DetachWaiterUnderLock(operation);
                 operation.ApplyTransition(SimulationOperationState.Canceled);
-                _activeOperations.Remove(operation.Id);
+                _activeOperations.Remove(operation);
                 needsUnwind = operation.Thread is not null;
             }
 
@@ -1904,7 +1910,7 @@ public IReadOnlyList<SimulationOperationStatus> CaptureStatus()
         lock (_gate)
         {
             var result = new List<SimulationOperationStatus>(_operations.Count);
-            foreach (var operation in _operations.Values)
+            foreach (var operation in _operations)
             {
                 result.Add(new SimulationOperationStatus(
                     operation.Id,
@@ -2001,7 +2007,7 @@ public SimulationDeadlockReport DetectDeadlock()
             var runnable = 0;
             var blocked = 0;
             var nonTerminal = 0;
-            foreach (var operation in _activeOperations.Values)
+            foreach (var operation in _activeOperations)
             {
                 nonTerminal++;
 
@@ -2054,7 +2060,7 @@ public string DescribeLiveness()
             report = DetectDeadlock();
             now = _clock.Now;
             operations = new List<(SimulationOperationId, SimulationOperationState, string, SimulationPauseReason?)>(_operations.Count);
-            foreach (var operation in _operations.Values)
+            foreach (var operation in _operations)
             {
                 operations.Add((operation.Id, operation.State, operation.WorkDescription, operation.PauseReason));
             }
@@ -2090,8 +2096,8 @@ private List<SimulationWaitCycle> FindWaitForCyclesUnderLock(Dictionary<Simulati
 
         // The wait-for graph is functional here (each paused operation waits on exactly one resource,
         // hence has at most one successor), so following the single successor chain from each start
-        // node detects every cycle deterministically. _activeOperations is id-sorted, so starts are ordered.
-        foreach (var operation in _activeOperations.Values)
+        // node detects every cycle deterministically. _activeOperations is in id order, so starts are ordered.
+        foreach (var operation in _activeOperations)
         {
             var start = operation.Id;
             if (!edges.ContainsKey(start) || globallyVisited.Contains(start))
@@ -2156,7 +2162,7 @@ private SimulationWaitCycle BuildCycle(List<SimulationOperationId> cycleIds, Dic
             var waiter = edge.Waiter;
             entries.Add(new SimulationWaitCycleEntry(
                 id,
-                _operations[id].WorkDescription,
+                GetOperation(id).WorkDescription,
                 waiter.Resource.Id,
                 waiter.Resource.Name,
                 edge.OwnerId,
@@ -2219,7 +2225,7 @@ public void Dispose()
                 _disposed = true;
                 victims = new List<SimulationOperation>();
                 registrations = new List<CancellationTokenRegistration>();
-                foreach (var operation in _activeOperations.Values)
+                foreach (var operation in _activeOperations)
                 {
                     operation.RequestTermination();
 
@@ -2267,12 +2273,16 @@ public void Dispose()
 
     private SimulationOperation? SelectRunnable()
     {
-        // Collect the runnable operations in ascending id order (_activeOperations is a SortedDictionary),
-        // then delegate the choice to the pluggable strategy. The default round-robin strategy
-        // reproduces the controlled-operation scheduler behavior exactly.
+        if (_strategy is RoundRobinSchedulingStrategy && _decisionLog is null && _replayValidator is null)
+        {
+            return SelectRoundRobinRunnable();
+        }
+
+        // Custom and instrumented strategies receive a stable snapshot which they may retain.
         List<SimulationOperation>? runnable = null;
-        foreach (var operation in _activeOperations.Values)
+        for (var index = 0; index < _activeOperations.Count; index++)
         {
+            SimulationOperation operation = _activeOperations[index];
             if (operation.State != SimulationOperationState.Runnable || !IsEligibleUnderLock(operation))
             {
                 continue;
@@ -2303,6 +2313,33 @@ public void Dispose()
         return chosen;
     }
 
+    private SimulationOperation? SelectRoundRobinRunnable()
+    {
+        SimulationOperation? wrapTarget = null;
+        for (var index = 0; index < _activeOperations.Count; index++)
+        {
+            SimulationOperation operation = _activeOperations[index];
+            if (operation.State != SimulationOperationState.Runnable || !IsEligibleUnderLock(operation))
+            {
+                continue;
+            }
+
+            wrapTarget ??= operation;
+            if (operation.Id > _lastSelected)
+            {
+                _lastSelected = operation.Id;
+                return operation;
+            }
+        }
+
+        if (wrapTarget is not null)
+        {
+            _lastSelected = wrapTarget.Id;
+        }
+
+        return wrapTarget;
+    }
+
     private bool IsEligibleUnderLock(SimulationOperation operation) =>
         operation.Node is not { } node || !_suspendedNodes.Contains(node.Address);
 
@@ -2567,7 +2604,7 @@ private void RunBody(SimulationOperation operation)
             var context = operation.CapturedContext;
             if (context is not null)
             {
-                ExecutionContext.Run(context, static s => ((BodyClosure)s!).Run(), new BodyClosure(this, operation));
+                ExecutionContext.Run(context, s_runBodyInCapturedContext, operation);
             }
             else
             {
@@ -2639,7 +2676,7 @@ private void HandOffTerminal(SimulationOperation operation, SimulationOperationS
                 }
 
                 operation.ApplyTransition(terminal, terminalException: terminalException);
-                _activeOperations.Remove(operation.Id);
+                _activeOperations.Remove(operation);
                 _pendingTerminalNotification = operation;
             }
         }
@@ -2725,17 +2762,9 @@ private void WaitForPendingTerminalNotificationUnderLock()
         }
     }
 
-    private sealed class TransitionPublicationScope(SimulationScheduler scheduler) : IDisposable
+    private readonly ref struct TransitionPublicationScope(SimulationScheduler scheduler)
     {
-        private int _disposed;
-
-        public void Dispose()
-        {
-            if (Interlocked.Exchange(ref _disposed, 1) == 0)
-            {
-                scheduler.ExitTransitionPublicationScope();
-            }
-        }
+        public void Dispose() => scheduler.ExitTransitionPublicationScope();
     }
 
     private static void UnwindParkedThread(SimulationOperation operation)
@@ -2778,12 +2807,23 @@ private void ValidateOwnership(SimulationOperation operation)
         }
     }
 
+    private SimulationOperation GetOperation(SimulationOperationId id)
+    {
+        var index = checked((int)(id.Value - 1));
+        if ((uint)index >= (uint)_operations.Count || _operations[index].Id != id)
+        {
+            throw new SimulationSchedulerException($"Unknown controlled operation id '{id}'.");
+        }
+
+        return _operations[index];
+    }
+
     private SimulationOperation[] SnapshotOperations()
     {
         lock (_gate)
         {
             var array = new SimulationOperation[_operations.Count];
-            _operations.Values.CopyTo(array, 0);
+            _operations.CopyTo(array, 0);
             return array;
         }
     }
@@ -2793,7 +2833,7 @@ private int CountActiveOperations(Func<SimulationOperation, bool> predicate)
         lock (_gate)
         {
             var count = 0;
-            foreach (var operation in _activeOperations.Values)
+            foreach (var operation in _activeOperations)
             {
                 if (predicate(operation))
                 {
@@ -2810,11 +2850,6 @@ private void Notify(SimulationOperation operation, SimulationOperationState stat
 
     private void ThrowIfDisposed() => ObjectDisposedException.ThrowIf(_disposed, this);
 
-    private sealed class BodyClosure(SimulationScheduler scheduler, SimulationOperation operation)
-    {
-        public void Run() => scheduler.RunBodyScoped(operation);
-    }
-
     private sealed class ReadinessWait : ISimulationWorkRegistration
     {
         private readonly SimulationScheduler _scheduler;