Re: [PATCH] lockdep: Speed up lockdep_unregister_key() with expedited RCU synchronization

[Boqun Feng <boqun.feng@gmail.com>]

On Mon, Mar 24, 2025 at 12:30:10PM -0700, Boqun Feng wrote:
> On Mon, Mar 24, 2025 at 12:21:07PM -0700, Boqun Feng wrote:
> > On Mon, Mar 24, 2025 at 01:23:50PM +0100, Eric Dumazet wrote:
> > [...]
> > > > > ---
> > > > >  kernel/locking/lockdep.c | 6 ++++--
> > > > >  1 file changed, 4 insertions(+), 2 deletions(-)
> > > > >
> > > > > diff --git a/kernel/locking/lockdep.c b/kernel/locking/lockdep.c
> > > > > index 4470680f02269..a79030ac36dd4 100644
> > > > > --- a/kernel/locking/lockdep.c
> > > > > +++ b/kernel/locking/lockdep.c
> > > > > @@ -6595,8 +6595,10 @@ void lockdep_unregister_key(struct lock_class_key *key)
> > > > >       if (need_callback)
> > > > >               call_rcu(&delayed_free.rcu_head, free_zapped_rcu);
> > > > >
> > > > > -     /* Wait until is_dynamic_key() has finished accessing k->hash_entry. */
> > > > > -     synchronize_rcu();
> > 
> > I feel a bit confusing even for the old comment, normally I would expect
> > the caller of lockdep_unregister_key() should guarantee the key has been
> > unpublished, in other words, there is no way a lockdep_unregister_key()
> > could race with a register_lock_class()/lockdep_init_map_type(). The
> > synchronize_rcu() is not needed then.
> > 
> > Let's say someone breaks my assumption above, then when doing a
> > register_lock_class() with a key about to be unregister, I cannot see
> > anything stops the following:
> > 
> > 	CPU 0				CPU 1
> > 	=====				=====
> > 	register_lock_class():
> > 	  ...
> > 	  } else if (... && !is_dynamic_key(lock->key)) {
> > 	  	// ->key is not unregistered yet, so this branch is not
> > 		// taken.
> > 	  	return NULL;
> > 	  }
> > 	  				lockdep_unregister_key(..);
> > 					// key unregister, can be free
> > 					// any time.
> > 	  key = lock->key->subkeys + subclass; // BOOM! UAF.
> > 
> > So either we don't need the synchronize_rcu() here or the
> > synchronize_rcu() doesn't help at all. Am I missing something subtle
> > here?
> > 
> 
> Oh! Maybe I was missing register_lock_class() must be called with irq
> disabled, which is also an RCU read-side critical section.
> 

Since register_lock_class() will be call with irq disabled, maybe hazard
pointers [1] is better because most of the case we only have nr_cpus
readers, so the potential hazard pointer slots are fixed.

So the below patch can reduce the time of the tc command from real ~1.7
second (v6.14) to real ~0.05 second (v6.14 + patch) in my test env,
which is not surprising given it's a dedicated hazard pointers for
lock_class_key.

Thoughts?

[1]: https://lore.kernel.org/lkml/20240917143402.930114-1-boqun.feng@gmail.com/

Regards,
Boqun

---------------------------------->8
From: Boqun Feng <boqun.feng@gmail.com>
Date: Mon, 24 Mar 2025 13:38:15 -0700
Subject: [PATCH] WIP

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 kernel/locking/lockdep.c | 65 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 4 deletions(-)

diff --git a/kernel/locking/lockdep.c b/kernel/locking/lockdep.c
index 4470680f0226..0b6e78d56cfe 100644
--- a/kernel/locking/lockdep.c
+++ b/kernel/locking/lockdep.c
@@ -111,6 +111,29 @@ late_initcall(kernel_lockdep_sysctls_init);
 DEFINE_PER_CPU(unsigned int, lockdep_recursion);
 EXPORT_PER_CPU_SYMBOL_GPL(lockdep_recursion);
 
