[RFC PATCH v2 03/13] sched: Split preempt-count bookkeeping from raw updates

[Karl Mehltretter <kmehltretter@gmail.com>]

Very early hardirq entry code needs to update the preempt count before
instrumentable debug and tracing code can run. Conversely, hardirq return
code needs to finish instrumentable validation and latency tracing before
raw teardown.

Add a small hardirq-entry header with irqentry_hardirq_count_start() to
stage the raw HARDIRQ_OFFSET addition and record whether it was applied in
an opaque cookie. With DEBUG_PREEMPT and without separate reschedule bits,
route an already-negative count through a noinstr slow helper which
preserves the existing DEBUG_LOCKS_WARN_ON() oops short circuit. Add
irqentry_hardirq_count_finish() to report a rejected addition or complete
post-add debug checks and latency tracing, and add
irqentry_hardirq_count_sub_prepare() for validation and tracing before a
raw subtraction. Keep preempt_count_add() and preempt_count_sub() behavior
and ordering unchanged.

Keep latency tracing in each calling path so CALLER_ADDR0 identifies the
caller. Force the start and stop helpers inline because trace provenance
depends on their stack frames.

Assisted-by: LLM
---
 include/linux/hardirq.h          |   1 +
 include/linux/irqentry_hardirq.h |  81 +++++++++++++++++++++++++
 kernel/sched/core.c              | 101 +++++++++++++++++++++++++++----
 3 files changed, 171 insertions(+), 12 deletions(-)
 create mode 100644 include/linux/irqentry_hardirq.h

diff --git a/include/linux/hardirq.h b/include/linux/hardirq.h
index 2648c49cd8f6c..dcbda3c1d5636 100644
--- a/include/linux/hardirq.h
+++ b/include/linux/hardirq.h
@@ -4,6 +4,7 @@
 
 #include <linux/context_tracking_state.h>
 #include <linux/preempt.h>
+#include <linux/irqentry_hardirq.h>
 #include <linux/lockdep.h>
 #include <linux/ftrace_irq.h>
 #include <linux/sched.h>
diff --git a/include/linux/irqentry_hardirq.h b/include/linux/irqentry_hardirq.h
new file mode 100644
index 0000000000000..1937fc972e846
--- /dev/null
+++ b/include/linux/irqentry_hardirq.h
@@ -0,0 +1,81 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __LINUX_IRQENTRY_HARDIRQ_H
+#define __LINUX_IRQENTRY_HARDIRQ_H
+
+#include <linux/preempt.h>
+
+struct irqentry_hardirq_cookie {
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+	bool applied;
+#endif
+};
+
+static __always_inline bool
+irqentry_hardirq_count_applied(struct irqentry_hardirq_cookie cookie)
+{
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+	return cookie.applied;
+#else
+	return true;
+#endif
+}
+
+static __always_inline bool irqentry_hardirq_count_valid(void)
+{
+#ifdef CONFIG_DEBUG_PREEMPT
+	/*
+	 * Without separate reschedule bits, preempt_count() excludes any
+	 * shared reschedule bit. A negative result therefore denotes an
+	 * existing underflow.
+	 */
+	if (!IS_ENABLED(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS) &&
+	    preempt_count() < 0)
+		return false;
+#endif
+	return true;
+}
+
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+struct irqentry_hardirq_cookie noinstr
+__irqentry_hardirq_count_start_slow(void);
+#endif
+
+/*
+ * Attempt the non-instrumentable part of a staged hardirq-count addition and
+ * record whether it was applied. The returned cookie must be passed to the
+ * matching irqentry_hardirq_count_finish() and subtraction preparation. A
+ * rejected cookie attempts deferred reporting and suppresses post-add
+ * bookkeeping and the matching subtraction regardless of warning
+ * suppression.
+ */
+static __always_inline struct irqentry_hardirq_cookie
+irqentry_hardirq_count_start(void)
+{
+	struct irqentry_hardirq_cookie cookie = { };
+
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+	cookie.applied = true;
+	/* DEBUG_LOCKS_WARN_ON() permits the addition during an oops. */
+	if (unlikely(!irqentry_hardirq_count_valid()))
+		return __irqentry_hardirq_count_start_slow();
+#endif
+	__preempt_count_add(HARDIRQ_OFFSET);
+
+	return cookie;
+}
+
+#if defined(CONFIG_PREEMPTION) && (defined(CONFIG_DEBUG_PREEMPT) || \
+					defined(CONFIG_TRACE_PREEMPT_TOGGLE))
+void irqentry_hardirq_count_finish(struct irqentry_hardirq_cookie cookie);
+bool irqentry_hardirq_count_sub_prepare(void);
+#else
+static __always_inline void
+irqentry_hardirq_count_finish(struct irqentry_hardirq_cookie cookie) { }
+static __always_inline bool irqentry_hardirq_count_sub_prepare(void) { return true; }
+#endif
+
+#endif /* __LINUX_IRQENTRY_HARDIRQ_H */
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 04b62c16fa8ca..7e2da8fc32b6d 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6028,7 +6028,7 @@ static inline void sched_tick_stop(int cpu) { }
  * If the value passed in is equal to the current preempt count
  * then we just disabled preemption. Start timing the latency.
  */
