Share the rescue clock of steal blocks between participants

src/pool.jl

@@ -182,15 +182,18 @@ mutable struct ChunkCounter
 end
 ChunkCounter() = ChunkCounter(pad_before(), 0, pad_after())
 
-# Cache-line-separated packed range (see `pack_range`). UInt64 also on 32-bit systems:
-# Julia aligns atomic fields to their size. i686 uses SSE2 loads/stores and lock cmpxchg8b
-# (checked on Julia 1.10/1.12); ARMv7 should use ldrexd/strexd (not verified).
+# Cache-line-separated packed range (see `pack_range`) and its owner's rescue clock
+# (`owner_due`), on one line. UInt64 also on 32-bit systems: Julia aligns atomic fields to
+# their size. i686 uses SSE2 loads/stores and lock cmpxchg8b (checked on Julia 1.10/1.12);
+# ARMv7 should use ldrexd/strexd (not verified).
+const RANGE_PAD = CACHE_LINE ÷ sizeof(Int) - 2 * (sizeof(UInt64) ÷ sizeof(Int))
 mutable struct RangeBlock
-    const _pad0::PadBefore
+    const _pad0::NTuple{RANGE_PAD, Int}
     @atomic word::UInt64
+    @atomic seen::UInt64
     const _pad1::PadAfter
 end
-RangeBlock() = RangeBlock(pad_before(), UInt64(0), pad_after())
+RangeBlock() = RangeBlock(ntuple(_ -> 0, Val(RANGE_PAD)), UInt64(0), UInt64(0), pad_after())
 
 # Per-pool range state: error/stop share a read-only line until failure; `active` has
 # its own line, read by idle workers and set/cleared by the dispatcher. Steal blocks
@@ -203,7 +206,7 @@ mutable struct RangeState
     @atomic active::Bool
     const _pad2::PadAfter
     const blocks::Vector{RangeBlock}
-    # Dispatcher-only first-awake timestamps; zero means not observed awake.
+    # Static mode's dispatcher-only first-awake timestamps; zero means not observed awake.
     const seen::Vector{UInt64}
 end
 RangeState(nw::Int) = RangeState(

src/ranges.jl

@@ -50,6 +50,21 @@ const RANGE_PARK_S = 10.0e-6
 @inline rescue_due(asleep::Bool, seen::UInt64, now::UInt64, deadline::UInt64) =
     now >= deadline || (!asleep && seen != 0 && now - seen >= RANGE_AWAKE_NS)
 
+# The same rule for the owner of steal block `b` with its command pending, shared by
+# thieves and the dispatcher: `b.seen` is the first awake observation of the region,
+# never reset before `init_blocks!`, so waits do not renew; the owner must still be awake.
+# Another observer's `seen` may follow our `now`, which is then not due yet.
+@inline function owner_due(b::RangeBlock, asleep::Bool, now::UInt64, deadline::UInt64)
+    now >= deadline && return true
+    asleep && return false
+    seen = @atomic :monotonic b.seen
+    if seen == 0
+        old, ok = @atomicreplace :monotonic :monotonic b.seen UInt64(0) => now
+        seen = ok ? now : old
+    end
+    return seen <= now && now - seen >= RANGE_AWAKE_NS
+end
+
 # Wake all sleeping workers; their usual child wake-ups remain harmless.
 @inline function wake_all(pool::Pool, n::Int, skip::Int, base::Int)
     for chunk in 2:n
@@ -66,8 +81,11 @@ end
     )
     st = job.state
     seen = st.seen
-    for chunk in 2:n
-        @inbounds seen[worker_for_chunk(chunk, skip)] = 0
+    # Steal jobs use their blocks' clocks
+    if !(job isa StealRanges)
+        for chunk in 2:n
+            @inbounds seen[worker_for_chunk(chunk, skip)] = 0
+        end
     end
     spin = pool.oversubscribed ? RANGE_SPIN_OVERSUBSCRIBED_NS : RANGE_SPIN_NS
     tstart = UInt64(0)
