From 0b6b1c285920563ba0c119d3190ac25af4731d02 Mon Sep 17 00:00:00 2001
From: Konstantin Belousov <kib@FreeBSD.org>
Date: Wed, 28 Aug 2024 00:33:38 +0300
Subject: [PATCH] Add rangelock_may_recurse(9)

Reviewed by:	markj
Tested by:	lwhsu
Sponsored by:	The FreeBSD Foundation
Differential revision:	https://reviews.freebsd.org/D46465
---
 sys/kern/kern_rangelock.c | 41 +++++++++++++++++++++++++++++++++++++++
 sys/sys/rangelock.h       |  1 +
 2 files changed, 42 insertions(+)

diff --git a/sys/kern/kern_rangelock.c b/sys/kern/kern_rangelock.c
index 0e62b91b4ee7b6..4d74c02302e7d3 100644
--- a/sys/kern/kern_rangelock.c
+++ b/sys/kern/kern_rangelock.c
@@ -752,6 +752,47 @@ rangelock_trywlock(struct rangelock *lock, vm_ooffset_t start, vm_ooffset_t end)
 	return (rangelock_lock_int(lock, true, start, end, RL_LOCK_WRITE));
 }
 
+/*
+ * If the caller asserts that it can obtain the range locks on the
+ * same lock simultaneously, switch to the non-cheat mode.  Cheat mode
+ * cannot handle it, hanging in drain or trylock retries.
+ */
+void
+rangelock_may_recurse(struct rangelock *lock)
+{
+	uintptr_t v, x;
+
+	v = atomic_load_ptr(&lock->head);
+	if ((v & RL_CHEAT_CHEATING) == 0)
+		return;
+
+	sleepq_lock(&lock->head);
+	for (;;) {
+		if ((v & RL_CHEAT_CHEATING) == 0) {
+			sleepq_release(&lock->head);
+			return;
+		}
+
+		/* Cheating and locked, drain. */
+		if ((v & RL_CHEAT_WLOCKED) != 0 ||
+		    (v & ~RL_CHEAT_MASK) >= RL_CHEAT_READER) {
+			x = v | RL_CHEAT_DRAINING;
+			if (atomic_fcmpset_ptr(&lock->head, &v, x) != 0) {
+				rangelock_cheat_drain(lock);
+				return;
+			}
+			continue;
+		}
+
+		/* Cheating and unlocked, clear RL_CHEAT_CHEATING. */
+		x = 0;
+		if (atomic_fcmpset_ptr(&lock->head, &v, x) != 0) {
+			sleepq_release(&lock->head);
+			return;
+		}
+	}
+}
+
 #ifdef INVARIANT_SUPPORT
 void
 _rangelock_cookie_assert(void *cookie, int what, const char *file, int line)
diff --git a/sys/sys/rangelock.h b/sys/sys/rangelock.h
index 127f101ddc2e7b..accf33d7296b7d 100644
--- a/sys/sys/rangelock.h
+++ b/sys/sys/rangelock.h
@@ -66,6 +66,7 @@ void	*rangelock_wlock(struct rangelock *lock, vm_ooffset_t start,
 void	*rangelock_trywlock(struct rangelock *lock, vm_ooffset_t start,
     vm_ooffset_t end);
 void	 rangelock_entry_free(struct rl_q_entry *e);
+void	rangelock_may_recurse(struct rangelock *lock);
 #if defined(INVARIANTS) || defined(INVARIANT_SUPPORT)
 void	_rangelock_cookie_assert(void *cookie, int what, const char *file,
     int line);