2 files changed, 45 insertions, 0 deletions

diff --git a/include/linux/sched.h b/include/linux/sched.h
index b94f3541f67b..a5b9a83065fa 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1993,8 +1993,10 @@ extern void set_task_comm(struct task_struct *tsk, char *from);
 extern char *get_task_comm(char *to, struct task_struct *tsk);
 
 #ifdef CONFIG_SMP
+extern void wait_task_context_switch(struct task_struct *p);
 extern unsigned long wait_task_inactive(struct task_struct *, long match_state);
 #else
+static inline void wait_task_context_switch(struct task_struct *p) {}
 static inline unsigned long wait_task_inactive(struct task_struct *p,
                                                long match_state)
 {
diff --git a/kernel/sched.c b/kernel/sched.c
index 6cc1fd5d5072..f91bc8141dc3 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -2003,6 +2003,49 @@ migrate_task(struct task_struct *p, int dest_cpu, struct migration_req *req)
 }
 
 /*
+ * wait_task_context_switch -   wait for a thread to complete at least one
+ *                              context switch.
+ *
+ * @p must not be current.
+ */
+void wait_task_context_switch(struct task_struct *p)
+{
+        unsigned long nvcsw, nivcsw, flags;
+        int running;
+        struct rq *rq;
+
+        nvcsw   = p->nvcsw;
+        nivcsw  = p->nivcsw;
+        for (;;) {
+                /*
+                 * The runqueue is assigned before the actual context
+                 * switch. We need to take the runqueue lock.
+                 *
+                 * We could check initially without the lock but it is
+                 * very likely that we need to take the lock in every
+                 * iteration.
+                 */
+                rq = task_rq_lock(p, &flags);
+                running = task_running(rq, p);
+                task_rq_unlock(rq, &flags);
+
+                if (likely(!running))
+                        break;
+                /*
+                 * The switch count is incremented before the actual
+                 * context switch. We thus wait for two switches to be
+                 * sure at least one completed.
+                 */
+                if ((p->nvcsw - nvcsw) > 1)
+                        break;
+                if ((p->nivcsw - nivcsw) > 1)
+                        break;
+
+                cpu_relax();
+        }
+}
+
+/*
  * wait_task_inactive - wait for a thread to unschedule.
  *
  * If @match_state is nonzero, it's the @p->state value just checked and