@@ -87,9 +105,15 @@ end
             (@atomic :monotonic ws.cmd) == cmd || continue
             asleep = @atomic :monotonic (@inbounds pool.sleepflags[k]).sleeping
             now == 0 && (now = time_ns())
-            t = asleep ? UInt64(0) : (@inbounds(seen[k]) == 0 ? now : @inbounds(seen[k]))
-            @inbounds seen[k] = t
-            if rescue_due(asleep, t, now, job.deadline)
+            if job isa StealRanges
+                # Chunk `chunk` runs block `k`
+                due = owner_due(@inbounds(job.blocks[k]), asleep, now, job.deadline)
+            else
+                t = asleep ? UInt64(0) : (@inbounds(seen[k]) == 0 ? now : @inbounds(seen[k]))
+                @inbounds seen[k] = t
+                due = rescue_due(asleep, t, now, job.deadline)
+            end
+            if due
                 (@atomicreplace :monotonic ws.cmd cmd => 0).success || continue
                 wake_children(pool, cmd)
                 copies && clear_slot!(@inbounds wslots[k])
@@ -221,7 +245,9 @@ const RANGE_STARTED = UInt64(1) << 63
 @inline function init_blocks!(blocks::Vector{RangeBlock}, q::Int, r::Int, p::Int)
     for k in 1:p
         lo, hi = block_bounds(q, r, k)
-        @atomic :monotonic (@inbounds blocks[k]).word = pack_range(false, lo - 1, hi)
+        b = @inbounds blocks[k]
+        @atomic :monotonic b.word = pack_range(false, lo - 1, hi)
+        @atomic :monotonic b.seen = UInt64(0)
     end
     return nothing
 end
@@ -265,25 +291,31 @@ end
 @inline block_started(j::StealRanges, v::Int) =
     ((@atomic :monotonic (@inbounds j.blocks[v]).word) & RANGE_STARTED) != 0
 
-# Started blocks are freely stealable. For unstarted blocks, use the rescue rule:
-# watch awake workers with pending commands for RANGE_AWAKE_NS; protect sleeping workers,
-# claimed commands and the caller's block until the deadline.
+# `steal_allowed` verdicts: take; wait for an awake owner to claim or become due;
+# protected (owner asleep or claimed, or the caller's block) until the deadline
+const STEAL_YES = 0
+const STEAL_SOON = 1
+const STEAL_NO = 2
+
+# The verdict for an unstarted, nonempty block `b` whose owner's command is `pending`
+@inline function unstarted_verdict(
+        b::RangeBlock, pending::Bool, asleep::Bool, now::UInt64, deadline::UInt64
+    )
+    pending || return now >= deadline ? STEAL_YES : STEAL_NO
+    owner_due(b, asleep, now, deadline) && return STEAL_YES
+    return asleep ? STEAL_NO : STEAL_SOON
+end
+
+# Started and empty blocks are freely stealable; others follow the rescue rule. No
+# waiting here: a private 2 µs wait per call and block made loaded launches take ~1 ms.
 @inline function steal_allowed(j::StealRanges, v::Int)
-    block_started(j, v) && return true
+    b = @inbounds j.blocks[v]
+    started, lo, hi = unpack_range(@atomic :monotonic b.word)
+    (started || lo >= hi) && return STEAL_YES
     pool = j.pool
-    seen = UInt64(0)
-    while true
-        now = time_ns()
-        pending = (@atomic :monotonic (@inbounds pool.workers[v]).cmd) > 0
-        asleep = @atomic :monotonic (@inbounds pool.sleepflags[v]).sleeping
-        awake_pending = pending && !asleep
-        awake_pending && seen == 0 && (seen = now)
-        rescue_due(!awake_pending, seen, now, j.deadline) && return true
-        awake_pending || return false
-        cpu_pause()
-        block_started(j, v) && return true
-    end
-    return
+    pending = (@atomic :monotonic (@inbounds pool.workers[v]).cmd) > 0
+    asleep = @atomic :monotonic (@inbounds pool.sleepflags[v]).sleeping
+    return unstarted_verdict(b, pending, asleep, time_ns(), j.deadline)
 end
 
 function (j::StealRanges)(chunk::Int)
