ReentrantLock: wakeup a single task on unlock and add a short spin

[andrebsguedes]

I propose a change in the implementation of the ReentrantLock to improve its overall throughput for short critical sections and fix the quadratic wake-up behavior where each unlock schedules all waiting tasks on the lock's wait queue.

This implementation follows the same principles of the Mutex in the parking_lot Rust crate which is based on the Webkit WTF::ParkingLot class. Only the basic working principle is implemented here, further improvements such as eventual fairness will be proposed separately.

The gist of the change is that we add one extra state to the lock, essentially going from:

0x0 => The lock is not locked
0x1 => The lock is locked by exactly one task. No other task is waiting for it.
0x2 => The lock is locked and some other task tried to lock but failed (conflict)

To:

# PARKED_BIT | LOCKED_BIT | Description
#     0      |     0      | The lock is not locked, nor is anyone waiting for it.
# -----------+------------+------------------------------------------------------------------
#     0      |     1      | The lock is locked by exactly one task. No other task is
#            |            | waiting for it.
# -----------+------------+------------------------------------------------------------------
#     1      |     0      | The lock is not locked. One or more tasks are parked.
# -----------+------------+------------------------------------------------------------------
#     1      |     1      | The lock is locked by exactly one task. One or more tasks are
#            |            | parked waiting for the lock to become available.
#            |            | In this state, PARKED_BIT is only ever cleared when the cond_wait lock
#            |            | is held (i.e. on unlock). This ensures that
#            |            | we never end up in a situation where there are parked tasks but
#            |            | PARKED_BIT is not set (which would result in those tasks
#            |            | potentially never getting woken up).

In the current implementation we must schedule all tasks to cause a conflict (state 0x2) because on unlock we only notify any task if the lock is in the conflict state. This behavior means that with high contention and a short critical section the tasks will be effectively spinning in the scheduler queue.

With the extra state the proposed implementation has enough information to know if there are other tasks to be notified or not, which means we can always notify one task at a time while preserving the optimized path of not notifying if there are no tasks waiting. To improve throughput for short critical sections we also introduce a bounded amount of spinning before attempting to park.

Results

Not spinning on the scheduler queue greatly reduces the CPU utilization of the following example:

function example()
    lock = ReentrantLock()
    @sync begin
        for i in 1:10000
            Threads.@spawn begin
                @lock lock begin
                    sleep(0.001)
                end
            end
        end
    end
end


@time example()

Current:

28.890623 seconds (101.65 k allocations: 7.646 MiB, 0.25% compilation time)

Proposed:

22.806669 seconds (101.65 k allocations: 7.814 MiB, 0.35% compilation time)

In a micro-benchmark where 8 threads contend for a single lock with a very short critical section we see a ~2x improvement.

Current:

8-element Vector{Int64}:
 6258688
 5373952
 6651904
 6389760
 6586368
 3899392
 5177344
 5505024
Total iterations: 45842432

Proposed:

8-element Vector{Int64}:
 12320768
 12976128
 10354688
 12845056
  7503872
 13598720
 13860864
 11993088
Total iterations: 95453184

In the uncontended scenario the extra bookkeeping causes a 10% throughput reduction:
EDIT: I reverted _trylock to the simple case to recover the uncontended throughput and now both implementations are on the same ballpark (without hurting the above numbers).

In the uncontended scenario:

Current:

Total iterations: 236748800

Proposed:

Total iterations: 237699072

Closes #56182

[oscardssmith]

is this better than just calling yield?

[andrebsguedes]

In this case we do not want to leave the core but instead just busy wait in-core for a small number of iterations before we attempt the compare_exchange again, this way if the critical section of the lock is small enough we have a chance to acquire the lock without paying for a OS thread context switch (or a Julia scheduler task switch if you mean Base.yield).
This is the same strategy employed by Rust here.

[oscardssmith]

I meant Base.yield. I think we're in a different situation than Rust since we have M:N threading and a user mode scheduler which Rust doesn't.

[andrebsguedes]

Yes, I am aware of the differences but the one that matters most in this case is that Julia has a cooperative scheduler. This means we have a tradeoff between micro-contention throughput (where we want to stay in-core) and being nice to the other tasks (by calling Base.yield).

So I wrote a benchmark that measures the total amount of forward progress that can be made both by the tasks participating in locking as well as other unrelated tasks to see where we reach a good balance in this tradeoff. It turns out yielding to the Julia scheduler up to a limited amount seems to work great and does not suffer from the pathological case of always spinning in the scheduler (the first example in the PR description).

I will update the code and post benchmark results soon-ish (my daughter arrived yesterday!).

Thanks, @gbaraldi for working on benchmarks also, maybe I can contribute this new benchmark to your LockBench.jl

[gbaraldi]

@andrebsguedes It would be lovely to add more benchmarks to it. Having a suite of benchmarks that stress locks in different ways would be great.

[adienes]

(this PR is fantastically written! professional, comprehensive, and easy to follow 👏👏)

[vchuravy]

Thanks! Very well written PR.

[gbaraldi]

Another point of comparison could be https://github.com/kprotty/usync which is just a "normal" lock

[gbaraldi]

I wrote https://github.com/gbaraldi/LockBench.jl and there it seems to be a good upgrade

[gbaraldi]

Master

PR

[andrebsguedes]

I updated the proposal to reflect the changes from tuning for fairness to non-locking tasks also.

The following benchmark introduces competing tasks that perform a unit of work and call Base.yield in a loop to understand how unrelated tasks are affected by locking tasks.

Current implementation (thundering heard, 8 threads, 8 tasks locking, 8 task yielding, short critical section)

Locking:
  488.60 kHz
  488.60 kHz
  482.49 kHz
  486.56 kHz
  482.49 kHz
  486.56 kHz
  484.52 kHz
  488.60 kHz
Yielding:
  348.12 kHz
  348.12 kHz
  348.12 kHz
  348.12 kHz
  348.12 kHz
  348.12 kHz
  348.12 kHz
  348.12 kHz
Total useful work: 10625024

Initial proposal (limited CPU spin, 8 threads, 8 tasks locking, 8 tasks yielding, short critical section)

Locking:
  916.67 kHz
  937.04 kHz
  910.56 kHz
  916.67 kHz
  918.71 kHz
  918.71 kHz
  926.86 kHz
  928.90 kHz
Yielding:
  38.70 kHz
  38.70 kHz
  38.70 kHz
  40.74 kHz
  38.70 kHz
  40.74 kHz
  40.74 kHz
  38.70 kHz
Total useful work: 15144960

New proposal (limited scheduler spin, 8 threads, 8 tasks locking, 8 tasks yielding, short critical section)

Locking:
  737.73 kHz
  749.96 kHz
  731.62 kHz
  739.77 kHz
  729.58 kHz
  735.69 kHz
  733.66 kHz
  735.69 kHz
Yielding:
  401.47 kHz
  403.51 kHz
  399.43 kHz
  403.51 kHz
  401.47 kHz
  403.51 kHz
  401.47 kHz
  401.47 kHz
Total useful work: 15077376

Some observations:

• The current implementation wastes more cycles in the scheduler as it is the one with the least amount of useful work.
• The initial proposal has the best micro-contention throughput but to achieve this it favors the locking tasks too much.
• The new proposal strikes the best balance between the locking and the yielding tasks while still being an improvement in micro-contention throughput and in the amount of useful work performed.
• The new proposal still performs as well as the previous one in the simple sleep example which means the limited spin in the scheduler is benign.