[Devel] [PATCH RFC PATCH hci-8.1 1/4] mm: support moving anonymous page tables between address spaces

[Andrey Drobyshev <andrey.drobyshev@virtuozzo.com>]

Teach move_page_tables and move_huge_pmd to use the destination mm for
page-table operations and transfer RSS, swap-list and huge-page table
accounting. Restrict cross-mm moves to same-address anonymous DOEXEC VMAs.

Add the parallel exec move wrapper, joining all workers before returning
on failure. Existing DOEXEC callers still use the copy path in this commit.

https://virtuozzo.atlassian.net/browse/VSTOR-103860

Signed-off-by: Andrey Drobyshev <andrey.drobyshev@virtuozzo.com>
---
 include/linux/huge_mm.h |   3 +-
 mm/huge_memory.c        |  28 +++++++---
 mm/internal.h           |   2 +
 mm/mremap.c             | 140 +++++++++++++++++++++++++++++++++++++++++++-----
 4 files changed, 152 insertions(+), 21 deletions(-)

diff --git a/include/linux/huge_mm.h b/include/linux/huge_mm.h
index b94c2e8ee918..e2ec0348bebd 100644
--- a/include/linux/huge_mm.h
+++ b/include/linux/huge_mm.h
@@ -31,7 +31,8 @@ int zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma, pmd_t *pmd,
 		 unsigned long addr);
 int zap_huge_pud(struct mmu_gather *tlb, struct vm_area_struct *vma, pud_t *pud,
 		 unsigned long addr);
-bool move_huge_pmd(struct vm_area_struct *vma, unsigned long old_addr,
+bool move_huge_pmd(struct vm_area_struct *vma, struct vm_area_struct *new_vma,
+		   unsigned long old_addr,
 		   unsigned long new_addr, pmd_t *old_pmd, pmd_t *new_pmd);
 int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
 		    pmd_t *pmd, unsigned long addr, pgprot_t newprot,
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 90868f69ee33..9954dac5e53e 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -2219,12 +2219,15 @@ static pmd_t clear_uffd_wp_pmd(pmd_t pmd)
 	return pmd;
 }
 