@@ -298,7 +330,7 @@ function (j::StealRanges)(chunk::Int)
     return nothing
 end
 
-# Retry protected blocks briefly before retiring.
+# Retry protected (`STEAL_NO`) blocks briefly before retiring.
 const PROTECTED_PASSES = 3
 
 # Run the participant's own block, by owner rules
@@ -318,12 +350,14 @@ end
     p = j.p
     retries = 0
     while true
-        found = protected = false
+        found = protected = soon = false
         v = me
         for _ in 2:p
             v = v == p ? 1 : v + 1
-            if !steal_allowed(j, v)
-                protected = true
+            verdict = steal_allowed(j, v)
+            if verdict != STEAL_YES
+                protected |= verdict == STEAL_NO
+                soon |= verdict == STEAL_SOON
                 continue
             end
             # Again from the same block while it has work: the next steal takes a
@@ -337,10 +371,11 @@ end
             end
         end
         ranges_stopped(st) && break
-        # Rescan after successful steals. Retry protected blocks briefly, then let their
-        # owners or the dispatcher handle them rather than polling until the deadline.
+        # Rescan after successful steals and while an owner may become due. Retry
+        # protected blocks briefly, then let their owners or the dispatcher handle them
+        # rather than polling until the deadline.
         if !found
-            (protected && (retries += 1) <= PROTECTED_PASSES) || break
+            soon || (protected && (retries += 1) <= PROTECTED_PASSES) || break
             j.pool.oversubscribed ? os_yield() : cpu_pause()
         end
     end

test/foreach_range.jl

@@ -231,6 +231,106 @@ end
         end
     end
 