-static inline void preempt_latency_start(int val)
+static __always_inline void preempt_latency_start(int val)
 {
 	if (preempt_count() == val) {
 		unsigned long ip = get_lock_parent_ip();
@@ -6039,6 +6039,62 @@ static inline void preempt_latency_start(int val)
 	}
 }
 
+static __always_inline void preempt_count_add_debug_checks(void)
+{
+#ifdef CONFIG_DEBUG_PREEMPT
+	DEBUG_LOCKS_WARN_ON((preempt_count() & PREEMPT_MASK) >=
+				PREEMPT_MASK - 10);
+#endif
+}
+
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+/*
+ * Preserve the DEBUG_LOCKS_WARN_ON() add-or-reject decision without invoking
+ * instrumented reporting from noinstr entry code. The caller reaches this
+ * only after observing an already-negative count. Recheck after testing oops
+ * state to retain the original short-circuit order.
+ */
+struct irqentry_hardirq_cookie noinstr
+__irqentry_hardirq_count_start_slow(void)
+{
+	struct irqentry_hardirq_cookie cookie = { };
+
+	if (!READ_ONCE(oops_in_progress) && !irqentry_hardirq_count_valid())
+		return cookie;
+
+	__preempt_count_add(HARDIRQ_OFFSET);
+	cookie.applied = true;
+	return cookie;
+}
+NOKPROBE_SYMBOL(__irqentry_hardirq_count_start_slow);
+#endif
+
+/**
+ * irqentry_hardirq_count_finish - Finish a staged hardirq-count addition
+ * @cookie: Result of the matching irqentry_hardirq_count_start()
+ *
+ * Attempt to report a rejected addition and skip its post-add work, or
+ * complete the post-add debug checks and latency tracing for an applied
+ * addition, after an entry path had to stage the update before instrumentable
+ * code could run.
+ */
+void irqentry_hardirq_count_finish(struct irqentry_hardirq_cookie cookie)
+{
+#if defined(CONFIG_DEBUG_PREEMPT) && \
+	!defined(CONFIG_HAS_SEPARATE_PREEMPT_RESCHED_BITS)
+	bool applied = irqentry_hardirq_count_applied(cookie);
+
+	if (!applied) {
+		DEBUG_LOCKS_WARN_ON(!applied);
+		return;
+	}
+#endif
+	preempt_count_add_debug_checks();
+	preempt_latency_start(HARDIRQ_OFFSET);
+}
+NOKPROBE_SYMBOL(irqentry_hardirq_count_finish);
+
 void preempt_count_add(int val)
 {
 #ifdef CONFIG_DEBUG_PREEMPT
@@ -6054,13 +6110,8 @@ void preempt_count_add(int val)
 		return;
 #endif
 	__preempt_count_add(val);
-#ifdef CONFIG_DEBUG_PREEMPT
-	/*
-	 * Spinlock count overflowing soon?
-	 */
-	DEBUG_LOCKS_WARN_ON((preempt_count() & PREEMPT_MASK) >=
-				PREEMPT_MASK - 10);
-#endif
+	preempt_count_add_debug_checks();
+	/* Keep latency attribution in the preempt_count_add() stack frame. */
 	preempt_latency_start(val);
 }
 EXPORT_SYMBOL(preempt_count_add);
@@ -6070,13 +6121,13 @@ NOKPROBE_SYMBOL(preempt_count_add);
  * If the value passed in equals to the current preempt count
  * then we just enabled preemption. Stop timing the latency.
  */
-static inline void preempt_latency_stop(int val)
+static __always_inline void preempt_latency_stop(int val)
 {
 	if (preempt_count() == val)
 		trace_preempt_on(CALLER_ADDR0, get_lock_parent_ip());
 }
 
-void preempt_count_sub(int val)
+static __always_inline bool preempt_count_sub_debug_checks(int val)
 {
 #ifdef CONFIG_DEBUG_PREEMPT
 	/*
@@ -6086,15 +6137,41 @@ void preempt_count_sub(int val)
 	unsigned int pc = preempt_count();
 
 	if (DEBUG_LOCKS_WARN_ON(pc - uval > pc))
-		return;
+		return false;
 	/*
 	 * Is the spinlock portion underflowing?
 	 */
 	if (DEBUG_LOCKS_WARN_ON((val < PREEMPT_MASK) &&
 			!(preempt_count() & PREEMPT_MASK)))
-		return;
+		return false;
 #endif
+	return true;
+}
+
+/**
+ * irqentry_hardirq_count_sub_prepare - Prepare a raw hardirq-count subtraction
+ *
+ * Run debug checks and latency tracing before an entry path performs a raw
+ * preempt count subtraction after all instrumentable code has returned.
+ *
+ * Return: true when the caller should perform the subtraction.
+ */
+bool irqentry_hardirq_count_sub_prepare(void)
+{
+	if (!preempt_count_sub_debug_checks(HARDIRQ_OFFSET))
+		return false;
+
+	preempt_latency_stop(HARDIRQ_OFFSET);
+	return true;
+}
+NOKPROBE_SYMBOL(irqentry_hardirq_count_sub_prepare);
+
+void preempt_count_sub(int val)
+{
+	if (!preempt_count_sub_debug_checks(val))
+		return;
 
+	/* Keep latency attribution in the preempt_count_sub() stack frame. */
 	preempt_latency_stop(val);
 	__preempt_count_sub(val);
 }
-- 
2.39.5 (Apple Git-154)