make our locks fast (JuliaLang#42964)

Now that we have the ability to express atomics (particularly the write barrier of the Task object), it is time for us to speed up and simplify our locks.

This PR does need someone to verify the synchronizations are correct, but I think this version is actually quite simple. There is a single byte that represents the state of the lock:

0x0 : unlocked
0x1 : locked (without other waiters)
0x2 : locked with waiters
And then when we have the locked bit set, we can manipulate any of the other data (e.g. owner and reentrancy) without needing any atomic work, and if we don't hold that bit, we should not write to any other field.

Author: vtjnash

base/lock.jl

@@ -26,11 +26,22 @@ end
 ```
 """
 mutable struct ReentrantLock <: AbstractLock
-    locked_by::Union{Task, Nothing}
-    cond_wait::ThreadSynchronizer
-    reentrancy_cnt::Int
-
-    ReentrantLock() = new(nothing, ThreadSynchronizer(), 0)
+    # offset = 16
+    @atomic locked_by::Union{Task, Nothing}
+    # offset32 = 20, offset64 = 24
+    reentrancy_cnt::UInt32
+    # offset32 = 24, offset64 = 28
+    @atomic havelock::UInt8 # 0x0 = none, 0x1 = lock, 0x2 = conflict
+    # offset32 = 28, offset64 = 32
+    cond_wait::ThreadSynchronizer # 2 words
+    # offset32 = 36, offset64 = 48
+    # sizeof32 = 20, sizeof64 = 32
+    # now add padding to make this a full cache line to minimize false sharing between objects
+    _::NTuple{Int === Int32 ? 2 : 3, Int}
+    # offset32 = 44, offset64 = 72 == sizeof+offset
+    # sizeof32 = 28, sizeof64 = 56
+
+    ReentrantLock() = new(nothing, 0x0000_0000, 0x00, ThreadSynchronizer())
 end
 
 assert_havelock(l::ReentrantLock) = assert_havelock(l, l.locked_by)
@@ -42,7 +53,7 @@ Check whether the `lock` is held by any task/thread.
 This should not be used for synchronization (see instead [`trylock`](@ref)).
 """
 function islocked(rl::ReentrantLock)
-    return rl.reentrancy_cnt != 0
+    return rl.havelock != 0
 end
 
 """
@@ -55,23 +66,26 @@ return `false`.
 
 Each successful `trylock` must be matched by an [`unlock`](@ref).
 """
-function trylock(rl::ReentrantLock)
-    t = current_task()
-    if t === rl.locked_by
-        rl.reentrancy_cnt += 1
+@inline function trylock(rl::ReentrantLock)
+    ct = current_task()
+    if rl.locked_by === ct
+        #@assert rl.havelock !== 0x00
+        rl.reentrancy_cnt += 0x0000_0001
         return true
     end
-    lock(rl.cond_wait)
-    if rl.reentrancy_cnt == 0
-        rl.locked_by = t
-        rl.reentrancy_cnt = 1
-        GC.disable_finalizers()
-        got = true
-    else
-        got = false
+    return _trylock(rl, ct)
+end
+@noinline function _trylock(rl::ReentrantLock, ct::Task)
+    GC.disable_finalizers()
+    if (@atomicreplace :acquire rl.havelock 0x00 => 0x01).success
+        #@assert rl.locked_by === nothing
+        #@assert rl.reentrancy_cnt === 0
+        rl.reentrancy_cnt = 0x0000_0001
+        @atomic :release rl.locked_by = ct
+        return true
     end
-    unlock(rl.cond_wait)
-    return got
+    GC.enable_finalizers()
+    return false
 end
 
 """
@@ -83,28 +97,23 @@ wait for it to become available.
 
 Each `lock` must be matched by an [`unlock`](@ref).
 """
-function lock(rl::ReentrantLock)
-    t = current_task()
-    if t === rl.locked_by
-        rl.reentrancy_cnt += 1
-    else
-        lock(rl.cond_wait)
-        while true
-            if rl.reentrancy_cnt == 0
-                rl.locked_by = t
-                rl.reentrancy_cnt = 1
-                GC.disable_finalizers()
-                break
-            end
-            try
-                wait(rl.cond_wait)
-            catch
-                unlock(rl.cond_wait)
-                rethrow()
+@inline function lock(rl::ReentrantLock)
+    trylock(rl) || (@noinline function slowlock(rl::ReentrantLock)
+        c = rl.cond_wait
+        lock(c.lock)
+        try
+            while true
+                if (@atomicreplace rl.havelock 0x01 => 0x02).old == 0x00 # :sequentially_consistent ? # now either 0x00 or 0x02
+                    # it was unlocked, so try to lock it ourself
+                    _trylock(rl, current_task()) && break
+                else # it was locked, so now wait for the release to notify us
+                    wait(c)
+                end
             end
+        finally
+            unlock(c.lock)
         end
-        unlock(rl.cond_wait)
-    end
+    end)(rl)
     return
 end
 
@@ -116,58 +125,42 @@ Releases ownership of the `lock`.
 If this is a recursive lock which has been acquired before, decrement an
 internal counter and return immediately.
 """
-function unlock(rl::ReentrantLock)
-    t = current_task()
-    n = rl.reentrancy_cnt
-    n == 0 && error("unlock count must match lock count")
-    rl.locked_by === t || error("unlock from wrong thread")
-    if n > 1
-        rl.reentrancy_cnt = n - 1
-    else
-        lock(rl.cond_wait)
-        rl.reentrancy_cnt = 0
-        rl.locked_by = nothing
-        if !isempty(rl.cond_wait.waitq)
-            try
-                notify(rl.cond_wait)
-            catch
-                unlock(rl.cond_wait)
-                rethrow()
+@inline function unlock(rl::ReentrantLock)
+    rl.locked_by === current_task() ||
+        error(rl.reentrancy_cnt == 0x0000_0000 ? "unlock count must match lock count" : "unlock from wrong thread")
+    (@noinline function _unlock(rl::ReentrantLock)
+        n = rl.reentrancy_cnt - 0x0000_0001
+        rl.reentrancy_cnt = n
+        if n == 0x0000_00000
+            @atomic :monotonic rl.locked_by = nothing
+            if (@atomicswap :release rl.havelock = 0x00) == 0x02
+                (@noinline function notifywaiters(rl)
+                    cond_wait = rl.cond_wait
+                    lock(cond_wait)
+                    try
+                        notify(cond_wait)
+                    finally
+                        unlock(cond_wait)
+                    end
+                end)(rl)
             end
+            return true
         end
-        GC.enable_finalizers()
-        unlock(rl.cond_wait)
-    end
-    return
+        return false
+    end)(rl) && GC.enable_finalizers()
+    nothing
 end
 
 function unlockall(rl::ReentrantLock)
-    t = current_task()
-    n = rl.reentrancy_cnt
-    rl.locked_by === t || error("unlock from wrong thread")
-    n == 0 && error("unlock count must match lock count")
-    lock(rl.cond_wait)
-    rl.reentrancy_cnt = 0
-    rl.locked_by = nothing
-    if !isempty(rl.cond_wait.waitq)
-        try
-            notify(rl.cond_wait)
-        catch
-            unlock(rl.cond_wait)
-            rethrow()
-        end
-    end
-    GC.enable_finalizers()
-    unlock(rl.cond_wait)
+    n = @atomicswap :not_atomic rl.reentrancy_cnt = 0x0000_0001
+    unlock(rl)
     return n
 end
 
-function relockall(rl::ReentrantLock, n::Int)
-    t = current_task()
+function relockall(rl::ReentrantLock, n::UInt32)
     lock(rl)
-    n1 = rl.reentrancy_cnt
-    rl.reentrancy_cnt = n
-    n1 == 1 || concurrency_violation()
+    old = @atomicswap :not_atomic rl.reentrancy_cnt = n
+    old == 0x0000_0001 || concurrency_violation()
     return
 end
 

base/locks-mt.jl

@@ -27,46 +27,19 @@ contending threads. If you have more contention than that, different
 synchronization approaches should be considered.
 """
 mutable struct SpinLock <: AbstractLock
-    owned::Int
+    # we make this much larger than necessary to minimize false-sharing
+    @atomic owned::Int
     SpinLock() = new(0)
 end
 
-import Base.Sys.WORD_SIZE
-
-@eval _xchg!(x::SpinLock, v::Int) =
-    llvmcall($"""
-             %ptr = inttoptr i$WORD_SIZE %0 to i$WORD_SIZE*
-             %rv = atomicrmw xchg i$WORD_SIZE* %ptr, i$WORD_SIZE %1 acq_rel
-             ret i$WORD_SIZE %rv
-             """, Int, Tuple{Ptr{Int}, Int}, unsafe_convert(Ptr{Int}, pointer_from_objref(x)), v)
-
-@eval _get(x::SpinLock) =
-    llvmcall($"""
-             %ptr = inttoptr i$WORD_SIZE %0 to i$WORD_SIZE*
-             %rv = load atomic i$WORD_SIZE, i$WORD_SIZE* %ptr monotonic, align $(gc_alignment(Int))
-             ret i$WORD_SIZE %rv
-             """, Int, Tuple{Ptr{Int}}, unsafe_convert(Ptr{Int}, pointer_from_objref(x)))
-
-@eval _set!(x::SpinLock, v::Int) =
-    llvmcall($"""
-             %ptr = inttoptr i$WORD_SIZE %0 to i$WORD_SIZE*
-             store atomic i$WORD_SIZE %1, i$WORD_SIZE* %ptr release, align $(gc_alignment(Int))
-             ret void
-             """, Cvoid, Tuple{Ptr{Int}, Int}, unsafe_convert(Ptr{Int}, pointer_from_objref(x)), v)
-
 # Note: this cannot assert that the lock is held by the correct thread, because we do not
 # track which thread locked it. Users beware.
 Base.assert_havelock(l::SpinLock) = islocked(l) ? nothing : Base.concurrency_violation()
 
 function lock(l::SpinLock)
     while true
-        if _get(l) == 0
-            GC.disable_finalizers()
-            p = _xchg!(l, 1)
-            if p == 0
-                return
-            end
-            GC.enable_finalizers()
+        if @inline trylock(l)
+            return
         end
         ccall(:jl_cpu_pause, Cvoid, ())
         # Temporary solution before we have gc transition support in codegen.
@@ -75,9 +48,9 @@ function lock(l::SpinLock)
 end
 
 function trylock(l::SpinLock)
-    if _get(l) == 0
+    if l.owned == 0
         GC.disable_finalizers()
-        p = _xchg!(l, 1)
+        p = @atomicswap :acquire l.owned = 1
         if p == 0
             return true
         end
@@ -87,13 +60,14 @@ function trylock(l::SpinLock)
 end
 
 function unlock(l::SpinLock)
-    _get(l) == 0 && error("unlock count must match lock count")
-    _set!(l, 0)
+    if (@atomicswap :release l.owned = 0) == 0
+        error("unlock count must match lock count")
+    end
     GC.enable_finalizers()
     ccall(:jl_cpu_wake, Cvoid, ())
     return
 end
 
 function islocked(l::SpinLock)
-    return _get(l) != 0
+    return l.owned != 0
 end

doc/src/base/multi-threading.md

@@ -9,14 +9,6 @@ Base.Threads.nthreads
 ```
 
 See also [Multi-Threading](@ref man-multithreading).
-## Synchronization
-
-```@docs
-Base.Threads.Condition
-Base.Threads.Event
-```
-
-See also [Synchronization](@ref lib-task-sync).
 
 ## Atomic operations
 

doc/src/base/parallel.md

@@ -26,6 +26,8 @@ Base.schedule
 
 ## [Synchronization](@id lib-task-sync)
 
+## Synchronization
+
 ```@docs
 Base.errormonitor
 Base.@sync
@@ -34,6 +36,8 @@ Base.fetch(t::Task)
 Base.timedwait
 
 Base.Condition
+Base.Threads.Condition
+Base.Threads.Event
 Base.notify
 
 Base.Semaphore

test/threads.jl

@@ -2,6 +2,27 @@
 
 using Test
 
+# simple sanity tests for locks under cooperative concurrent access
+let lk = ReentrantLock()
+    c1 = Base.Event()
+    c2 = Base.Event()
+    @test trylock(lk)
+    @test trylock(lk)
+    t1 = @async (notify(c1); lock(lk); unlock(lk); trylock(lk))
+    t2 = @async (notify(c2); trylock(lk))
+    wait(c1)
+    wait(c2)
+    @test t1.queue === lk.cond_wait.waitq
+    @test t2.queue !== lk.cond_wait.waitq
+    @test istaskdone(t2)
+    @test !fetch(t2)
+    unlock(lk)
+    @test t1.queue === lk.cond_wait.waitq
+    unlock(lk)
+    @test t1.queue !== lk.cond_wait.waitq
+    @test fetch(t1)
+end
+
 let cmd = `$(Base.julia_cmd()) --depwarn=error --rr-detach --startup-file=no threads_exec.jl`
     for test_nthreads in (1, 2, 4, 4) # run once to try single-threaded mode, then try a couple times to trigger bad races
         new_env = copy(ENV)

test/threads_exec.jl

@@ -27,6 +27,27 @@ end
 # (expected test duration is about 18-180 seconds)
 Timer(t -> killjob("KILLING BY THREAD TEST WATCHDOG\n"), 1200)
 
+# basic lock check
+if nthreads() > 1
+    let lk = Base.Threads.SpinLock()
+        c1 = Base.Event()
+        c2 = Base.Event()
+        @test trylock(lk)
+        @test !trylock(lk)
+        t1 = Threads.@spawn (notify(c1); lock(lk); unlock(lk); trylock(lk))
+        t2 = Threads.@spawn (notify(c2); trylock(lk))
+        Libc.systemsleep(0.1) # block our thread from scheduling for a bit
+        wait(c1)
+        wait(c2)
+        @test !fetch(t2)
+        @test istaskdone(t2)
+        @test !istaskdone(t1)
+        unlock(lk)
+        @test fetch(t1)
+        @test istaskdone(t1)
+    end
+end
+
 # threading constructs
 
 let a = zeros(Int, 2 * nthreads())