4 files changed, 36 insertions, 28 deletions

diff --git a/include/linux/kthread.h b/include/linux/kthread.h
index 2803264c512f..c1961761311d 100644
--- a/include/linux/kthread.h
+++ b/include/linux/kthread.h
@@ -62,7 +62,6 @@ void *kthread_probe_data(struct task_struct *k);
 int kthread_park(struct task_struct *k);
 void kthread_unpark(struct task_struct *k);
 void kthread_parkme(void);
-void kthread_park_complete(struct task_struct *k);
 
 int kthreadd(void *unused);
 extern struct task_struct *kthreadd_task;
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 9256118bd40c..43731fe51c97 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -118,7 +118,7 @@ struct task_group;
  * the comment with set_special_state().
  */
 #define is_special_task_state(state)                            \
-        ((state) & (__TASK_STOPPED | __TASK_TRACED | TASK_DEAD))
+        ((state) & (__TASK_STOPPED | __TASK_TRACED | TASK_PARKED | TASK_DEAD))
 
 #define __set_current_state(state_value)                        \
         do {                                                    \
diff --git a/kernel/kthread.c b/kernel/kthread.c
index 481951bf091d..750cb8082694 100644
--- a/kernel/kthread.c
+++ b/kernel/kthread.c
@@ -177,9 +177,20 @@ void *kthread_probe_data(struct task_struct *task)
 static void __kthread_parkme(struct kthread *self)
 {
         for (;;) {
-                set_current_state(TASK_PARKED);
+                /*
+                 * TASK_PARKED is a special state; we must serialize against
+                 * possible pending wakeups to avoid store-store collisions on
+                 * task->state.
+                 *
+                 * Such a collision might possibly result in the task state
+                 * changin from TASK_PARKED and us failing the
+                 * wait_task_inactive() in kthread_park().
+                 */
+                set_special_state(TASK_PARKED);
                 if (!test_bit(KTHREAD_SHOULD_PARK, &self->flags))
                         break;
+
+                complete_all(&self->parked);
                 schedule();
         }
         __set_current_state(TASK_RUNNING);
@@ -191,11 +202,6 @@ void kthread_parkme(void)
 }
 EXPORT_SYMBOL_GPL(kthread_parkme);
 
-void kthread_park_complete(struct task_struct *k)
-{
-        complete_all(&to_kthread(k)->parked);
-}
-
 static int kthread(void *_create)
 {
         /* Copy data: it's on kthread's stack */
@@ -461,6 +467,9 @@ void kthread_unpark(struct task_struct *k)
 
         reinit_completion(&kthread->parked);
         clear_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
+        /*
+         * __kthread_parkme() will either see !SHOULD_PARK or get the wakeup.
+         */
         wake_up_state(k, TASK_PARKED);
 }
 EXPORT_SYMBOL_GPL(kthread_unpark);
@@ -487,7 +496,16 @@ int kthread_park(struct task_struct *k)
         set_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
         if (k != current) {
                 wake_up_process(k);
+                /*
+                 * Wait for __kthread_parkme() to complete(), this means we
+                 * _will_ have TASK_PARKED and are about to call schedule().
+                 */
                 wait_for_completion(&kthread->parked);
+                /*
+                 * Now wait for that schedule() to complete and the task to
+                 * get scheduled out.
+                 */
+                WARN_ON_ONCE(!wait_task_inactive(k, TASK_PARKED));
         }
 
         return 0;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 22fce36426c0..fe365c9a08e9 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7,7 +7,6 @@
  */
 #include "sched.h"
 
-#include <linux/kthread.h>
 #include <linux/nospec.h>
 
 #include <linux/kcov.h>
@@ -2724,28 +2723,20 @@ static struct rq *finish_task_switch(struct task_struct *prev)
                 membarrier_mm_sync_core_before_usermode(mm);
                 mmdrop(mm);
         }
-        if (unlikely(prev_state & (TASK_DEAD|TASK_PARKED))) {
-                switch (prev_state) {
-                case TASK_DEAD:
-                        if (prev->sched_class->task_dead)
-                                prev->sched_class->task_dead(prev);
+        if (unlikely(prev_state == TASK_DEAD)) {
+                if (prev->sched_class->task_dead)
+                        prev->sched_class->task_dead(prev);
 
-                        /*
-                         * Remove function-return probe instances associated with this
-                         * task and put them back on the free list.
-                         */
-                        kprobe_flush_task(prev);
-
-                        /* Task is done with its stack. */
-                        put_task_stack(prev);
+                /*
+                 * Remove function-return probe instances associated with this
+                 * task and put them back on the free list.
+                 */
+                kprobe_flush_task(prev);
 
-                        put_task_struct(prev);
-                        break;
+                /* Task is done with its stack. */
+                put_task_stack(prev);
 
-                case TASK_PARKED:
-                        kthread_park_complete(prev);
-                        break;
-                }
+                put_task_struct(prev);
         }
 
         tick_nohz_task_switch();