[PATCH 2/4] mm: don't report OOM when a huge PMD race blocks a pfn insert

[Eliav Farber <farbere@amazon.com>]

insert_pfn() maps the PTE with get_locked_pte() and treats a NULL return
as an allocation failure. pte_offset_map_lock() also returns NULL when
the PMD is a leaf, which a VM_PFNMAP mapping with both ->fault and
->huge_fault can hit: an order-0 fault that passed handle_pte_fault()'s
pmd_none() check can reach insert_pfn() after a concurrent huge fault
installed a PMD leaf over the same address. The task then takes a
VM_FAULT_OOM for a mapping that is perfectly healthy. Changing an mm's
THP policy while another thread faults the same PMD, for example with
prctl(PR_SET_THP_DISABLE), is enough to produce the mixed orders.

Walk to the PMD first and tell the two cases apart: a failed PTE
allocation stays VM_FAULT_OOM, while a PMD leaf returns VM_FAULT_NOPAGE
so that the access refaults onto the huge entry.

Fixes: f9e54c3a2f5b ("vfio/pci: implement huge_fault support")
Signed-off-by: Eliav Farber <farbere@amazon.com>
---
 mm/memory.c | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/mm/memory.c b/mm/memory.c
index 0f54ee0a6d88..b7b89d041f42 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -2866,10 +2866,25 @@ static vm_fault_t insert_pfn(struct vm_area_struct *vma, unsigned long addr,
 	struct mm_struct *mm = vma->vm_mm;
 	pte_t *pte, entry;
 	spinlock_t *ptl;
+	pmd_t *pmd;
 
-	pte = get_locked_pte(mm, addr, &ptl);
-	if (!pte)
+	pmd = walk_to_pmd(mm, addr);
+	if (!pmd)
+		return VM_FAULT_OOM;
+
+	pte = pte_alloc_map_lock(mm, pmd, addr, &ptl);
+	if (!pte) {
+		/*
+		 * A concurrent ->huge_fault may have installed a PMD leaf over
+		 * addr, which pte_offset_map_lock() refuses to walk. The
+		 * mapping is fine, so let the access refault instead of
+		 * reporting a failed allocation.
+		 */
+		if (pmd_trans_huge(pmdp_get_lockless(pmd)))
+			return VM_FAULT_NOPAGE;
 		return VM_FAULT_OOM;
+	}
+
 	entry = ptep_get(pte);
 	if (!pte_none(entry)) {
 		if (mkwrite) {
-- 
2.47.3