Skip the reference pool mutex on attach when no decrefs are pending

[tobni]

Closes #6199.

[codspeed-hq]

Merging this PR will not alter performance

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 140 untouched benchmarks
🆕 7 new benchmarks

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
🆕	Simulation	empty_pool_attach	N/A	5.8 µs	N/A
🆕	WallTime	nested_attach_scaling/empty_pool[1]	N/A	4 ns	N/A
🆕	WallTime	nested_attach_scaling/empty_pool[2]	N/A	4 ns	N/A
🆕	WallTime	nested_attach_scaling/empty_pool[4]	N/A	6 ns	N/A
🆕	WallTime	nested_attach_scaling/sparse_pool[1]	N/A	4 ns	N/A
🆕	WallTime	nested_attach_scaling/sparse_pool[2]	N/A	4 ns	N/A
🆕	WallTime	nested_attach_scaling/sparse_pool[4]	N/A	7 ns	N/A

---

_{Comparing tobni:pool-dirty-flag (792c1ed) with main (d1e3be6)}

[ngoldbaum]

The 128 bit conversion benchmarks are the failure can be ignored.

[ngoldbaum]

FWIW having a lock-free fast path for a mutex makes sense to me as a general thing to help scaling but I haven't thought deeply about the implications here.

[tobni]

FWIW having a lock-free fast path for a mutex makes sense to me as a general thing to help scaling but I haven't thought deeply about the implications here.

Good to hear!

It is the de-facto bottleneck for pants concurrency model, so I am keen on this issue being resolved.

[davidhewitt]

Thanks for the PR! Is it possible to build a benchmark which demonstrates the pathological case here? As per the issue we've flip-flopped on this a fair bit, I'm definitely open to having the complexity if we can prove it's worth it.

[davidhewitt]

Should this pair use compare_exchange?

[tobni]

My understanding is that compare_exchange would re-introduce the contention point of every attach becoming a writer. I think it is benign that threads CAN fall through to pending_decrefs turning out empty.

TLDR;
No, I dont think it should.

[alex]

I think it's worth a comment making that explicit.

[tobni]

Thanks for the PR! Is it possible to build a benchmark which demonstrates the pathological case here? As per the issue we've flip-flopped on this a fair bit, I'm definitely open to having the complexity if we can prove it's worth it.

I am struggling a bit to build a benchmark that codspeed can execute that proves contention is reduced with this change. Would a bench that measures the fastpath be sufficient?

Edit: I have become more familiar with codspeed and will push a suggestion of walltime benchmark. This might be a larger maintenance burden than this change warrants, so please scrutinize 8530cd7.

[tobni]

@davidhewitt PTAL

[tobni]

These are aggregate results I net using the pushed benchmark locally. I have an 8 pysical-core machine.
The 2 cores is to simulate the CI runner. The unit is total attaches per second

empty_pool:

┌────────────────────┬───────┬────────┬───────┐
│      threads       │ main  │ branch │ ratio │
├────────────────────┼───────┼────────┼───────┤
│ 16 cores           │       │        │       │
├────────────────────┼───────┼────────┼───────┤
│ 1                  │ 104.0 │ 374.0  │ 3.6×  │
├────────────────────┼───────┼────────┼───────┤
│ 2                  │ 35.6  │ 725.9  │ 20.4× │
├────────────────────┼───────┼────────┼───────┤
│ 4                  │ 26.3  │ 1269.1 │ 48.3× │
├────────────────────┼───────┼────────┼───────┤
│ 2 cores            │       │        │       │
├────────────────────┼───────┼────────┼───────┤
│ 1                  │ 103.9 │ 373.1  │ 3.6×  │
├────────────────────┼───────┼────────┼───────┤
│ 2                  │ 35.6  │ 553.0  │ 15.5× │
├────────────────────┼───────┼────────┼───────┤
│ 4                  │ 35.0  │ 606.0  │ 17.3× │
└────────────────────┴───────┴────────┴───────┘

sparse_pool is ~ equivalent

[alex]

@davidhewitt I feel like we used to have this, and then it got removed at some point, do you remember the history? Am I mixing this up with something else?

[tobni]

