2 files changed, 3 insertions, 4 deletions

diff --git a/include/linux/kernel.h b/include/linux/kernel.h
index 5449d2f4a1ef..64ce58bee6f5 100644
--- a/include/linux/kernel.h
+++ b/include/linux/kernel.h
@@ -176,7 +176,7 @@ extern int _cond_resched(void);
  */
 # define might_sleep() \
         do { __might_sleep(__FILE__, __LINE__, 0); might_resched(); } while (0)
-# define sched_annotate_sleep() __set_current_state(TASK_RUNNING)
+# define sched_annotate_sleep() (current->task_state_change = 0)
 #else
   static inline void ___might_sleep(const char *file, int line,
                                    int preempt_offset) { }
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index c0accc00566e..e628cb11b560 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7292,13 +7292,12 @@ void __might_sleep(const char *file, int line, int preempt_offset)
          * since we will exit with TASK_RUNNING make sure we enter with it,
          * otherwise we will destroy state.
          */
-        if (WARN_ONCE(current->state != TASK_RUNNING,
+        WARN_ONCE(current->state != TASK_RUNNING && current->task_state_change,
                         "do not call blocking ops when !TASK_RUNNING; "
                         "state=%lx set at [<%p>] %pS\n",
                         current->state,
                         (void *)current->task_state_change,
-                        (void *)current->task_state_change))
-                __set_current_state(TASK_RUNNING);
+                        (void *)current->task_state_change);
 
         ___might_sleep(file, line, preempt_offset);
 }