+/* hazptr free always clears the slot */
+DEFINE_FREE(lockdep_key_hazptr, struct lock_class_key **, if (_T) smp_store_release((_T), NULL));
+DEFINE_PER_CPU(struct lock_class_key *, lockdep_key_hazptr);
+
+static void synchronize_lockdep_key_hazptr(struct lock_class_key *key)
+{
+	int cpu;
+
+	if (!key)
+		return;
+	/*
+	 * Synchronizes with the smp_mb() after protecting the pointer with
+	 * WRITE_ONCE().
+	 */
+	smp_mb();
+
+	for_each_possible_cpu(cpu) {
+		struct lock_class_key **hazptr = per_cpu_ptr(&lockdep_key_hazptr, cpu);
+
+		smp_cond_load_acquire(hazptr, VAL != key);
+	}
+}
+
 static __always_inline bool lockdep_enabled(void)
 {
 	if (!debug_locks)
@@ -1283,6 +1306,7 @@ static struct lock_class *
 register_lock_class(struct lockdep_map *lock, unsigned int subclass, int force)
 {
 	struct lockdep_subclass_key *key;
+	struct lock_class_key **key_hazptr __free(lockdep_key_hazptr) = NULL;
 	struct hlist_head *hash_head;
 	struct lock_class *class;
 	int idx;
@@ -1293,11 +1317,25 @@ register_lock_class(struct lockdep_map *lock, unsigned int subclass, int force)
 	if (likely(class))
 		goto out_set_class_cache;
 
+	/* Interrupts are disabled hence it's safe to acquire the hazptr slot */
+	key_hazptr = this_cpu_ptr(&lockdep_key_hazptr);
+
+	/* hazptr slot must be unusued */
+	BUG_ON(*key_hazptr);
+
 	if (!lock->key) {
 		if (!assign_lock_key(lock))
 			return NULL;
-	} else if (!static_obj(lock->key) && !is_dynamic_key(lock->key)) {
-		return NULL;
+	} else {
+		/* hazptr: protect the key */
+		WRITE_ONCE(*key_hazptr, lock->key);
+
+		/* Synchronizes with the smp_mb() in synchronize_lockdep_key_hazptr() */
+		smp_mb();
+
+		if (!static_obj(lock->key) && !is_dynamic_key(lock->key)) {
+			return NULL;
+		}
 	}
 
 	key = lock->key->subkeys + subclass;
@@ -4964,16 +5002,35 @@ void lockdep_init_map_type(struct lockdep_map *lock, const char *name,
 	 */
 	if (DEBUG_LOCKS_WARN_ON(!key))
 		return;
+
+	preempt_disable();
+	struct lock_class_key **hazptr __free(lockdep_key_hazptr) = this_cpu_ptr(&lockdep_key_hazptr);
+
+	/* hapztr: Protect the key */
+	WRITE_ONCE(*hazptr, key);
+
+	/* Synchronizes with the smp_mb() in synchronize_lockdep_key_hazptr() */
+	smp_mb();
+
 	/*
 	 * Sanity check, the lock-class key must either have been allocated
 	 * statically or must have been registered as a dynamic key.
 	 */
 	if (!static_obj(key) && !is_dynamic_key(key)) {
+		preempt_enable();
 		if (debug_locks)
 			printk(KERN_ERR "BUG: key %px has not been registered!\n", key);
 		DEBUG_LOCKS_WARN_ON(1);
 		return;
 	}
+
+	/* hazptr: Release the key */
+	smp_store_release(hazptr, NULL);
+
+	/* Do not auto clean up*/
+	hazptr = NULL;
+	preempt_enable();
+
 	lock->key = key;
 
 	if (unlikely(!debug_locks))
@@ -6595,8 +6652,8 @@ void lockdep_unregister_key(struct lock_class_key *key)
 	if (need_callback)
 		call_rcu(&delayed_free.rcu_head, free_zapped_rcu);
 
-	/* Wait until is_dynamic_key() has finished accessing k->hash_entry. */
-	synchronize_rcu();
+	/* Wait until register_lock_class() has finished accessing key. */
+	synchronize_lockdep_key_hazptr(key);
 }
 EXPORT_SYMBOL_GPL(lockdep_unregister_key);
 
-- 
2.48.1