[RFC PATCH v4 1/3] mm: allow page faults to request VMA-lock retry

From: Hongru Zhang <[email protected]>

Page faults handled under the per-VMA lock currently fall back to the
mmap_lock path whenever handle_mm_fault() returns VM_FAULT_RETRY. This
means that lower-level fault handlers have no way to tell the
architecture fault handler that the retry can safely continue under the
per-VMA lock.

Add VM_FAULT_MAY_USE_VMA_LOCK as an advisory bit that can be returned
together with VM_FAULT_RETRY. Architecture fault handlers use this bit
to allow at most one retry under the per-VMA lock.

This preserves the existing mmap_lock fallback behaviour for fault
handlers that continue to return VM_FAULT_RETRY without
VM_FAULT_MAY_USE_VMA_LOCK: major retries still enter the mmap_lock path
with FAULT_FLAG_TRIED set, while minor retries still enter it as a fresh
first attempt.

The difference is limited to fault handlers that return VM_FAULT_RETRY
with VM_FAULT_MAY_USE_VMA_LOCK. For them, both major and minor retries
enter the VMA-lock retry with FAULT_FLAG_TRIED set. For major faults
this follows the existing mmap_lock retry handling, but the retried
fault runs under the VMA lock rather than the mmap_lock. For minor
faults this replaces the fresh mmap_lock retry with a VMA-lock retry
that has FAULT_FLAG_TRIED set. This can avoid an extra retry round for
short waits. The cost is that a retried fault that blocks for an extended
period may wait while holding the VMA lock. Fault handlers should return
VM_FAULT_RETRY with VM_FAULT_MAY_USE_VMA_LOCK when that tradeoff is
preferable to falling back to mmap_lock immediately.

No current code sets VM_FAULT_MAY_USE_VMA_LOCK yet; this patch only
prepares the retry plumbing for later users.

No functional change is intended.

Signed-off-by: Hongru Zhang <[email protected]>
Suggested-by: Barry Song <[email protected]>
Suggested-by: Suren Baghdasaryan <[email protected]>
---
 arch/arm/mm/fault.c       |  6 ++++--
 arch/arm64/mm/fault.c     |  7 +++++--
 arch/loongarch/mm/fault.c |  6 ++++--
 arch/powerpc/mm/fault.c   |  6 ++++--
 arch/riscv/mm/fault.c     |  6 ++++--
 arch/s390/mm/fault.c      |  5 +++--
 arch/x86/mm/fault.c       |  6 ++++--
 include/linux/mm.h        | 28 ++++++++++++++++++++++++++++
 include/linux/mm_types.h  |  4 ++++
 9 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/arch/arm/mm/fault.c b/arch/arm/mm/fault.c
index e62cc4be5adf..158923b70901 100644
--- a/arch/arm/mm/fault.c
+++ b/arch/arm/mm/fault.c
@@ -391,6 +391,7 @@ do_page_fault(unsigned long addr, unsigned int fsr, struct pt_regs *regs)
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, addr);
 	if (!vma)
 		goto lock_mmap;
@@ -411,8 +412,6 @@ do_page_fault(unsigned long addr, unsigned int fsr, struct pt_regs *regs)
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 
 	/* Quick path to respond to signals */
 	if (fault_signal_pending(fault, regs)) {
@@ -420,6 +419,9 @@ do_page_fault(unsigned long addr, unsigned int fsr, struct pt_regs *regs)
 			goto no_context;
 		return 0;
 	}
+
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
 lock_mmap:
 
 retry:
