[RFC lockdep 0/4] Fixes and self testcases for wait context detection

[RFC lockdep 0/4] Fixes and self testcases for wait context detection

[Boqun Feng]

Hi Peter,

Recently I looked into the wait context check feature and found some
places could use fixes, besides a suite of test cases is also added to
verify these fixes and future development.

Note: I'm not 100% sure all the expected results of the test cases are
correct, please do have a look at the the comment of patch #4 in case I
miss something subtle.

Suggestion and comments are welcome!

Regards,
Boqun


Boqun Feng (4):
  lockdep/selftest: Make HARDIRQ context threaded
  lockdep: Allow wait context checking with empty ->held_locks
  rcu/lockdep: Annotate the rcu_callback_map with proper wait contexts
  lockdep/selftest: Add wait context selftests

 kernel/locking/lockdep.c |   6 +-
 kernel/rcu/update.c      |   8 +-
 lib/locking-selftest.c   | 233 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 244 insertions(+), 3 deletions(-)

-- 
2.29.2

[RFC lockdep 1/4] lockdep/selftest: Make HARDIRQ context threaded

[Boqun Feng]

Since we now use spinlock_t instead of raw_spinlock_t in lockdep self
tests, we should make the emulated HARDIRQ context threaded, otherwise,
spinlock_t cannot be used in the HARDIRQ context and some test cases
will fail because of wait context checking when
PROVE_RAW_LOCK_NESTING=y.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 lib/locking-selftest.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lib/locking-selftest.c b/lib/locking-selftest.c
index afa7d4bb291f..0af91a07fd18 100644
--- a/lib/locking-selftest.c
+++ b/lib/locking-selftest.c
@@ -186,6 +186,7 @@ static void init_shared_classes(void)
 #define HARDIRQ_ENTER()				\
 	local_irq_disable();			\
 	__irq_enter();				\
+	lockdep_hardirq_threaded();		\
 	WARN_ON(!in_irq());
 
 #define HARDIRQ_EXIT()				\
-- 
2.29.2

[RFC lockdep 2/4] lockdep: Allow wait context checking with empty ->held_locks

[Boqun Feng]

Currently, the guard for !curr->lockdep_depth in check_wait_context() is
unnecessary, because the code will work well without it. Moreover, there
are cases that we will miss if we skip for curr->lockdep_depth == 0. For
example:

	<in hardirq context>
	some_irq_handler():
	  // curr->lockdep_depth == 0
	  mutex_lock(&some_mutex):
	    check_wait_context() // skip the check!

Clearly, it's a bug, but due to the skip for !curr->lockdep_depth, we
cannot detect it in check_wait_context().

Therefore, remove the !curr->lockdep_depth checks and add comments on
why it's still working without it. The idea is that if we currently
don't hold any lock, then the current context is the only one we should
use to check.

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

diff --git a/kernel/locking/lockdep.c b/kernel/locking/lockdep.c
index c1418b47f625..d4fd52b22804 100644
--- a/kernel/locking/lockdep.c
+++ b/kernel/locking/lockdep.c
@@ -4508,7 +4508,7 @@ static int check_wait_context(struct task_struct *curr, struct held_lock *next)
 	short curr_inner;
 	int depth;
 
-	if (!curr->lockdep_depth || !next_inner || next->trylock)
+	if (!next_inner || next->trylock)
 		return 0;
 
 	if (!next_outer)