@davidhewitt I feel like we used to have this, and then it got removed at some point, do you remember the history? Am I mixing this up with something else?

I did some digging and documented my archeology findings in #6199. I believe that illuminates the history.

Edit:
It is specifically this change #3250 that removed the dirty boolean prior. I note that the implementation that was removed was using swap, a read-modify-write instruction (i.e every attach is a write). I'm not familiar with the 2023 code, but it looks to have tracked increfs as well as decrefs.

[tobni]

@davidhewitt is there anything I can do to move this along? pants is moving ahead with our free-threaded migration for next release, and this change is a large performance improvement for our use case.

[alex]

Code itself looks reasonable to me, I didn't review the benchmark/CI changes.

[davidhewitt]

Thanks for pushing this forward, I'm sufficiently convinced this is a good idea.

Testing locally the scaling benchmark I was able to push performance further by another ~20% with the following patch. The ideas:

• I used Relaxed ordering for all ops as the mutex already ensures data synchronization
• I swapped the store to a compare_exchange which helps to ensure that only one drainer ever gets to lock the mutex.
  • We don't write to the dirty flag under the mutex, so it's possible for false positives and multiple drainers to still get to the mutex, but in practice very rare. I tested moving the dirty flag writes under the mutex but this was slower in my testing.
• I split the function after the first load to a #[cold] slow function.

I tested each of these in isolation and am sufficiently convinced that each had a positive impact; the 20% is from the net effect of the three.

diff --git a/src/internal/state.rs b/src/internal/state.rs
index 71fc2b5f5..f7f37dda2 100644
--- a/src/internal/state.rs
+++ b/src/internal/state.rs
@@ -192,7 +192,9 @@ type PyObjVec = Vec<NonNull<ffi::PyObject>>;
 #[cfg(not(pyo3_disable_reference_pool))]
 /// Thread-safe storage for objects which were dec_ref while not attached.
 struct ReferencePool {
-    // Whether any decrefs are (or may be) pending
+    // Whether any decrefs are (or may be) pending. The `Mutex` performs
+    // synchronization so we can use `Relaxed` ordering for all operations
+    // on this flag.
     dirty: AtomicBool,
     pending_decrefs: Mutex<PyObjVec>,
 }
@@ -208,16 +210,41 @@ impl ReferencePool {
 
     fn register_decref(&self, obj: NonNull<ffi::PyObject>) {
         self.pending_decrefs.lock().unwrap().push(obj);
-        self.dirty.store(true, Ordering::Release);
+        self.dirty.store(true, Ordering::Relaxed);
     }
 
-    fn drop_deferred_references(&self, _py: Python<'_>) {
-        if !self.dirty.load(Ordering::Acquire) {
+    fn drop_deferred_references(&self, py: Python<'_>) {
+        // Check the dirty flag first to avoid any possible contention from atomic
+        // RMW operation to update the dirty flag on a hit.
+        if !self.dirty.load(Ordering::Relaxed) {
+            return;
+        }
+
+        // dirty flag is set, we _probably_ need to drop references (the flag is
+        // not updated under the mutex so false positives are possible but rare)
+        self.drop_deferred_references_slow(py);
+    }
+
+    #[cold]
+    fn drop_deferred_references_slow(&self, _py: Python<'_>) {
+        // Compare and swap the dirty flag to false avoids multiple threads from having
+        // contention on the mutex.
+        if self
+            .dirty
+            .compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed)
+            .is_err()
+        {
+            // Another thread is already dropping the references, so we can return early.
             return;
         }
-        self.dirty.store(false, Ordering::Relaxed);
 
         let mut pending_decrefs = self.pending_decrefs.lock().unwrap();
+        if pending_decrefs.is_empty() {
+            // We don't set the dirty flag under the mutex so it's possible to reach
+            // this case as a false positive. Returning early avoids a store on false
+            // positives.
+            return;
+        }
         let decrefs = mem::take(&mut *pending_decrefs);
         drop(pending_decrefs);

@tobni perhaps willing to accept this patch on top, and then let's merge?

[tobni]

Yep!

[tobni]

Seems I'm bitten by flaky CI again?

[alex]

I think this is caused by the base branch being out of date, can you rebase on a recent main?