Merge pull request #8 from lazydynamics/6-threading-scheduler

Add threaded scheduling to `TickedScheduler`

Author: wouterwln

Project.toml

@@ -1,15 +1,17 @@
 name = "EnvironmentEngine"
 uuid = "12b47476-e81a-4b1e-bd5d-7199c2c66cc7"
 authors = ["wouterwln"]
-version = "0.0.1-DEV"
+version = "0.0.1"
 
 [deps]
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
 Aqua = "0.8"
 InteractiveUtils = "1.10.0"
 JET = "0.9, 0.10"
+Static = "1.2.0"
 Test = "1.10"
 TestItemRunner = "1.1"
 Unitful = "1.24.0"

docs/make.jl

@@ -5,7 +5,7 @@ DocMeta.setdocmeta!(EnvironmentEngine, :DocTestSetup, :(using EnvironmentEngine)
 
 makedocs(;
     modules = [EnvironmentEngine],
-    authors = "lazydynamics and contributors",
+    authors = "Lazy Dynamics and contributors",
     sitename = "EnvironmentEngine.jl",
     format = Documenter.HTML(;
         canonical = "https://lazydynamics.github.io/EnvironmentEngine.jl", edit_link = "main", assets = String[]

src/implementations/schedulers/ticked.jl

@@ -1,4 +1,6 @@
 using Unitful
+using Base.Threads
+using Static
 
 export TickedScheduler
 
@@ -12,20 +14,26 @@ A scheduler that processes jobs at regular tick intervals.
 - `jobs::Vector{Job}`: The jobs to be scheduled
 - `ticker::Ticker{T}`: The ticker that manages timing
 """
-struct TickedScheduler{C <: Clock, T, V} <: Scheduler
+struct TickedScheduler{C <: Clock, T, V, B} <: Scheduler
     clock::C
     jobs::V
     ticker::Ticker{T}
+    threading::B
 end
 
 """
     TickedScheduler(clock::Clock, tick_period::T) where {T}
 
 Create a new TickedScheduler with the specified clock and tick period.
 """
-function TickedScheduler(clock::Clock, tick_period::T) where {T}
+function TickedScheduler(clock::Clock, tick_period::T; threading::Bool = false) where {T}
     tick_period = convert(Quantity{Float64}, tick_period)
-    TickedScheduler(clock, Job[], Ticker(tick_period, convert(typeof(tick_period), now(clock)), zero(tick_period)))
+    TickedScheduler(
+        clock,
+        Job[],
+        Ticker(tick_period, convert(typeof(tick_period), now(clock)), zero(tick_period)),
+        static(threading)
+    )
 end
 
 """
@@ -47,17 +55,28 @@ function update!(scheduler::TickedScheduler)
     advance_to!(scheduler.ticker, current_time)
 
     while can_tick(scheduler.ticker)
-        foreach(scheduler.jobs) do job
-            progress!(job, scheduler.ticker.period)
-        end
+        _process_jobs!(scheduler)
         consume_tick!(scheduler.ticker)
     end
 end
 
+function _process_jobs!(scheduler::TickedScheduler{C, T, V, <:False}) where {C, T, V}
+    foreach(scheduler.jobs) do job
+        progress!(job, scheduler.ticker.period)
+    end
+end
+
+function _process_jobs!(scheduler::TickedScheduler{C, T, V, <:True}) where {C, T, V}
+    tasks = map(scheduler.jobs) do job
+        @spawn progress!(job, scheduler.ticker.period)
+    end
+    wait.(tasks)
+end
+
 function get_period(scheduler::TickedScheduler)
     return scheduler.ticker.period
 end
 
 function compile(scheduler::TickedScheduler)
-    return TickedScheduler(scheduler.clock, (scheduler.jobs...,), scheduler.ticker)
+    return TickedScheduler(scheduler.clock, (scheduler.jobs...,), scheduler.ticker, scheduler.threading)
 end

test/implementations/scheduler/tickedscheduler_tests.jl

@@ -55,3 +55,81 @@
         @test order == [1, 2]
     end
 end
+
+@testitem "TickedScheduler with threading enabled" setup = [MockJobs, MockClocks] begin
+    import EnvironmentEngine: TickedScheduler, update!, set_time!, advance_time!, schedule!
+    using Unitful
+    using Base.Threads
+    if Base.Threads.nthreads() > 1
+        @testset "Threading preserves execution count" begin
+            clock = MockClocks.MockClock()
+
+            # Test both modes
+            scheduler_seq = TickedScheduler(clock, 1.0ms; threading = false)
+            scheduler_par = TickedScheduler(clock, 1.0ms; threading = true)
+
+            # Create counting jobs
+            jobs_seq = [MockJobs.CountingJob() for _ in 1:3]
+            jobs_par = [MockJobs.CountingJob() for _ in 1:3]
+
+            # Schedule jobs
+            for job in jobs_seq
+                schedule!(scheduler_seq, job)
+            end
+            for job in jobs_par
+                schedule!(scheduler_par, job)
+            end
+
+            # Advance time and update
+            advance_time!(clock, 10.0u"ms")
+            update!(scheduler_seq)
+            update!(scheduler_par)
+
+            # Both should execute the same number of times
+            seq_total = sum(job.counter[] for job in jobs_seq)
+            par_total = sum(job.counter[] for job in jobs_par)
+            @test seq_total == par_total
+            @test seq_total > 0
+        end
+
+        @testset "All jobs execute per tick" begin
+            clock = MockClocks.MockClock()
+            scheduler = TickedScheduler(clock, 1.0u"ms"; threading = true)
+
+            jobs = [MockJobs.CountingJob() for _ in 1:5]
+            for job in jobs
+                schedule!(scheduler, job)
+            end
+
+            advance_time!(clock, 1.0u"ms")
+            update!(scheduler)
+
+            # All jobs should have executed exactly once
+            for job in jobs
+                @test job.counter[] == 1
+            end
+        end
+
+        @testset "Multiple ticks work correctly" begin
+            clock = MockClocks.MockClock()
+            scheduler = TickedScheduler(clock, 1.0u"ms"; threading = true)
+
+            jobs = [MockJobs.CountingJob() for _ in 1:3]
+            for job in jobs
+                schedule!(scheduler, job)
+            end
+
+            # Advance time for 3 ticks
+            advance_time!(clock, 3.0u"ms")
+            update!(scheduler)
+
+            # All jobs should have executed 3 times
+            for job in jobs
+                @test job.counter[] == 3
+            end
+        end
+    else
+        multiple_threads = false
+        @test_broken multiple_threads
+    end
+end

test/integration/integration_tests.jl

@@ -91,3 +91,20 @@
         @test all(obs == 1 for obs in test_state.observations_received)
     end
 end
+
+@testitem "Multithread TimedJob scheduling" begin
+    import EnvironmentEngine: TimedJob
+    using Base.Threads: Atomic, atomic_add!
+    clock = VirtualClock()
+    scheduler = TickedScheduler(clock, 5.0ms; threading = true)
+
+    counter = Atomic{Int}(0)
+    every(scheduler, 1ms) do dt
+        atomic_add!(counter, 1)
+    end
+
+    advance_time!(clock, 1.0s)
+    update!(scheduler)
+
+    @test counter[] == 1000
+end

test/mocks/job.jl

@@ -2,11 +2,24 @@
     import EnvironmentEngine: Job, progress!
     using Unitful
     import Unitful: 𝐓
+    using Base.Threads: Atomic
+    import Base.Threads: atomic_add!
+
     struct MockJob <: Job
         f
     end
 
     function progress!(job::MockJob, dt::Quantity{<:Number, 𝐓})
         job.f()
     end
+
+    struct CountingJob <: Job
+        counter::Atomic{Int}
+    end
+
+    CountingJob() = CountingJob(Atomic{Int}(0))
+
+    function progress!(job::CountingJob, dt::Quantity{<:Number, 𝐓})
+        atomic_add!(job.counter, 1)
+    end
 end