@@ -4516,6 +4516,10 @@ static int check_wait_context(struct task_struct *curr, struct held_lock *next)
 
 	/*
 	 * Find start of current irq_context..
+	 *
+	 * Note: if curr->lockdep_depth == 0, we have depth == 0 after the
+	 * "depth++" below, and will skip the second for loop, i.e. using
+	 * the current task context as the curr_inner.
 	 */
 	for (depth = curr->lockdep_depth - 1; depth >= 0; depth--) {
 		struct held_lock *prev = curr->held_locks + depth;
-- 
2.29.2

[RFC lockdep 3/4] rcu/lockdep: Annotate the rcu_callback_map with proper wait contexts

[Boqun Feng]

rcu_callback_map is a virtual lock to annotate a context where RCU
callbacks are executed. RCU callbacks are required in softirq disable
contexts because with some config combination (use_softirq is true and
nocb is disabled) RCU callbacks only execute in softirq contexts.
Therefore wait context annotations can be added to detect bugs like
using mutex in a RCU callback.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 kernel/rcu/update.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/kernel/rcu/update.c b/kernel/rcu/update.c
index 39334d2d2b37..dd59e6412f61 100644
--- a/kernel/rcu/update.c
+++ b/kernel/rcu/update.c
@@ -269,8 +269,12 @@ EXPORT_SYMBOL_GPL(rcu_sched_lock_map);
 
 // Tell lockdep when RCU callbacks are being invoked.
 static struct lock_class_key rcu_callback_key;
-struct lockdep_map rcu_callback_map =
-	STATIC_LOCKDEP_MAP_INIT("rcu_callback", &rcu_callback_key);
+struct lockdep_map rcu_callback_map = {
+	.name = "rcu_callback",
+	.key = &rcu_callback_key,
+	.wait_type_outer = LD_WAIT_FREE,
+	.wait_type_inner = LD_WAIT_CONFIG, /* RCU callbacks are handled in softirq context */
+};
 EXPORT_SYMBOL_GPL(rcu_callback_map);
 
 noinstr int notrace debug_lockdep_rcu_enabled(void)
-- 
2.29.2

[RFC lockdep 4/4] lockdep/selftest: Add wait context selftests

[Boqun Feng]

These tests are added for two purposes:

*	Test the implementation of wait context checks and related
	annotations.

*	Semi-document the rules for wait context nesting when
	PROVE_RAW_LOCK_NESTING=y.

The test cases are only avaible for PROVE_RAW_LOCK_NESTING=y, as wait
context checking makes more sense for that configuration.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
---
 lib/locking-selftest.c | 232 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 232 insertions(+)

diff --git a/lib/locking-selftest.c b/lib/locking-selftest.c
index 0af91a07fd18..c00ef4e69637 100644
--- a/lib/locking-selftest.c
+++ b/lib/locking-selftest.c
@@ -63,6 +63,9 @@ static DEFINE_SPINLOCK(lock_B);
 static DEFINE_SPINLOCK(lock_C);
 static DEFINE_SPINLOCK(lock_D);
 
+static DEFINE_RAW_SPINLOCK(raw_lock_A);
+static DEFINE_RAW_SPINLOCK(raw_lock_B);
+
 static DEFINE_RWLOCK(rwlock_A);
 static DEFINE_RWLOCK(rwlock_B);
 static DEFINE_RWLOCK(rwlock_C);
@@ -1306,6 +1309,7 @@ GENERATE_PERMUTATIONS_3_EVENTS(irq_read_recursion3_soft_wlock)
 
 #ifdef CONFIG_DEBUG_LOCK_ALLOC
 # define I_SPINLOCK(x)	lockdep_reset_lock(&lock_##x.dep_map)
+# define I_RAW_SPINLOCK(x)	lockdep_reset_lock(&raw_lock_##x.dep_map)
 # define I_RWLOCK(x)	lockdep_reset_lock(&rwlock_##x.dep_map)
 # define I_MUTEX(x)	lockdep_reset_lock(&mutex_##x.dep_map)
 # define I_RWSEM(x)	lockdep_reset_lock(&rwsem_##x.dep_map)
@@ -1315,6 +1319,7 @@ GENERATE_PERMUTATIONS_3_EVENTS(irq_read_recursion3_soft_wlock)
 #endif
 #else
 # define I_SPINLOCK(x)
+# define I_RAW_SPINLOCK(x)
 # define I_RWLOCK(x)
 # define I_MUTEX(x)
 # define I_RWSEM(x)
@@ -1358,9 +1363,12 @@ static void reset_locks(void)
 	I1(A); I1(B); I1(C); I1(D);
 	I1(X1); I1(X2); I1(Y1); I1(Y2); I1(Z1); I1(Z2);
 	I_WW(t); I_WW(t2); I_WW(o.base); I_WW(o2.base); I_WW(o3.base);
+	I_RAW_SPINLOCK(A); I_RAW_SPINLOCK(B);
 	lockdep_reset();
 	I2(A); I2(B); I2(C); I2(D);
 	init_shared_classes();
+	raw_spin_lock_init(&raw_lock_A);
+	raw_spin_lock_init(&raw_lock_B);
 
 	ww_mutex_init(&o, &ww_lockdep); ww_mutex_init(&o2, &ww_lockdep); ww_mutex_init(&o3, &ww_lockdep);
 	memset(&t, 0, sizeof(t)); memset(&t2, 0, sizeof(t2));
@@ -2358,6 +2366,226 @@ static void queued_read_lock_tests(void)
 	pr_cont("\n");
 }
 
