The recent rework of mmu lazy mode has resulted in problems when
running as a Xen PV guest. Enabling lazy mmu mode for the new context
during context switch is done from the arch_end_context_switch() hook,
but when calling this hook current hasn't been changed yet, so the
lazy mmu mode state of the wrong task is modified.

Additionally it is much cleaner to use lazy_mmu_mode_pause() and
lazy_mmu_mode_resume() in the Xen context switch hooks, as it avoids
conditionals in those hooks.

In order not having to add another hook to be called after switching
current, modify lazy_mmu_mode_resume() to use a new sub-function which
takes a task pointer as parameter. This new sub-function can then be
used in the xen_end_context_switch() hook.

Fixes: 291b3abed657 ("x86/xen: use lazy_mmu_state when context-switching")
Signed-off-by: Juergen Gross <jgross@suse.com>
---
 arch/x86/xen/enlighten_pv.c |  7 ++-----
 include/linux/pgtable.h     | 33 ++++++++++++++++++++++++---------
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/arch/x86/xen/enlighten_pv.c b/arch/x86/xen/enlighten_pv.c
index ed2d7a3756ce..67bb6bf6d240 100644
--- a/arch/x86/xen/enlighten_pv.c
+++ b/arch/x86/xen/enlighten_pv.c
@@ -424,9 +424,7 @@ static void xen_start_context_switch(struct task_struct *prev)
 {
 	BUG_ON(preemptible());
 
-	if (this_cpu_read(xen_lazy_mode) == XEN_LAZY_MMU) {
-		arch_leave_lazy_mmu_mode();
-	}
+	lazy_mmu_mode_pause();
 	enter_lazy(XEN_LAZY_CPU);
 }
 
@@ -436,8 +434,7 @@ static void xen_end_context_switch(struct task_struct *next)
 
 	xen_mc_flush();
 	leave_lazy(XEN_LAZY_CPU);
-	if (__task_lazy_mmu_mode_active(next))
-		arch_enter_lazy_mmu_mode();
+	lazy_mmu_mode_resume_task(next);
 }
 
 static unsigned long xen_store_tr(void)
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index cdd68ed3ae1a..83a099bf2038 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -326,6 +326,28 @@ static inline void lazy_mmu_mode_pause(void)
 		arch_leave_lazy_mmu_mode();
 }
 
+/**
+ * lazy_mmu_mode_resume_task() - Resume the lazy MMU mode for a specific task.
+ *
+ * Like lazy_mmu_mode_resume() below, but with a task specified.
+ * Must be called only by lazy_mmu_mode_resume() or during context switch.
+ * Must never be called in interrupt context.
+ *
+ * Must match a call to lazy_mmu_mode_pause().
+ *
+ * Has no effect if called:
+ * - While paused (inside another pause()/resume() pair)
+ */
+static inline void lazy_mmu_mode_resume_task(struct task_struct *task)
+{
+	struct lazy_mmu_state *state = &task->lazy_mmu_state;
+
+	VM_WARN_ON_ONCE(state->pause_count == 0);
+
+	if (--state->pause_count == 0 && state->enable_count > 0)
+		arch_enter_lazy_mmu_mode();
+}
+
 /**
  * lazy_mmu_mode_resume() - Resume the lazy MMU mode.
  *
@@ -341,15 +363,8 @@ static inline void lazy_mmu_mode_pause(void)
  */
 static inline void lazy_mmu_mode_resume(void)
 {
-	struct lazy_mmu_state *state = &current->lazy_mmu_state;
-
-	if (in_interrupt())
-		return;
-
-	VM_WARN_ON_ONCE(state->pause_count == 0);
-
-	if (--state->pause_count == 0 && state->enable_count > 0)
-		arch_enter_lazy_mmu_mode();
+	if (!in_interrupt())
+		lazy_mmu_mode_resume_task(current);
 }
 #else
 static inline void lazy_mmu_mode_enable(void) {}
-- 
2.54.0