[PATCH 12/15] mm: handle ANON_VMA_LAZY during migration

[tao]

To ensure the atomicity of folio migration, introduce
folio_trylock_get_anon_rmap().

This helper guarantees that the migration operation is mutually
exclusive with free_pgtables(). For ANON_VMA_LAZY, it uses
vma_start_read() to prevent the VMA from being modified or removed
during migration.

Signed-off-by: tao <tao.wangtao@xxxxxxxxx>
---
 include/linux/rmap.h | 12 ++++++++
 mm/migrate.c         | 71 +++++++++++++++++++++++++-------------------
 mm/rmap.c            | 40 +++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 31 deletions(-)

diff --git a/include/linux/rmap.h b/include/linux/rmap.h
index ebe9f3f61170..59244481a8c1 100644
--- a/include/linux/rmap.h
+++ b/include/linux/rmap.h
@@ -1042,6 +1042,18 @@ bool folio_maybe_same_anon_vma(const struct folio *folio,
 anon_rmap_t folio_get_anon_rmap(const struct folio *folio);
 anon_rmap_t folio_lock_anon_rmap_read(const struct folio *folio,
 	struct rmap_walk_control *rwc);
+/*
+ * folio_trylock_get_anon_rmap ensures that the migration operation
+ * completes atomically and is mutually exclusive with free_pgtables().
+ *
+ * Note: for ANON_VMA_LAZY, this is not equivalent to
+ * anon_rmap_trylock_read() + folio_get_anon_rmap(), because
+ * anon_rmap_trylock_read() only increments the VMA reference count,
+ * while this helper uses vma_start_read() to prevent the VMA from
+ * being modified or removed.
+ */
+anon_rmap_t folio_trylock_get_anon_rmap(const struct folio *folio);
+void anon_rmap_unlock_put(anon_rmap_t anon_rmap);
 
 static inline struct vm_area_struct *anon_rmap_iter_first_vma(
 	anon_rmap_t anon_rmap, unsigned long start, unsigned long last,
diff --git a/mm/migrate.c b/mm/migrate.c
index b397cdeab09a..4abbfd1faea2 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -1173,10 +1173,11 @@ static void migrate_folio_undo_src(struct folio *src,
 				   struct list_head *ret)
 {
 	if (page_was_mapped)
-		remove_migration_ptes(src, src, 0);
+		remove_migration_ptes(src, src,
+			anon_rmap_value(anon_rmap) ? TTU_RMAP_LOCKED : 0);
 	/* Drop an anon_rmap reference if we took one */
 	if (anon_rmap_value(anon_rmap))
-		put_anon_rmap(anon_rmap);
+		anon_rmap_unlock_put(anon_rmap);
 	if (locked)
 		folio_unlock(src);
 	if (ret)
@@ -1279,6 +1280,18 @@ static int migrate_folio_unmap(new_folio_t get_new_folio,
 		folio_wait_writeback(src);
 	}
 
+	/*
+	 * Block others from accessing the new page when we get around to
+	 * establishing additional references. We are usually the only one
+	 * holding a reference to dst at this point. We used to have a BUG
+	 * here if folio_trylock(dst) fails, but would like to allow for
+	 * cases where there might be a race with the previous use of dst.
+	 * This is much like races on refcount of oldpage: just don't BUG().
+	 */
+	if (unlikely(!folio_trylock(dst)))
+		goto out;
+	dst_locked = true;
+
 	/*
 	 * By try_to_migrate(), src->mapcount goes down to 0 here. In this case,
 	 * we cannot notice that anon_vma is freed while we migrate a page.
@@ -1287,26 +1300,17 @@ static int migrate_folio_unmap(new_folio_t get_new_folio,
 	 * File Caches may use write_page() or lock_page() in migration, then,
 	 * just care Anon page here.
 	 *
-	 * Only folio_get_anon_rmap() understands the subtleties of
-	 * getting a hold on an anon_rmap from outside one of its mms.
+	 * Only folio_trylock_get_anon_rmap() understands the subtleties of
+	 * getting and locking an anon_rmap from outside one of its mms.
 	 * But if we cannot get anon_rmap, then we won't need it anyway,
 	 * because that implies that the anon page is no longer mapped
 	 * (and cannot be remapped so long as we hold the page lock).
 	 */
-	if (folio_test_anon(src) && !folio_test_ksm(src))
-		anon_rmap = folio_get_anon_rmap(src);
-
-	/*
-	 * Block others from accessing the new page when we get around to
-	 * establishing additional references. We are usually the only one
-	 * holding a reference to dst at this point. We used to have a BUG
-	 * here if folio_trylock(dst) fails, but would like to allow for
-	 * cases where there might be a race with the previous use of dst.
-	 * This is much like races on refcount of oldpage: just don't BUG().
-	 */
-	if (unlikely(!folio_trylock(dst)))
-		goto out;
-	dst_locked = true;
+	if (folio_test_anon(src) && !folio_test_ksm(src)) {
+		anon_rmap = folio_trylock_get_anon_rmap(src);
+		if (!anon_rmap_value(anon_rmap))
+			goto out;
+	}
 
 	if (unlikely(page_has_movable_ops(&src->page))) {
 		__migrate_folio_record(dst, old_page_state, anon_rmap);
@@ -1331,10 +1335,14 @@ static int migrate_folio_unmap(new_folio_t get_new_folio,
 			goto out;
 		}
 	} else if (folio_mapped(src)) {
+		enum ttu_flags ttu = mode == MIGRATE_ASYNC ? TTU_BATCH_FLUSH : 0;
+
+		if (anon_rmap_value(anon_rmap))
+			ttu |= TTU_RMAP_LOCKED;
 		/* Establish migration ptes */
 		VM_BUG_ON_FOLIO(folio_test_anon(src) &&
 			       !folio_test_ksm(src) && !anon_rmap_value(anon_rmap), src);
-		try_to_migrate(src, mode == MIGRATE_ASYNC ? TTU_BATCH_FLUSH : 0);
+		try_to_migrate(src, ttu);
 		old_page_state |= PAGE_WAS_MAPPED;
 	}
 
@@ -1415,7 +1423,8 @@ static int migrate_folio_move(free_folio_t put_new_folio, unsigned long private,
 		lru_add_drain();
 
 	if (old_page_state & PAGE_WAS_MAPPED)
-		remove_migration_ptes(src, dst, 0);
+		remove_migration_ptes(src, dst,
+			anon_rmap_value(anon_rmap) ? TTU_RMAP_LOCKED : 0);
 
 out_unlock_both:
 	folio_unlock(dst);
@@ -1434,7 +1443,7 @@ static int migrate_folio_move(free_folio_t put_new_folio, unsigned long private,
 	list_del(&src->lru);
 	/* Drop an anon_rmap reference if we took one */
 	if (anon_rmap_value(anon_rmap))
-		put_anon_rmap(anon_rmap);
+		anon_rmap_unlock_put(anon_rmap);
 	folio_unlock(src);
 	migrate_folio_done(src, reason);
 
@@ -1485,7 +1494,7 @@ static int unmap_and_move_huge_page(new_folio_t get_new_folio,
 	int page_was_mapped = 0;
 	anon_rmap_t anon_rmap = ANON_RMAP_NULL;
 	struct address_space *mapping = NULL;
-	enum ttu_flags ttu = 0;
+	enum ttu_flags ttu = TTU_RMAP_LOCKED;
 
 	if (folio_ref_count(src) == 1) {
 		/* page was freed from under us. So we are done. */
@@ -1519,11 +1528,14 @@ static int unmap_and_move_huge_page(new_folio_t get_new_folio,
 		goto out_unlock;
 	}
 
-	if (folio_test_anon(src))
-		anon_rmap = folio_get_anon_rmap(src);
-
 	if (unlikely(!folio_trylock(dst)))
-		goto put_anon;
+		goto out_unlock;
+
+	if (folio_test_anon(src)) {
+		anon_rmap = folio_trylock_get_anon_rmap(src);
+		if (!anon_rmap_value(anon_rmap))
+			goto unlock_put_anon;
+	}
 
 	if (folio_mapped(src)) {
 		if (!folio_test_anon(src)) {
@@ -1536,8 +1548,6 @@ static int unmap_and_move_huge_page(new_folio_t get_new_folio,
 			mapping = hugetlb_folio_mapping_lock_write(src);
 			if (unlikely(!mapping))
 				goto unlock_put_anon;
-
-			ttu = TTU_RMAP_LOCKED;
 		}
 
 		try_to_migrate(src, ttu);
@@ -1550,15 +1560,14 @@ static int unmap_and_move_huge_page(new_folio_t get_new_folio,
 	if (page_was_mapped)
 		remove_migration_ptes(src, !rc ? dst : src, ttu);
 
-	if (ttu & TTU_RMAP_LOCKED)
+	if (page_was_mapped && !folio_test_anon(src))
 		i_mmap_unlock_write(mapping);
 
 unlock_put_anon:
 	folio_unlock(dst);
 
-put_anon:
 	if (anon_rmap_value(anon_rmap))
-		put_anon_rmap(anon_rmap);
+		anon_rmap_unlock_put(anon_rmap);
 
 	if (!rc) {
 		move_hugetlb_state(src, dst, reason);
diff --git a/mm/rmap.c b/mm/rmap.c
index 57cd85efc50a..46876b3dbfbc 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1223,6 +1223,46 @@ anon_rmap_t folio_lock_anon_rmap_read(const struct folio *folio,
 	return anon_vma ? anon_vma_to_anon_rmap(anon_vma) : ANON_RMAP_NULL;
 }
 
+anon_rmap_t folio_trylock_get_anon_rmap(const struct folio *folio)
+{
+	struct anon_vma *anon_vma;
+	struct vm_area_struct *vma;
+
+	if (folio_test_anon_vma_lazy(folio)) {
+		vma = folio_get_anon_vma_lazy(folio);
+		if (vma && !lock_vma_under_rcu(vma->vm_mm, vma->vm_start)) {
+			vma_put(vma);
+			vma = NULL;
+		}
+		if (vma)
+			return vma_to_anon_rmap(vma);
+	}
+
+	anon_vma = folio_get_anon_vma(folio);
+	if (anon_vma && !anon_vma_trylock_read(anon_vma)) {
+		put_anon_vma(anon_vma);
+		anon_vma = NULL;
+	}
+	return anon_vma ? anon_vma_to_anon_rmap(anon_vma) : ANON_RMAP_NULL;
+}
+
+void anon_rmap_unlock_put(anon_rmap_t anon_rmap)
+{
+	struct anon_vma *anon_vma;
+
+	if (!anon_rmap_is_anon_vma(anon_rmap)) {
+		struct vm_area_struct *vma = anon_rmap_to_vma(anon_rmap);
+
+		vma_end_read(vma);
+		vma_put(vma);
+		return;
+	}
+
+	anon_vma = anon_rmap_to_anon_vma(anon_rmap);
+	anon_vma_unlock_read(anon_vma);
+	put_anon_vma(anon_vma);
+}
+
 #ifdef CONFIG_ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH
 /*
  * Flush TLB entries for recently unmapped pages from remote CPUs. It is
-- 
2.17.1