+#define __guard(cleanup) __maybe_unused __attribute__((__cleanup__(cleanup)))
+
+static void hardirq_exit(int *_)
+{
+	HARDIRQ_EXIT();
+}
+
+#define HARDIRQ_CONTEXT(name, ...)					\
+	int hardirq_guard_##name __guard(hardirq_exit);			\
+	HARDIRQ_ENTER();
+
+#define NOTTHREADED_HARDIRQ_CONTEXT(name, ...)				\
+	int notthreaded_hardirq_guard_##name __guard(hardirq_exit);	\
+	local_irq_disable();						\
+	__irq_enter();							\
+	WARN_ON(!in_irq());
+
+static void softirq_exit(int *_)
+{
+	SOFTIRQ_EXIT();
+}
+
+#define SOFTIRQ_CONTEXT(name, ...)				\
+	int softirq_guard_##name __guard(softirq_exit);		\
+	SOFTIRQ_ENTER();
+
+static void rcu_exit(int *_)
+{
+	rcu_read_unlock();
+}
+
+#define RCU_CONTEXT(name, ...)					\
+	int rcu_guard_##name __guard(rcu_exit);			\
+	rcu_read_lock();
+
+static void rcu_bh_exit(int *_)
+{
+	rcu_read_unlock_bh();
+}
+
+#define RCU_BH_CONTEXT(name, ...)				\
+	int rcu_bh_guard_##name __guard(rcu_bh_exit);		\
+	rcu_read_lock_bh();
+
+static void rcu_sched_exit(int *_)
+{
+	rcu_read_unlock_sched();
+}
+
+#define RCU_SCHED_CONTEXT(name, ...)				\
+	int rcu_sched_guard_##name __guard(rcu_sched_exit);	\
+	rcu_read_lock_sched();
+
+static void rcu_callback_exit(int *_)
+{
+	rcu_lock_release(&rcu_callback_map);
+}
+
+#define RCU_CALLBACK_CONTEXT(name, ...)					\
+	int rcu_callback_guard_##name __guard(rcu_callback_exit);	\
+	rcu_lock_acquire(&rcu_callback_map);
+
+
+static void raw_spinlock_exit(raw_spinlock_t **lock)
+{
+	raw_spin_unlock(*lock);
+}
+
+#define RAW_SPINLOCK_CONTEXT(name, lock)						\
+	raw_spinlock_t *raw_spinlock_guard_##name __guard(raw_spinlock_exit) = &(lock);	\
+	raw_spin_lock(&(lock));
+
+static void spinlock_exit(spinlock_t **lock)
+{
+	spin_unlock(*lock);
+}
+
+#define SPINLOCK_CONTEXT(name, lock)						\
+	spinlock_t *spinlock_guard_##name __guard(spinlock_exit) = &(lock);	\
+	spin_lock(&(lock));
+
+static void mutex_exit(struct mutex **lock)
+{
+	mutex_unlock(*lock);
+}
+
+#define MUTEX_CONTEXT(name, lock)					\
+	struct mutex *mutex_guard_##name __guard(mutex_exit) = &(lock);	\
+	mutex_lock(&(lock));
+
+#define GENERATE_2_CONTEXT_TESTCASE(outer, outer_lock, inner, inner_lock)	\
+										\
+static void __maybe_unused inner##_in_##outer(void)				\
+{										\
+	outer##_CONTEXT(_, outer_lock);						\
+	{									\
+		inner##_CONTEXT(_, inner_lock);					\
+	}									\
+}
+
+/*
+ * wait contexts (considering PREEMPT_RT)
+ *
+ * o: inner is allowed in outer
+ * x: inner is disallowed in outer
+ *
+ *       \  inner |  RCU  | RAW_SPIN | SPIN | MUTEX
+ * outer  \       |       |          |      |
+ * ---------------+-------+----------+------+-------
+ * HARDIRQ        |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * NOTTHREADED_IRQ|   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * SOFTIRQ        |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU            |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_BH         |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_CALLBACK   |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_SCHED      |   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * RAW_SPIN       |   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * SPIN           |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * MUTEX          |   o   |    o     |  o   |  o
+ * ---------------+-------+----------+------+-------
+ */
+
+#define GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(HARDIRQ, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(NOTTHREADED_HARDIRQ, , inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(SOFTIRQ, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU, , inner, inner_lock)				\
+GENERATE_2_CONTEXT_TESTCASE(RCU_BH, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU_CALLBACK, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU_SCHED, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RAW_SPINLOCK, raw_lock_A, inner, inner_lock)	\
+GENERATE_2_CONTEXT_TESTCASE(SPINLOCK, lock_A, inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(MUTEX, mutex_A, inner, inner_lock)
+
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(RCU, )
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(RAW_SPINLOCK, raw_lock_B)
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(SPINLOCK, lock_B)
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(MUTEX, mutex_B)
+
+/* the outer context allows all kinds of preemption */
+#define DO_CONTEXT_TESTCASE_OUTER_PREEMPTIBLE(outer)			\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, SUCCESS, LOCKTYPE_MUTEX);		\
+
+/*
+ * the outer context only allows the preemption introduced by spinlock_t (which
+ * is a sleepable lock for PREEMPT_RT)
+ */
+#define DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(outer)		\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, FAILURE, LOCKTYPE_MUTEX);		\
+
+/* the outer doesn't allows any kind of preemption */
+#define DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(outer)			\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, FAILURE, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, FAILURE, LOCKTYPE_MUTEX);		\
+
+static void wait_context_tests(void)
+{
+	printk("  --------------------------------------------------------------------------\n");
+	printk("  | wait context tests |\n");
+	printk("  --------------------------------------------------------------------------\n");
+	printk("                                 | rcu  | raw  | spin |mutex |\n");
+	printk("  --------------------------------------------------------------------------\n");
+	print_testname("in hardirq context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(HARDIRQ);
+	pr_cont("\n");
+
+	print_testname("in hardirq context (not threaded)");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(NOTTHREADED_HARDIRQ);
+	pr_cont("\n");
+
+	print_testname("in softirq context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(SOFTIRQ);
+	pr_cont("\n");
+
+	print_testname("in RCU context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU);
+	pr_cont("\n");
+
+	print_testname("in RCU-bh context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU_BH);
+	pr_cont("\n");
+
+	print_testname("in RCU callback context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU_CALLBACK);
+	pr_cont("\n");
+
+	print_testname("in RCU-sched context");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(RCU_SCHED);
+	pr_cont("\n");
+
+	print_testname("in RAW_SPINLOCK context");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(RAW_SPINLOCK);
+	pr_cont("\n");
+
+	print_testname("in SPINLOCK context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(SPINLOCK);
+	pr_cont("\n");
+
+	print_testname("in MUTEX context");
+	DO_CONTEXT_TESTCASE_OUTER_PREEMPTIBLE(MUTEX);
+	pr_cont("\n");
+}
+
 void locking_selftest(void)
 {
 	/*
@@ -2479,6 +2707,10 @@ void locking_selftest(void)
 	if (IS_ENABLED(CONFIG_QUEUED_RWLOCKS))
 		queued_read_lock_tests();
 
+	/* Wait context test cases that are specific for RAW_LOCK_NESTING */
+	if (IS_ENABLED(CONFIG_PROVE_RAW_LOCK_NESTING))
+		wait_context_tests();
+
 	if (unexpected_testcase_failures) {
 		printk("-----------------------------------------------------------------\n");
 		debug_locks = 0;
-- 
2.29.2

Re: [RFC lockdep 4/4] lockdep/selftest: Add wait context selftests

[Peter Zijlstra]

On Tue, Dec 08, 2020 at 06:31:12PM +0800, Boqun Feng wrote:
> These tests are added for two purposes:
> 
> *	Test the implementation of wait context checks and related
> 	annotations.
> 
> *	Semi-document the rules for wait context nesting when
> 	PROVE_RAW_LOCK_NESTING=y.

Documentation/locking/locktypes.rst should have that.

> The test cases are only avaible for PROVE_RAW_LOCK_NESTING=y, as wait
> context checking makes more sense for that configuration.

Looks about right ;-)

Re: [RFC lockdep 4/4] lockdep/selftest: Add wait context selftests

[Boqun Feng]

On Tue, Dec 08, 2020 at 03:33:24PM +0100, Peter Zijlstra wrote:
> On Tue, Dec 08, 2020 at 06:31:12PM +0800, Boqun Feng wrote:
> > These tests are added for two purposes:
> > 
> > *	Test the implementation of wait context checks and related
> > 	annotations.
> > 
> > *	Semi-document the rules for wait context nesting when
> > 	PROVE_RAW_LOCK_NESTING=y.
> 
> Documentation/locking/locktypes.rst should have that.
> 

Thanks for the pointer!

I miss it before, and it's really a comprehensive document for lock
nesting rules. Still I think more rules can be (and should be) put in
that document: a broader idea is the context nesting rule (e.g. whether
a spinlock_t is allowed in a hard irq handler). And the document
reminders me that I'm missing some locks (e.g local_lock) in the test
cases. So will improve both the document and the test cases in the next
version. In the meanwhile, feel free to point out any mistake or
misunderstanding of mine in the rules or the tests, I'm really still
learning through these locks with PREEMPT_RT into consideration, thanks!

Regards,
Boqun

> > The test cases are only avaible for PROVE_RAW_LOCK_NESTING=y, as wait
> > context checking makes more sense for that configuration.
> 
> Looks about right ;-)

[tip: locking/core] lockdep/selftest: Add wait context selftests

[tip-bot2 for Boqun Feng]

The following commit has been merged into the locking/core branch of tip:

Commit-ID:     9271a40d2a1429113160ccc4c16150921600bcc1
Gitweb:        https://git.kernel.org/tip/9271a40d2a1429113160ccc4c16150921600bcc1
Author:        Boqun Feng <boqun.feng@gmail.com>
AuthorDate:    Tue, 08 Dec 2020 18:31:12 +08:00
Committer:     Peter Zijlstra <peterz@infradead.org>
CommitterDate: Thu, 14 Jan 2021 11:20:16 +01:00

lockdep/selftest: Add wait context selftests

These tests are added for two purposes:

*	Test the implementation of wait context checks and related
	annotations.

*	Semi-document the rules for wait context nesting when
	PROVE_RAW_LOCK_NESTING=y.

The test cases are only avaible for PROVE_RAW_LOCK_NESTING=y, as wait
context checking makes more sense for that configuration.

Signed-off-by: Boqun Feng <boqun.feng@gmail.com>
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Link: https://lkml.kernel.org/r/20201208103112.2838119-5-boqun.feng@gmail.com
---
 lib/locking-selftest.c | 232 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 232 insertions(+)

diff --git a/lib/locking-selftest.c b/lib/locking-selftest.c
index 9959ea2..23376ee 100644
--- a/lib/locking-selftest.c
+++ b/lib/locking-selftest.c
@@ -64,6 +64,9 @@ static DEFINE_SPINLOCK(lock_B);
 static DEFINE_SPINLOCK(lock_C);
 static DEFINE_SPINLOCK(lock_D);
 
+static DEFINE_RAW_SPINLOCK(raw_lock_A);
+static DEFINE_RAW_SPINLOCK(raw_lock_B);
+
 static DEFINE_RWLOCK(rwlock_A);
 static DEFINE_RWLOCK(rwlock_B);
 static DEFINE_RWLOCK(rwlock_C);
@@ -1306,6 +1309,7 @@ GENERATE_PERMUTATIONS_3_EVENTS(irq_read_recursion3_soft_wlock)
 
 #ifdef CONFIG_DEBUG_LOCK_ALLOC
 # define I_SPINLOCK(x)	lockdep_reset_lock(&lock_##x.dep_map)
+# define I_RAW_SPINLOCK(x)	lockdep_reset_lock(&raw_lock_##x.dep_map)
 # define I_RWLOCK(x)	lockdep_reset_lock(&rwlock_##x.dep_map)
 # define I_MUTEX(x)	lockdep_reset_lock(&mutex_##x.dep_map)
 # define I_RWSEM(x)	lockdep_reset_lock(&rwsem_##x.dep_map)
@@ -1315,6 +1319,7 @@ GENERATE_PERMUTATIONS_3_EVENTS(irq_read_recursion3_soft_wlock)
 #endif
 #else
 # define I_SPINLOCK(x)
+# define I_RAW_SPINLOCK(x)
 # define I_RWLOCK(x)
 # define I_MUTEX(x)
 # define I_RWSEM(x)
@@ -1358,9 +1363,12 @@ static void reset_locks(void)
 	I1(A); I1(B); I1(C); I1(D);
 	I1(X1); I1(X2); I1(Y1); I1(Y2); I1(Z1); I1(Z2);
 	I_WW(t); I_WW(t2); I_WW(o.base); I_WW(o2.base); I_WW(o3.base);
+	I_RAW_SPINLOCK(A); I_RAW_SPINLOCK(B);
 	lockdep_reset();
 	I2(A); I2(B); I2(C); I2(D);
 	init_shared_classes();
+	raw_spin_lock_init(&raw_lock_A);
+	raw_spin_lock_init(&raw_lock_B);
 
 	ww_mutex_init(&o, &ww_lockdep); ww_mutex_init(&o2, &ww_lockdep); ww_mutex_init(&o3, &ww_lockdep);
 	memset(&t, 0, sizeof(t)); memset(&t2, 0, sizeof(t2));
@@ -2419,6 +2427,226 @@ static void fs_reclaim_tests(void)
 	pr_cont("\n");
 }
 
+#define __guard(cleanup) __maybe_unused __attribute__((__cleanup__(cleanup)))
+
+static void hardirq_exit(int *_)
+{
+	HARDIRQ_EXIT();
+}
+
+#define HARDIRQ_CONTEXT(name, ...)					\
+	int hardirq_guard_##name __guard(hardirq_exit);			\
+	HARDIRQ_ENTER();
+
+#define NOTTHREADED_HARDIRQ_CONTEXT(name, ...)				\
+	int notthreaded_hardirq_guard_##name __guard(hardirq_exit);	\
+	local_irq_disable();						\
+	__irq_enter();							\
+	WARN_ON(!in_irq());
+
+static void softirq_exit(int *_)
+{
+	SOFTIRQ_EXIT();
+}
+
+#define SOFTIRQ_CONTEXT(name, ...)				\
+	int softirq_guard_##name __guard(softirq_exit);		\
+	SOFTIRQ_ENTER();
+
+static void rcu_exit(int *_)
+{
+	rcu_read_unlock();
+}
+
+#define RCU_CONTEXT(name, ...)					\
+	int rcu_guard_##name __guard(rcu_exit);			\
+	rcu_read_lock();
+
+static void rcu_bh_exit(int *_)
+{
+	rcu_read_unlock_bh();
+}
+
+#define RCU_BH_CONTEXT(name, ...)				\
+	int rcu_bh_guard_##name __guard(rcu_bh_exit);		\
+	rcu_read_lock_bh();
+
+static void rcu_sched_exit(int *_)
+{
+	rcu_read_unlock_sched();
+}
+
+#define RCU_SCHED_CONTEXT(name, ...)				\
+	int rcu_sched_guard_##name __guard(rcu_sched_exit);	\
+	rcu_read_lock_sched();
+
+static void rcu_callback_exit(int *_)
+{
+	rcu_lock_release(&rcu_callback_map);
+}
+
+#define RCU_CALLBACK_CONTEXT(name, ...)					\
+	int rcu_callback_guard_##name __guard(rcu_callback_exit);	\
+	rcu_lock_acquire(&rcu_callback_map);
+
+
+static void raw_spinlock_exit(raw_spinlock_t **lock)
+{
+	raw_spin_unlock(*lock);
+}
+
+#define RAW_SPINLOCK_CONTEXT(name, lock)						\
+	raw_spinlock_t *raw_spinlock_guard_##name __guard(raw_spinlock_exit) = &(lock);	\
+	raw_spin_lock(&(lock));
+
+static void spinlock_exit(spinlock_t **lock)
+{
+	spin_unlock(*lock);
+}
+
+#define SPINLOCK_CONTEXT(name, lock)						\
+	spinlock_t *spinlock_guard_##name __guard(spinlock_exit) = &(lock);	\
+	spin_lock(&(lock));
+
+static void mutex_exit(struct mutex **lock)
+{
+	mutex_unlock(*lock);
+}
+
+#define MUTEX_CONTEXT(name, lock)					\
+	struct mutex *mutex_guard_##name __guard(mutex_exit) = &(lock);	\
+	mutex_lock(&(lock));
+
+#define GENERATE_2_CONTEXT_TESTCASE(outer, outer_lock, inner, inner_lock)	\
+										\
+static void __maybe_unused inner##_in_##outer(void)				\
+{										\
+	outer##_CONTEXT(_, outer_lock);						\
+	{									\
+		inner##_CONTEXT(_, inner_lock);					\
+	}									\
+}
+
+/*
+ * wait contexts (considering PREEMPT_RT)
+ *
+ * o: inner is allowed in outer
+ * x: inner is disallowed in outer
+ *
+ *       \  inner |  RCU  | RAW_SPIN | SPIN | MUTEX
+ * outer  \       |       |          |      |
+ * ---------------+-------+----------+------+-------
+ * HARDIRQ        |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * NOTTHREADED_IRQ|   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * SOFTIRQ        |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU            |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_BH         |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_CALLBACK   |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * RCU_SCHED      |   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * RAW_SPIN       |   o   |    o     |  x   |  x
+ * ---------------+-------+----------+------+-------
+ * SPIN           |   o   |    o     |  o   |  x
+ * ---------------+-------+----------+------+-------
+ * MUTEX          |   o   |    o     |  o   |  o
+ * ---------------+-------+----------+------+-------
+ */
+
+#define GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(HARDIRQ, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(NOTTHREADED_HARDIRQ, , inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(SOFTIRQ, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU, , inner, inner_lock)				\
+GENERATE_2_CONTEXT_TESTCASE(RCU_BH, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU_CALLBACK, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RCU_SCHED, , inner, inner_lock)			\
+GENERATE_2_CONTEXT_TESTCASE(RAW_SPINLOCK, raw_lock_A, inner, inner_lock)	\
+GENERATE_2_CONTEXT_TESTCASE(SPINLOCK, lock_A, inner, inner_lock)		\
+GENERATE_2_CONTEXT_TESTCASE(MUTEX, mutex_A, inner, inner_lock)
+
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(RCU, )
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(RAW_SPINLOCK, raw_lock_B)
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(SPINLOCK, lock_B)
+GENERATE_2_CONTEXT_TESTCASE_FOR_ALL_OUTER(MUTEX, mutex_B)
+
+/* the outer context allows all kinds of preemption */
+#define DO_CONTEXT_TESTCASE_OUTER_PREEMPTIBLE(outer)			\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, SUCCESS, LOCKTYPE_MUTEX);		\
+
+/*
+ * the outer context only allows the preemption introduced by spinlock_t (which
+ * is a sleepable lock for PREEMPT_RT)
+ */
+#define DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(outer)		\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, FAILURE, LOCKTYPE_MUTEX);		\
+
+/* the outer doesn't allows any kind of preemption */
+#define DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(outer)			\
+	dotest(RCU_in_##outer, SUCCESS, LOCKTYPE_RWLOCK);		\
+	dotest(RAW_SPINLOCK_in_##outer, SUCCESS, LOCKTYPE_SPIN);	\
+	dotest(SPINLOCK_in_##outer, FAILURE, LOCKTYPE_SPIN);		\
+	dotest(MUTEX_in_##outer, FAILURE, LOCKTYPE_MUTEX);		\
+
+static void wait_context_tests(void)
+{
+	printk("  --------------------------------------------------------------------------\n");
+	printk("  | wait context tests |\n");
+	printk("  --------------------------------------------------------------------------\n");
+	printk("                                 | rcu  | raw  | spin |mutex |\n");
+	printk("  --------------------------------------------------------------------------\n");
+	print_testname("in hardirq context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(HARDIRQ);
+	pr_cont("\n");
+
+	print_testname("in hardirq context (not threaded)");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(NOTTHREADED_HARDIRQ);
+	pr_cont("\n");
+
+	print_testname("in softirq context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(SOFTIRQ);
+	pr_cont("\n");
+
+	print_testname("in RCU context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU);
+	pr_cont("\n");
+
+	print_testname("in RCU-bh context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU_BH);
+	pr_cont("\n");
+
+	print_testname("in RCU callback context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(RCU_CALLBACK);
+	pr_cont("\n");
+
+	print_testname("in RCU-sched context");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(RCU_SCHED);
+	pr_cont("\n");
+
+	print_testname("in RAW_SPINLOCK context");
+	DO_CONTEXT_TESTCASE_OUTER_NOT_PREEMPTIBLE(RAW_SPINLOCK);
+	pr_cont("\n");
+
+	print_testname("in SPINLOCK context");
+	DO_CONTEXT_TESTCASE_OUTER_LIMITED_PREEMPTIBLE(SPINLOCK);
+	pr_cont("\n");
+
+	print_testname("in MUTEX context");
+	DO_CONTEXT_TESTCASE_OUTER_PREEMPTIBLE(MUTEX);
+	pr_cont("\n");
+}
+
 void locking_selftest(void)
 {
 	/*
@@ -2542,6 +2770,10 @@ void locking_selftest(void)
 
 	fs_reclaim_tests();
 
+	/* Wait context test cases that are specific for RAW_LOCK_NESTING */
+	if (IS_ENABLED(CONFIG_PROVE_RAW_LOCK_NESTING))
+		wait_context_tests();
+
 	if (unexpected_testcase_failures) {
 		printk("-----------------------------------------------------------------\n");
 		debug_locks = 0;