[PATCH v3 07/13] mm: enable lazy_mmu sections to nest

[Kevin Brodsky <kevin.brodsky@xxxxxxx>]

Despite recent efforts to prevent lazy_mmu sections from nesting, it
remains difficult to ensure that it never occurs - and in fact it
does occur on arm64 in certain situations (CONFIG_DEBUG_PAGEALLOC).
Commit 1ef3095b1405 ("arm64/mm: Permit lazy_mmu_mode to be nested")
made nesting tolerable on arm64, but without truly supporting it:
the inner call to leave() disables the batching optimisation before
the outer section ends.

This patch actually enables lazy_mmu sections to nest by tracking
the nesting level in task_struct, in a similar fashion to e.g.
pagefault_{enable,disable}(). This is fully handled by the generic
lazy_mmu helpers that were recently introduced.

lazy_mmu sections were not initially intended to nest, so we need to
clarify the semantics w.r.t. the arch_*_lazy_mmu_mode() callbacks.
This patch takes the following approach:

* The outermost calls to lazy_mmu_mode_{enable,disable}() trigger
  calls to arch_{enter,leave}_lazy_mmu_mode() - this is unchanged.

* Nested calls to lazy_mmu_mode_{enable,disable}() are not forwarded
  to the arch via arch_{enter,leave} - lazy MMU remains enabled so
  the assumption is that these callbacks are not relevant. However,
  existing code may rely on a call to disable() to flush any batched
  state, regardless of nesting. arch_flush_lazy_mmu_mode() is
  therefore called in that situation.

A separate interface was recently introduced to temporarily pause
the lazy MMU mode: lazy_mmu_mode_{pause,resume}(). pause() fully
exits the mode *regardless of the nesting level*, and resume()
restores the mode at the same nesting level.

Whether the mode is actually enabled or not at any point is tracked
by a separate "enabled" field in task_struct; this makes it possible
to check invariants in the generic API, and to expose a new
in_lazy_mmu_mode() helper to replace the various ways arch's
currently track whether the mode is enabled (this will be done in
later patches).

In summary (count/enabled represent the values *after* the call):

lazy_mmu_mode_enable()          -> arch_enter()     count=1 enabled=1
    lazy_mmu_mode_enable()      -> ø                count=2 enabled=1
        lazy_mmu_mode_pause()   -> arch_leave()     count=2 enabled=0
        lazy_mmu_mode_resume()  -> arch_enter()     count=2 enabled=1
    lazy_mmu_mode_disable()     -> arch_flush()     count=1 enabled=1
lazy_mmu_mode_disable()         -> arch_leave()     count=0 enabled=0

Note: in_lazy_mmu_mode() is added to <linux/sched.h> to allow arch
headers included by <linux/pgtable.h> to use it.

Signed-off-by: Kevin Brodsky <kevin.brodsky@xxxxxxx>
---
Alexander Gordeev suggested that a future optimisation may need
lazy_mmu_mode_{pause,resume}() to call distinct arch callbacks [1]. For
now arch_{leave,enter}() are called directly, but introducing new arch
callbacks should be straightforward.

[1] 
https://lore.kernel.org/all/5a0818bb-75d4-47df-925c-0102f7d598f4-agordeev@xxxxxxxxxxxxx/
---
 arch/arm64/include/asm/pgtable.h | 12 ------
 include/linux/mm_types_task.h    |  5 +++
 include/linux/pgtable.h          | 69 ++++++++++++++++++++++++++++++--
 include/linux/sched.h            | 16 ++++++++
 4 files changed, 86 insertions(+), 16 deletions(-)

diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index e3cbb10288c4..f15ca4d62f09 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -82,18 +82,6 @@ static inline void queue_pte_barriers(void)
 
 static inline void arch_enter_lazy_mmu_mode(void)
 {
-       /*
-        * lazy_mmu_mode is not supposed to permit nesting. But in practice this
-        * does happen with CONFIG_DEBUG_PAGEALLOC, where a page allocation
-        * inside a lazy_mmu_mode section (such as zap_pte_range()) will change
-        * permissions on the linear map with apply_to_page_range(), which
-        * re-enters lazy_mmu_mode. So we tolerate nesting in our
-        * implementation. The first call to arch_leave_lazy_mmu_mode() will
-        * flush and clear the flag such that the remainder of the work in the
-        * outer nest behaves as if outside of lazy mmu mode. This is safe and
-        * keeps tracking simple.
-        */
-
        if (in_interrupt())
                return;
 
diff --git a/include/linux/mm_types_task.h b/include/linux/mm_types_task.h
index a82aa80c0ba4..2ff83b85fef0 100644
--- a/include/linux/mm_types_task.h
+++ b/include/linux/mm_types_task.h
@@ -88,4 +88,9 @@ struct tlbflush_unmap_batch {
 #endif
 };
 
+struct lazy_mmu_state {
+       u8 count;
+       bool enabled;
+};
+
 #endif /* _LINUX_MM_TYPES_TASK_H */
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index 194b2c3e7576..269225a733de 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -228,28 +228,89 @@ static inline int pmd_dirty(pmd_t pmd)
  * of the lazy mode. So the implementation must assume preemption may be 
enabled
  * and cpu migration is possible; it must take steps to be robust against this.
  * (In practice, for user PTE updates, the appropriate page table lock(s) are
- * held, but for kernel PTE updates, no lock is held). Nesting is not permitted
- * and the mode cannot be used in interrupt context.
+ * held, but for kernel PTE updates, no lock is held). The mode cannot be used
+ * in interrupt context.
+ *
+ * The lazy MMU mode is enabled for a given block of code using:
+ *
+ *   lazy_mmu_mode_enable();
+ *   <code>
+ *   lazy_mmu_mode_disable();
+ *
+ * Nesting is permitted: <code> may itself use an enable()/disable() pair.
+ * A nested call to enable() has no functional effect; however disable() causes
+ * any batched architectural state to be flushed regardless of nesting. After a
+ * call to disable(), the caller can therefore rely on all previous page table
+ * modifications to have taken effect, but the lazy MMU mode may still be
+ * enabled.
+ *
+ * In certain cases, it may be desirable to temporarily pause the lazy MMU 
mode.
+ * This can be done using:
+ *
+ *   lazy_mmu_mode_pause();
+ *   <code>
+ *   lazy_mmu_mode_resume();
+ *
+ * This sequence must only be used if the lazy MMU mode is already enabled.
+ * pause() ensures that the mode is exited regardless of the nesting level;
+ * resume() re-enters the mode at the same nesting level. <code> must not 
modify
+ * the lazy MMU state (i.e. it must not call any of the lazy_mmu_mode_*
+ * helpers).
+ *
+ * in_lazy_mmu_mode() can be used to check whether the lazy MMU mode is
+ * currently enabled.
  */
 #ifdef CONFIG_ARCH_LAZY_MMU
 static inline void lazy_mmu_mode_enable(void)
 {
-       arch_enter_lazy_mmu_mode();
+       struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+       VM_BUG_ON(state->count == U8_MAX);
+       /* enable() must not be called while paused */
+       VM_WARN_ON(state->count > 0 && !state->enabled);
+
+       if (state->count == 0) {
+               arch_enter_lazy_mmu_mode();
+               state->enabled = true;
+       }
+       ++state->count;
 }
 
 static inline void lazy_mmu_mode_disable(void)
 {
-       arch_leave_lazy_mmu_mode();
+       struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+       VM_BUG_ON(state->count == 0);
+       VM_WARN_ON(!state->enabled);
+
+       --state->count;
+       if (state->count == 0) {
+               state->enabled = false;
+               arch_leave_lazy_mmu_mode();
+       } else {
+               /* Exiting a nested section */
+               arch_flush_lazy_mmu_mode();
+       }
 }
 
 static inline void lazy_mmu_mode_pause(void)
 {
+       struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+       VM_WARN_ON(state->count == 0 || !state->enabled);
+
+       state->enabled = false;
        arch_leave_lazy_mmu_mode();
 }
 
 static inline void lazy_mmu_mode_resume(void)
 {
+       struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+       VM_WARN_ON(state->count == 0 || state->enabled);
+
        arch_enter_lazy_mmu_mode();
+       state->enabled = true;
 }
 #else
 static inline void lazy_mmu_mode_enable(void) {}
diff --git a/include/linux/sched.h b/include/linux/sched.h
index cbb7340c5866..2862d8bf2160 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1441,6 +1441,10 @@ struct task_struct {
 
        struct page_frag                task_frag;
 
+#ifdef CONFIG_ARCH_LAZY_MMU
+       struct lazy_mmu_state           lazy_mmu_state;
+#endif
+
 #ifdef CONFIG_TASK_DELAY_ACCT
        struct task_delay_info          *delays;
 #endif
@@ -1724,6 +1728,18 @@ static inline char task_state_to_char(struct task_struct 
*tsk)
        return task_index_to_char(task_state_index(tsk));
 }
 
+#ifdef CONFIG_ARCH_LAZY_MMU
+static inline bool in_lazy_mmu_mode(void)
+{
+       return current->lazy_mmu_state.enabled;
+}
+#else
+static inline bool in_lazy_mmu_mode(void)
+{
+       return false;
+}
+#endif
+
 extern struct pid *cad_pid;
 
 /*
-- 
2.47.0