2 files changed, 14 insertions, 7 deletions

diff --git a/include/linux/kernel.h b/include/linux/kernel.h
index b804f694ae6a..2b5b1e0899a8 100644
--- a/include/linux/kernel.h
+++ b/include/linux/kernel.h
@@ -125,7 +125,7 @@ extern int _cond_resched(void);
 #endif
 
 #ifdef CONFIG_DEBUG_SPINLOCK_SLEEP
-  void __might_sleep(char *file, int line);
+  void __might_sleep(char *file, int line, int preempt_offset);
 /**
  * might_sleep - annotation for functions that can sleep
  *
@@ -137,9 +137,9 @@ extern int _cond_resched(void);
  * supposed to.
  */
 # define might_sleep() \
-        do { __might_sleep(__FILE__, __LINE__); might_resched(); } while (0)
+        do { __might_sleep(__FILE__, __LINE__, 0); might_resched(); } while (0)
 #else
-  static inline void __might_sleep(char *file, int line) { }
+  static inline void __might_sleep(char *file, int line, int preempt_offset) { }
 # define might_sleep() do { might_resched(); } while (0)
 #endif
 
diff --git a/kernel/sched.c b/kernel/sched.c
index 370a6c31c5e1..3ff4d004bd95 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -6610,7 +6610,7 @@ static inline int should_resched(void)
 
 static void __cond_resched(void)
 {
-        __might_sleep(__FILE__, __LINE__);
+        __might_sleep(__FILE__, __LINE__, 0);
 
         add_preempt_count(PREEMPT_ACTIVE);
         schedule();
@@ -9429,13 +9429,20 @@ void __init sched_init(void)
 }
 
 #ifdef CONFIG_DEBUG_SPINLOCK_SLEEP
-void __might_sleep(char *file, int line)
+static inline int preempt_count_equals(int preempt_offset)
+{
+        int nested = preempt_count() & ~PREEMPT_ACTIVE;
+
+        return (nested == PREEMPT_INATOMIC_BASE + preempt_offset);
+}
+
+void __might_sleep(char *file, int line, int preempt_offset)
 {
 #ifdef in_atomic
         static unsigned long prev_jiffy;        /* ratelimiting */
 
-        if ((!in_atomic() && !irqs_disabled()) ||
-                    system_state != SYSTEM_RUNNING || oops_in_progress)
+        if ((preempt_count_equals(preempt_offset) && !irqs_disabled()) ||
+            system_state != SYSTEM_RUNNING || oops_in_progress)
                 return;
         if (time_before(jiffies, prev_jiffy + HZ) && prev_jiffy)
                 return;