+    @testset "steal clocks" begin
+        TT = ThreadTeams
+        YES, SOON, NO = TT.STEAL_YES, TT.STEAL_SOON, TT.STEAL_NO
+        A = UInt64(TT.RANGE_AWAKE_NS)
+        t, d = UInt64(10^9), UInt64(10^9 + 10^6)
+        # `owner_due` and `unstarted_verdict` with explicit clocks, on a block of its own.
+        # Observations in sequence: no simultaneous observers.
+        b = TT.RangeBlock()
+        seen() = @atomic b.seen
+        verdict(pending, asleep, now) = TT.unstarted_verdict(b, pending, asleep, now, d)
+        # an awake owner with its command pending: soon, then due, from the first
+        # observation, which later ones do not renew
+        @test verdict(true, false, t) == SOON && seen() == t
+        @test verdict(true, false, t + A - 1) == SOON && seen() == t
+        @test verdict(true, false, t + A) == YES && seen() == t
+        @test TT.owner_due(b, false, t + A, d) && !TT.owner_due(b, false, t + A - 1, d)
+        # asleep: protected until the deadline; the clock survives the observation
+        @test verdict(true, true, t + 2A) == NO && seen() == t
+        @test verdict(true, false, t + 2A) == YES
+        @test verdict(true, true, d) == YES
+        b = TT.RangeBlock()
+        @test verdict(true, true, t) == NO && seen() == 0
+        @test !TT.owner_due(b, true, t, d) && seen() == 0
+        # claimed (or no command, as the caller's block): protected until the deadline
+        @test verdict(false, false, t) == NO && verdict(false, false, d - 1) == NO
+        @test verdict(false, false, d) == YES && seen() == 0
+        # an observer that read the clock before another's timestamp waits, rather than
+        # wrapping around to due; later ones use the first timestamp
+        @test verdict(true, false, t + 5) == SOON && seen() == t + 5
+        @test verdict(true, false, t) == SOON && seen() == t + 5
+        @test !TT.owner_due(b, false, t, d)
+        @test verdict(true, false, t + 4 + A) == SOON
+        @test verdict(true, false, t + 5 + A) == YES
+        # a timestamp in the future is not due before the deadline
+        @atomic b.seen = d - 1
+        @test verdict(true, false, t) == SOON && !TT.owner_due(b, false, t, d)
+
+        # A pool whose workers never run: its commands stay pending, its owners awake
+        function idle_pool(nw)
+            return TT.Pool(
+                nothing, [TT.WorkerState() for _ in 1:nw], [Base.Event(true) for _ in 1:nw],
+                [Task(() -> nothing) for _ in 1:nw], Any[nothing for _ in 1:nw],
+                TT.pad_before(), false, TT.pad_after(), zeros(Int, Threads.maxthreadid()),
+                [TT.SleepFlag() for _ in 1:nw], TT.ChunkCounter(), Any[nothing for _ in 1:nw],
+                zeros(Int, nw), false, 0, TT.RangeState(nw)
+            )
+        end
+        # `1:100` in blocks 1:50 and 51:100; the caller (`skip = 1`) has run block 1, and
+        # worker 2 has chunk 2 pending
+        function fixture(deadline)
+            pool = idle_pool(2)
+            st = pool.ranges
+            calls = Tuple{Int, Int}[]
+            j = TT.StealRanges(
+                (lo, hi) -> (push!(calls, (lo, hi)); nothing), st.blocks, 50, 0, 2, 1, 10,
+                deadline, st, pool
+            )
+            TT.init_blocks!(st.blocks, 50, 0, 2)
+            @atomic st.blocks[1].word = TT.pack_range(true, 50, 50)
+            base = TT.cmd_base(2, 1, false)
+            @atomic pool.workers[2].cmd = base | 2
+            return pool, j, base, calls
+        end
+        # thieves rescan while an owner may become due, then take its block
+        pool, j, base, calls = fixture(typemax(UInt64))
+        @test TT.steal_allowed(j, 2) == SOON
+        @atomic j.blocks[2].seen = UInt64(0)   # the scan's first observation: soon
+        TT.run_steal(j, 1)
+        @test ((_, lo, hi) = TT.unpack_range(@atomic j.blocks[2].word); lo == hi)
+        @test sum(c -> c[2] - c[1] + 1, calls) == 50 && minimum(first, calls) == 51
+        # the dispatcher's rescue starts the block's shared clock...
+        pool, j, base, calls = fixture(typemax(UInt64))
+        slot = TT.Slot{typeof(j)}()
+        slot[] = j
+        TT.rescue_ranges(pool, j, slot, 2, 1, base, typeof(slot)[], false)
+        @test (@atomic pool.workers[2].cmd) == 0 && (@atomic j.blocks[2].seen) != 0
+        @test sum(c -> c[2] - c[1] + 1, calls) == 50
+        # ...and waits for one a thief started, here as if in the future
+        pool, j, base, calls = fixture(typemax(UInt64))
+        slot = TT.Slot{typeof(j)}()
+        slot[] = j
+        t0 = time_ns()
+        @atomic j.blocks[2].seen = t0 + 3_000_000
+        TT.rescue_ranges(pool, j, slot, 2, 1, base, typeof(slot)[], false)
+        @test time_ns() - t0 >= 3_000_000 && (@atomic pool.workers[2].cmd) == 0
+        @test sum(c -> c[2] - c[1] + 1, calls) == 50
+
+        # each region starts with fresh clocks
+        pool = @atomic TT.POOL.pool
+        if pool !== nothing
+            st = pool.ranges
+            for k in 1:NW
+                @atomic st.blocks[k].seen = UInt64(1)
+            end
+            t0 = time_ns()
+            foreach_range((lo, hi) -> nothing, 100NW; schedule = :steal)
+            @test all(k -> (s = @atomic st.blocks[k].seen; s == 0 || s >= t0), 1:NW)
+        end
+    end
+
     @testset "argument checks" begin
         @test_throws ArgumentError foreach_range((lo, hi) -> nothing, -1)
         @test_throws ArgumentError foreach_range((lo, hi) -> nothing, (typemax(Int) >> 2) + 1)