-bool move_huge_pmd(struct vm_area_struct *vma, unsigned long old_addr,
+bool move_huge_pmd(struct vm_area_struct *vma, struct vm_area_struct *new_vma,
+		  unsigned long old_addr,
 		  unsigned long new_addr, pmd_t *old_pmd, pmd_t *new_pmd)
 {
 	spinlock_t *old_ptl, *new_ptl;
 	pmd_t pmd;
 	struct mm_struct *mm = vma->vm_mm;
+	struct mm_struct *new_mm = new_vma->vm_mm;
+	bool cross_mm = mm != new_mm;
 	bool force_flush = false;
 
 	/*
@@ -2243,7 +2246,7 @@ bool move_huge_pmd(struct vm_area_struct *vma, unsigned long old_addr,
 	 */
 	old_ptl = __pmd_trans_huge_lock(old_pmd, vma);
 	if (old_ptl) {
-		new_ptl = pmd_lockptr(mm, new_pmd);
+		new_ptl = pmd_lockptr(new_mm, new_pmd);
 		if (new_ptl != old_ptl)
 			spin_lock_nested(new_ptl, SINGLE_DEPTH_NESTING);
 		pmd = pmdp_huge_get_and_clear(mm, old_addr, old_pmd);
@@ -2251,15 +2254,28 @@ bool move_huge_pmd(struct vm_area_struct *vma, unsigned long old_addr,
 			force_flush = true;
 		VM_BUG_ON(!pmd_none(*new_pmd));
 
-		if (pmd_move_must_withdraw(new_ptl, old_ptl, vma)) {
+		if (cross_mm || pmd_move_must_withdraw(new_ptl, old_ptl, vma)) {
 			pgtable_t pgtable;
 			pgtable = pgtable_trans_huge_withdraw(mm, old_pmd);
-			pgtable_trans_huge_deposit(mm, new_pmd, pgtable);
+			pgtable_trans_huge_deposit(new_mm, new_pmd, pgtable);
+			if (cross_mm) {
+				mm_dec_nr_ptes(mm);
+				mm_inc_nr_ptes(new_mm);
+			}
+		}
+		if (cross_mm) {
+			if (is_huge_zero_pmd(pmd)) {
+				/* The source retains its per-mm zero-page reference. */
+				mm_get_huge_zero_folio(new_mm);
+			} else {
+				add_mm_counter(mm, MM_ANONPAGES, -HPAGE_PMD_NR);
+				add_mm_counter(new_mm, MM_ANONPAGES, HPAGE_PMD_NR);
+			}
 		}
 		pmd = move_soft_dirty_pmd(pmd);
-		if (vma_has_uffd_without_event_remap(vma))
+		if (cross_mm || vma_has_uffd_without_event_remap(vma))
 			pmd = clear_uffd_wp_pmd(pmd);
-		set_pmd_at(mm, new_addr, new_pmd, pmd);
+		set_pmd_at(new_mm, new_addr, new_pmd, pmd);
 		if (force_flush)
 			flush_pmd_tlb_range(vma, old_addr, old_addr + PMD_SIZE);
 		if (new_ptl != old_ptl)
diff --git a/mm/internal.h b/mm/internal.h
index d57e67a83d70..b0fed4393315 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1514,6 +1514,8 @@ extern struct list_lru shadow_nodes;
 } while (0)
 
 /* mremap.c */
+int move_page_range_exec(struct vm_area_struct *new_vma,
+			 struct vm_area_struct *old_vma);
 unsigned long move_page_tables(struct vm_area_struct *vma,
 	unsigned long old_addr, struct vm_area_struct *new_vma,
 	unsigned long new_addr, unsigned long len,
diff --git a/mm/mremap.c b/mm/mremap.c
index 1e6d40a240c9..6f64ff7a8aaa 100644
--- a/mm/mremap.c
+++ b/mm/mremap.c
@@ -25,6 +25,7 @@
 #include <linux/uaccess.h>
 #include <linux/userfaultfd_k.h>
 #include <linux/mempolicy.h>
+#include <linux/padata.h>
 
 #include <asm/cacheflush.h>
 #include <asm/tlb.h>
@@ -133,13 +134,59 @@ static pte_t move_soft_dirty_pte(pte_t pte)
 	return pte;
 }
 
+/* Transfer RSS without duplicating page references, rmaps or swap entries. */
+static void move_pte_account(struct vm_area_struct *vma, unsigned long addr,
+			     struct mm_struct *new_mm, pte_t pte)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	struct folio *folio;
+	swp_entry_t entry;
+	int counter;
+
+	if (pte_present(pte)) {
+		folio = vm_normal_folio(vma, addr, pte);
+		if (!folio) {
+#ifdef CONFIG_KSM
+			if (is_ksm_zero_pte(pte)) {
+				atomic_long_dec(&mm->ksm_zero_pages);
+				atomic_long_inc(&new_mm->ksm_zero_pages);
+			}
+#endif
+			return;
+		}
+		counter = mm_counter(folio);
+	} else {
+		entry = pte_to_swp_entry(pte);
+		if (!non_swap_entry(entry)) {
+			/* swapoff must also visit the destination mm. */
+			if (list_empty(&new_mm->mmlist)) {
+				spin_lock(&mmlist_lock);
+				if (list_empty(&new_mm->mmlist))
+					list_add(&new_mm->mmlist, &mm->mmlist);
+				spin_unlock(&mmlist_lock);
+			}
+			counter = MM_SWAPENTS;
+		} else if (is_migration_entry(entry) ||
+			   is_device_private_entry(entry) ||
+			   is_device_exclusive_entry(entry)) {
+			counter = mm_counter(pfn_swap_entry_folio(entry));
+		} else {
+			return;
+		}
+	}
+	dec_mm_counter(mm, counter);
+	inc_mm_counter(new_mm, counter);
+}
+
 static int move_ptes(struct vm_area_struct *vma, pmd_t *old_pmd,
 		unsigned long old_addr, unsigned long old_end,
 		struct vm_area_struct *new_vma, pmd_t *new_pmd,
 		unsigned long new_addr, bool need_rmap_locks)
 {
-	bool need_clear_uffd_wp = vma_has_uffd_without_event_remap(vma);
 	struct mm_struct *mm = vma->vm_mm;
+	struct mm_struct *new_mm = new_vma->vm_mm;
+	bool cross_mm = mm != new_mm;
+	bool need_clear_uffd_wp = cross_mm || vma_has_uffd_without_event_remap(vma);
 	pte_t *old_pte, *new_pte, pte;
 	pmd_t dummy_pmdval;
 	spinlock_t *old_ptl, *new_ptl;
@@ -184,7 +231,7 @@ static int move_ptes(struct vm_area_struct *vma, pmd_t *old_pmd,
 	 * mmap_lock, so this new_pte page is stable, so there is no need to get
 	 * pmdval and do pmd_same() check.
 	 */
-	new_pte = pte_offset_map_rw_nolock(mm, new_pmd, new_addr, &dummy_pmdval,
+	new_pte = pte_offset_map_rw_nolock(new_mm, new_pmd, new_addr, &dummy_pmdval,
 					   &new_ptl);
 	if (!new_pte) {
 		pte_unmap_unlock(old_pte, old_ptl);
@@ -202,6 +249,8 @@ static int move_ptes(struct vm_area_struct *vma, pmd_t *old_pmd,
 			continue;
 
 		pte = ptep_get_and_clear(mm, old_addr, old_pte);
+		if (cross_mm)
+			move_pte_account(vma, old_addr, new_mm, pte);
 		/*
 		 * If we are remapping a valid PTE, make sure
 		 * to flush TLB before we drop the PTL for the
@@ -219,7 +268,7 @@ static int move_ptes(struct vm_area_struct *vma, pmd_t *old_pmd,
 		pte = move_soft_dirty_pte(pte);
 
 		if (need_clear_uffd_wp && pte_marker_uffd_wp(pte))
-			pte_clear(mm, new_addr, new_pte);
+			pte_clear(new_mm, new_addr, new_pte);
 		else {
 			if (need_clear_uffd_wp) {
 				if (pte_present(pte))
@@ -227,7 +276,7 @@ static int move_ptes(struct vm_area_struct *vma, pmd_t *old_pmd,
 				else if (is_swap_pte(pte))
 					pte = pte_swp_clear_uffd_wp(pte);
 			}
-			set_pte_at(mm, new_addr, new_pte, pte);
+			set_pte_at(new_mm, new_addr, new_pte, pte);
 		}
 	}
 
@@ -495,6 +544,7 @@ static __always_inline unsigned long get_extent(enum pgt_entry entry,
  * pgt_entry. Returns true if the move was successful, else false.
  */
 static bool move_pgt_entry(enum pgt_entry entry, struct vm_area_struct *vma,
+			struct vm_area_struct *new_vma,
 			unsigned long old_addr, unsigned long new_addr,
 			void *old_entry, void *new_entry, bool need_rmap_locks)
 {
@@ -515,7 +565,7 @@ static bool move_pgt_entry(enum pgt_entry entry, struct vm_area_struct *vma,
 		break;
 	case HPAGE_PMD:
 		moved = IS_ENABLED(CONFIG_TRANSPARENT_HUGEPAGE) &&
-			move_huge_pmd(vma, old_addr, new_addr, old_entry,
+			move_huge_pmd(vma, new_vma, old_addr, new_addr, old_entry,
 				      new_entry);
 		break;
 	case HPAGE_PUD:
@@ -596,10 +646,19 @@ unsigned long move_page_tables(struct vm_area_struct *vma,
 	struct mmu_notifier_range range;
 	pmd_t *old_pmd, *new_pmd;
 	pud_t *old_pud, *new_pud;
+	bool cross_mm = vma->vm_mm != new_vma->vm_mm;
 
 	if (!len)
 		return 0;
 
+	/* Cross-mm moves are only used for same-address anonymous DOEXEC VMAs. */
+	if (cross_mm) {
+		if (WARN_ON_ONCE(!vma_is_anonymous(vma) ||
+				 !(vma->vm_flags & VM_EXEC_KEEP) ||
+				 old_addr != new_addr || !need_rmap_locks))
+			return 0;
+	}
+
 	old_end = old_addr + len;
 
 	if (is_vm_hugetlb_page(vma))
@@ -610,7 +669,7 @@ unsigned long move_page_tables(struct vm_area_struct *vma,
 	 * If possible, realign addresses to PMD boundary for faster copy.
 	 * Only realign if the mremap copying hits a PMD boundary.
 	 */
-	if (len >= PMD_SIZE - (old_addr & ~PMD_MASK))
+	if (!cross_mm && len >= PMD_SIZE - (old_addr & ~PMD_MASK))
 		try_realign_addr(&old_addr, vma, &new_addr, new_vma, PMD_MASK,
 				 for_stack);
 
@@ -630,19 +689,23 @@ unsigned long move_page_tables(struct vm_area_struct *vma,
 		old_pud = get_old_pud(vma->vm_mm, old_addr);
 		if (!old_pud)
 			continue;
-		new_pud = alloc_new_pud(vma->vm_mm, vma, new_addr);
+		new_pud = alloc_new_pud(new_vma->vm_mm, new_vma, new_addr);
 		if (!new_pud)
 			break;
 		if (pud_trans_huge(*old_pud) || pud_devmap(*old_pud)) {
+			/* Anonymous PUD-sized THPs are not supported. */
+			if (cross_mm)
+				break;
 			if (extent == HPAGE_PUD_SIZE) {
-				move_pgt_entry(HPAGE_PUD, vma, old_addr, new_addr,
+				move_pgt_entry(HPAGE_PUD, vma, new_vma, old_addr, new_addr,
 					       old_pud, new_pud, need_rmap_locks);
 				/* We ignore and continue on error? */
 				continue;
 			}
-		} else if (IS_ENABLED(CONFIG_HAVE_MOVE_PUD) && extent == PUD_SIZE) {
+		} else if (!cross_mm && IS_ENABLED(CONFIG_HAVE_MOVE_PUD) &&
+			   extent == PUD_SIZE) {
 
-			if (move_pgt_entry(NORMAL_PUD, vma, old_addr, new_addr,
+			if (move_pgt_entry(NORMAL_PUD, vma, new_vma, old_addr, new_addr,
 					   old_pud, new_pud, true))
 				continue;
 		}
@@ -651,24 +714,24 @@ unsigned long move_page_tables(struct vm_area_struct *vma,
 		old_pmd = get_old_pmd(vma->vm_mm, old_addr);
 		if (!old_pmd)
 			continue;
-		new_pmd = alloc_new_pmd(vma->vm_mm, vma, new_addr);
+		new_pmd = alloc_new_pmd(new_vma->vm_mm, new_vma, new_addr);
 		if (!new_pmd)
 			break;
 again:
 		if (is_swap_pmd(*old_pmd) || pmd_trans_huge(*old_pmd) ||
 		    pmd_devmap(*old_pmd)) {
 			if (extent == HPAGE_PMD_SIZE &&
-			    move_pgt_entry(HPAGE_PMD, vma, old_addr, new_addr,
+			    move_pgt_entry(HPAGE_PMD, vma, new_vma, old_addr, new_addr,
 					   old_pmd, new_pmd, need_rmap_locks))
 				continue;
 			split_huge_pmd(vma, old_pmd, old_addr);
-		} else if (IS_ENABLED(CONFIG_HAVE_MOVE_PMD) &&
+		} else if (!cross_mm && IS_ENABLED(CONFIG_HAVE_MOVE_PMD) &&
 			   extent == PMD_SIZE) {
 			/*
 			 * If the extent is PMD-sized, try to speed the move by
 			 * moving at the PMD level if possible.
 			 */
-			if (move_pgt_entry(NORMAL_PMD, vma, old_addr, new_addr,
+			if (move_pgt_entry(NORMAL_PMD, vma, new_vma, old_addr, new_addr,
 					   old_pmd, new_pmd, true))
 				continue;
 		}
@@ -693,6 +756,55 @@ unsigned long move_page_tables(struct vm_area_struct *vma,
 	return len + old_addr - old_end;	/* how much done */
 }
 
+struct move_page_range_args {
+	struct vm_area_struct *old_vma;
+	struct vm_area_struct *new_vma;
+};
+
+static int move_page_range_chunk(unsigned long addr, unsigned long end, void *arg)
+{
+	struct move_page_range_args *args = arg;
+	unsigned long len = end - addr;
+
+	return move_page_tables(args->old_vma, addr, args->new_vma, addr,
+				len, true, false) == len ? 0 : -ENOMEM;
+}
+
+/*
+ * The caller holds both mmap locks while workers transfer disjoint PMDs.
+ * A failure can leave mappings in both mms; there is no rollback. This is
+ * only usable after exec's point of no return, where errors are fatal.
+ * padata_do_multithreaded() joins all workers, including on error, before
+ * either mm may be torn down.
+ */
+int move_page_range_exec(struct vm_area_struct *new_vma,
+			 struct vm_area_struct *old_vma)
+{
+	struct move_page_range_args args = { old_vma, new_vma };
+#ifdef CONFIG_PADATA
+	struct padata_mt_job job = {
+		.thread_fn = move_page_range_chunk,
+		.fn_arg = &args,
+		.start = old_vma->vm_start,
+		.size = old_vma->vm_end - old_vma->vm_start,
+		.align = PMD_SIZE,
+		.min_chunk = max(1ul << 27, PMD_SIZE),
+		.max_threads = 16,
+	};
+#endif
+
+	mmap_assert_write_locked(old_vma->vm_mm);
+	mmap_assert_write_locked(new_vma->vm_mm);
+	if (WARN_ON_ONCE(!(old_vma->vm_flags & VM_EXEC_KEEP)))
+		return -EINVAL;
+
+#ifdef CONFIG_PADATA
+	return padata_do_multithreaded(&job);
+#else
+	return move_page_range_chunk(old_vma->vm_start, old_vma->vm_end, &args);
+#endif
+}
+
 static unsigned long move_vma(struct vm_area_struct *vma,
 		unsigned long old_addr, unsigned long old_len,
 		unsigned long new_len, unsigned long new_addr,

-- 
2.47.1

_______________________________________________
Devel mailing list
Devel@openvz.org
https://lists.openvz.org/mailman/listinfo/devel