kthread: fix possible infinite wait for parking when kthread exits meanwhile

index e8adc10556e0..a001c1ec489b 100644
@@ -77,6 +77,12 @@ static void put_kthread_cb(struct callback_head *work)
 	struct kthread *kthread;
 
 	kthread = container_of(work, struct kthread, put_work);
+	/*
+	 * Kick out possible waiter on a parked completion before
+	 * ref put.  That will force them to spin in a loop once
+	 * more and eventually get the NULL kthread pointer.
+	 */
+	complete(&kthread->parked);
 	put_kthread(kthread);
 }
 

@@ -449,6 +455,11 @@ static void __kthread_unpark(struct task_struct *k, struct kthread *kthread)
 	}
 }
 
+static bool __kthread_isparked(struct kthread *kthread)
+{
+	return test_bit(KTHREAD_IS_PARKED, &kthread->flags);
+}
+
 /**
  * kthread_unpark - unpark a thread created by kthread_create().
  * @k:		thread created by kthread_create().

@@ -479,23 +490,43 @@ EXPORT_SYMBOL_GPL(kthread_unpark);
  *
  * Returns 0 if the thread is parked, -ENOSYS if the thread exited.
  * If called by the kthread itself just the park bit is set.
+ *
+ * BEWARE: The caller is responsible for ensuring the validity of @k when
+ *         calling this function.
+ *
+ * BEWARE: Only one simultaneous caller is possible.  Others will hang
+ *         forever.  You have been warned.
  */
 int kthread_park(struct task_struct *k)
 {
-	struct kthread *kthread = to_live_kthread_and_get(k);
+	struct kthread *kthread;
 	int ret = -ENOSYS;
 
-	if (kthread) {
-		if (!test_bit(KTHREAD_IS_PARKED, &kthread->flags)) {
+	/*
+	 * Here we try to be careful and handle the case, when kthread
+	 * is going to die and will never park.  In that particular case
+	 * put_kthread_cb() is called when kthread->vfork_done is already
+	 * NULL.  put_kthread_cb() does the last completion on kthread->parked,
+	 * thus we will spin once more and next attempt to get an alive
+	 * kthread will fail.
+	 */
+	do {
+		kthread = to_live_kthread_and_get(k);
+		if (!kthread)
+			break;
+		if (!__kthread_isparked(kthread)) {
 			set_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
 			if (k != current) {
 				wake_up_process(k);
 				wait_for_completion(&kthread->parked);
 			}
 		}
+		if (k == current || __kthread_isparked(kthread))
+			/* The way out */
+			ret = 0;
 		put_kthread(kthread);
-		ret = 0;
-	}
+	} while (ret);
+
 	return ret;
 }
 EXPORT_SYMBOL_GPL(kthread_park);