diff --git a/arch/arm64/mm/fault.c b/arch/arm64/mm/fault.c
index 0b52557652be..b17986b40ac3 100644
--- a/arch/arm64/mm/fault.c
+++ b/arch/arm64/mm/fault.c
@@ -678,6 +678,7 @@ static int __kprobes do_page_fault(unsigned long far, unsigned long esr,
 	if (!(mm_flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, addr);
 	if (!vma)
 		goto lock_mmap;
@@ -715,8 +716,6 @@ static int __kprobes do_page_fault(unsigned long far, unsigned long esr,
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		mm_flags |= FAULT_FLAG_TRIED;
 
 	/* Quick path to respond to signals */
 	if (fault_signal_pending(fault, regs)) {
@@ -724,6 +723,10 @@ static int __kprobes do_page_fault(unsigned long far, unsigned long esr,
 			goto no_context;
 		return 0;
 	}
+
+	if (fault_should_retry_under_vma_lock(fault, &mm_flags))
+		goto retry_vma;
+
 lock_mmap:
 
 retry:
diff --git a/arch/loongarch/mm/fault.c b/arch/loongarch/mm/fault.c
index 2c93d33356e5..6a946838b54b 100644
--- a/arch/loongarch/mm/fault.c
+++ b/arch/loongarch/mm/fault.c
@@ -219,6 +219,7 @@ static void __kprobes __do_page_fault(struct pt_regs *regs,
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, address);
 	if (!vma)
 		goto lock_mmap;
@@ -256,8 +257,6 @@ static void __kprobes __do_page_fault(struct pt_regs *regs,
 	}
 
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 
 	/* Quick path to respond to signals */
 	if (fault_signal_pending(fault, regs)) {
@@ -265,6 +264,9 @@ static void __kprobes __do_page_fault(struct pt_regs *regs,
 			no_context(regs, write, address);
 		return;
 	}
+
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
 lock_mmap:
 
 retry:
diff --git a/arch/powerpc/mm/fault.c b/arch/powerpc/mm/fault.c
index 806c74e0d5ab..e2a128fba408 100644
--- a/arch/powerpc/mm/fault.c
+++ b/arch/powerpc/mm/fault.c
@@ -487,6 +487,7 @@ static int ___do_page_fault(struct pt_regs *regs, unsigned long address,
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, address);
 	if (!vma)
 		goto lock_mmap;
@@ -511,12 +512,13 @@ static int ___do_page_fault(struct pt_regs *regs, unsigned long address,
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 
 	if (fault_signal_pending(fault, regs))
 		return user_mode(regs) ? 0 : SIGBUS;
 
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
+
 lock_mmap:
 
 	/* When running in the kernel we expect faults to occur only to
diff --git a/arch/riscv/mm/fault.c b/arch/riscv/mm/fault.c
index 04ed6f8acae4..87b061feba51 100644
--- a/arch/riscv/mm/fault.c
+++ b/arch/riscv/mm/fault.c
@@ -347,6 +347,7 @@ void handle_page_fault(struct pt_regs *regs)
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, addr);
 	if (!vma)
 		goto lock_mmap;
@@ -368,14 +369,15 @@ void handle_page_fault(struct pt_regs *regs)
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 
 	if (fault_signal_pending(fault, regs)) {
 		if (!user_mode(regs))
 			no_context(regs, addr);
 		return;
 	}
+
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
 lock_mmap:
 
 retry:
diff --git a/arch/s390/mm/fault.c b/arch/s390/mm/fault.c
index 028aeb9c48d6..8e90e522436b 100644
--- a/arch/s390/mm/fault.c
+++ b/arch/s390/mm/fault.c
@@ -294,6 +294,7 @@ static void do_exception(struct pt_regs *regs, int access)
 		flags |= FAULT_FLAG_WRITE;
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
+retry_vma:
 	vma = lock_vma_under_rcu(mm, address);
 	if (!vma)
 		goto lock_mmap;
@@ -310,14 +311,14 @@ static void do_exception(struct pt_regs *regs, int access)
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 	/* Quick path to respond to signals */
 	if (fault_signal_pending(fault, regs)) {
 		if (!user_mode(regs))
 			handle_fault_error_nolock(regs, 0);
 		return;
 	}
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
 lock_mmap:
 retry:
 	vma = lock_mm_and_find_vma(mm, address, regs);
diff --git a/arch/x86/mm/fault.c b/arch/x86/mm/fault.c
index 45b99c3b1442..53c8f003fe53 100644
--- a/arch/x86/mm/fault.c
+++ b/arch/x86/mm/fault.c
@@ -1331,6 +1331,7 @@ void do_user_addr_fault(struct pt_regs *regs,
 	if (!(flags & FAULT_FLAG_USER))
 		goto lock_mmap;
 
+retry_vma:
 	vma = lock_vma_under_rcu(mm, address);
 	if (!vma)
 		goto lock_mmap;
@@ -1349,8 +1350,6 @@ void do_user_addr_fault(struct pt_regs *regs,
 		goto done;
 	}
 	count_vm_vma_lock_event(VMA_LOCK_RETRY);
-	if (fault & VM_FAULT_MAJOR)
-		flags |= FAULT_FLAG_TRIED;
 
 	/* Quick path to respond to signals */
 	if (fault_signal_pending(fault, regs)) {
@@ -1360,6 +1359,9 @@ void do_user_addr_fault(struct pt_regs *regs,
 						 ARCH_DEFAULT_PKEY);
 		return;
 	}
+
+	if (fault_should_retry_under_vma_lock(fault, &flags))
+		goto retry_vma;
 lock_mmap:
 
 retry:
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 7fabe6c66b4b..27ec6673acfe 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -727,6 +727,34 @@ static inline bool fault_flag_allow_retry_first(enum fault_flag flags)
 	    (!(flags & FAULT_FLAG_TRIED));
 }
 
+/**
+ * fault_should_retry_under_vma_lock - decide whether to retry with VMA lock
+ * @fault: fault result from handle_mm_fault() under FAULT_FLAG_VMA_LOCK
+ * @flags: fault flags for the current fault, updated on retry
+ *
+ * Architecture page fault handlers call this after a VMA-lock fault returns
+ * VM_FAULT_RETRY. If the fault result also has VM_FAULT_MAY_USE_VMA_LOCK,
+ * allow one bounded retry under the VMA lock and set FAULT_FLAG_TRIED.
+ *
+ * When the fault must fall back to the mmap_lock path, preserve the existing
+ * VM_FAULT_MAJOR behavior by marking FAULT_FLAG_TRIED before the retry.
+ *
+ * Return: true if the caller should retry under the VMA lock, false if it
+ * should fall back to the mmap_lock fault path.
+ */
+static inline bool fault_should_retry_under_vma_lock(vm_fault_t fault, unsigned int *flags)
+{
+	if ((fault & VM_FAULT_MAY_USE_VMA_LOCK) && !(*flags & FAULT_FLAG_TRIED)) {
+		*flags |= FAULT_FLAG_TRIED;
+		return true;
+	}
+
+	if (fault & VM_FAULT_MAJOR)
+		*flags |= FAULT_FLAG_TRIED;
+
+	return false;
+}
+
 #define FAULT_FLAG_TRACE \
 	{ FAULT_FLAG_WRITE,		"WRITE" }, \
 	{ FAULT_FLAG_MKWRITE,		"MKWRITE" }, \
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index b5d4cd3b067b..46a832757109 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -1684,6 +1684,8 @@ typedef __bitwise unsigned int vm_fault_t;
  * @VM_FAULT_NOPAGE:		->fault installed the pte, not return page
  * @VM_FAULT_LOCKED:		->fault locked the returned page
  * @VM_FAULT_RETRY:		->fault blocked, must retry
+ * @VM_FAULT_MAY_USE_VMA_LOCK:	->fault blocked, retry may be handled under
+ *				the VMA lock
  * @VM_FAULT_FALLBACK:		huge page fault failed, fall back to small
  * @VM_FAULT_DONE_COW:		->fault has fully handled COW
  * @VM_FAULT_NEEDDSYNC:		->fault did not modify page tables and needs
@@ -1707,6 +1709,7 @@ enum vm_fault_reason {
 	VM_FAULT_DONE_COW       = (__force vm_fault_t)0x001000,
 	VM_FAULT_NEEDDSYNC      = (__force vm_fault_t)0x002000,
 	VM_FAULT_COMPLETED      = (__force vm_fault_t)0x004000,
+	VM_FAULT_MAY_USE_VMA_LOCK = (__force vm_fault_t)0x008000,
 	VM_FAULT_HINDEX_MASK    = (__force vm_fault_t)0x0f0000,
 };
 
@@ -1731,6 +1734,7 @@ enum vm_fault_reason {
 	{ VM_FAULT_FALLBACK,            "FALLBACK" },	\
 	{ VM_FAULT_DONE_COW,            "DONE_COW" },	\
 	{ VM_FAULT_NEEDDSYNC,           "NEEDDSYNC" },	\
+	{ VM_FAULT_MAY_USE_VMA_LOCK,    "MAY_USE_VMA_LOCK" },	\
 	{ VM_FAULT_COMPLETED,           "COMPLETED" }
 
 struct vm_special_mapping {
-- 
2.43.0