[RFC PATCH 2/2] mm/huge_memory: detect partially mapped shared folios without per-page mapcounts

["Jose A. Perez de Azpillaga" <azpijr@gmail.com>]

With CONFIG_NO_PAGE_MAPCOUNT, __folio_remove_rmap() takes an anon folio
as partially mapped when its mapcount drops below the number of pages
and no entire mapping is left. With several MMs mapping the folio, the
mapcount can stay at or above the number of pages while some pages are
mapped by nobody:

 - parent and child zap the same half of a THP after fork();

 - two children zapped the same half and the parent removes the last
   PMD mapping.

The folio is never queued for deferred split, and the unused pages stay
allocated until every mapper is gone.

Queue the folio without PG_partially_mapped when the mapcount can't
tell, and have the shrinker walk the rmap of these candidates. If a page
is mapped by nobody while another one still is, mark the folio partially
mapped and split it as usual. Otherwise drop it from the queue, unless
the underused check wants it.

The walk needs one bit per page, which is too much for the stack with
large base pages, so it uses a static bitmap under a mutex.

Nothing changes with per-page mapcounts.

512 MB of THPs with the same half zapped by every mapper, memory freed
by the shrinker:

                                before    after
  2M,  2 processes                0 MB   252 MB
  2M,  3 processes                0 MB   237 MB
  2M,  2 children + parent PMD    0 MB   244 MB
  64K, 2 processes                0 MB   252 MB
  64K, 3 processes                0 MB   252 MB

The walk takes about 4 us per 64K folio in a TCG guest.

Signed-off-by: Jose A. Perez de Azpillaga <azpijr@gmail.com>
---
 mm/huge_memory.c | 84 +++++++++++++++++++++++++++++++++++++++++++++---
 mm/rmap.c        | 21 ++++++++++--
 2 files changed, 98 insertions(+), 7 deletions(-)

diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 8c4a69c797ed..f80fa028dd34 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -4632,6 +4632,10 @@ void deferred_split_folio(struct folio *folio, bool partially_mapped)
 	if (folio_order(folio) <= 1)
 		return;
 
+	/* Already queued, and there is no flag to set. */
+	if (!partially_mapped && !data_race(list_empty(&folio->_deferred_list)))
+		return;
+
 	/*
 	 * Exclude swapcache: originally to avoid a corrupt deferred split
 	 * queue. Nowadays that is fully prevented by __memcg1_swapout();
@@ -4655,9 +4659,6 @@ void deferred_split_folio(struct folio *folio, bool partially_mapped)
 			count_mthp_stat(folio_order(folio), MTHP_STAT_SPLIT_DEFERRED);
 			mod_mthp_stat(folio_order(folio), MTHP_STAT_NR_ANON_PARTIALLY_MAPPED, 1);
 		}
-	} else {
-		/* partially mapped folios cannot become non-partially mapped */
-		VM_WARN_ON_FOLIO(folio_test_partially_mapped(folio), folio);
 	}
 	__list_lru_add(&deferred_split_lru, lru, &folio->_deferred_list, nid, memcg);
 	list_lru_unlock_irqrestore(lru, &flags);
@@ -4700,6 +4701,72 @@ static bool thp_underused(struct folio *folio)
 	return false;
 }
 
+/*
+ * One bit per page of the folio being checked. Too big for the stack of a
+ * shrinker with large base pages, so there is a single one.
+ */
+static DEFINE_MUTEX(mapped_pages_lock);
+static DECLARE_BITMAP(mapped_pages, MAX_PTRS_PER_PTE);
+
+static bool folio_mapped_pages_one(struct folio *folio,
+		struct vm_area_struct *vma, unsigned long addr, void *arg)
+{
+	const unsigned long nr_pages = folio_nr_pages(folio);
+	DEFINE_FOLIO_VMA_WALK(pvmw, folio, vma, addr, 0);
+	unsigned long *nr_mapped = arg;
+	pgoff_t idx;
+
+	while (page_vma_mapped_walk(&pvmw)) {
+		if (!pvmw.pte) {
+			*nr_mapped = nr_pages;
+			page_vma_mapped_walk_done(&pvmw);
+			break;
+		}
+		idx = linear_anon_page_index(vma, pvmw.address) -
+		      folio_pgoff(folio);
+		if (idx < nr_pages && !__test_and_set_bit(idx, mapped_pages))
+			(*nr_mapped)++;
+	}
+
+	return *nr_mapped < nr_pages;
+}
+
+/*
+ * Without per-page mapcounts, a partial unmap of a folio that is mapped by
+ * several MMs is only detected at unmap time if the total mapcount drops
+ * below the number of pages. Check the page tables to find out whether a
+ * page is mapped by nobody while another one is still mapped.
+ */
+static bool folio_mapped_shared_partially(struct folio *folio)
+{
+	const unsigned long nr_pages = folio_nr_pages(folio);
+	unsigned long nr_mapped = 0;
+	struct rmap_walk_control rwc = {
+		.rmap_one = folio_mapped_pages_one,
+		.arg = &nr_mapped,
+		.anon_lock = folio_lock_anon_vma_read,
+	};
+
+	VM_WARN_ON_FOLIO(!folio_test_locked(folio), folio);
+
+	if (nr_pages > MAX_PTRS_PER_PTE)
+		return false;
+
+	mutex_lock(&mapped_pages_lock);
+	bitmap_zero(mapped_pages, nr_pages);
+	rmap_walk(folio, &rwc);
+	mutex_unlock(&mapped_pages_lock);
+
+	return nr_mapped && nr_mapped < nr_pages;
+}
+
+static bool folio_needs_mapped_check(struct folio *folio)
+{
+	return IS_ENABLED(CONFIG_NO_PAGE_MAPCOUNT) && folio_test_anon(folio) &&
+	       !folio_test_swapcache(folio) && !folio_entire_mapcount(folio) &&
+	       folio_maybe_mapped_shared(folio);
+}
+
 static enum lru_status deferred_split_isolate(struct list_head *item,
 					      struct list_lru_one *lru,
 					      void *cb_arg)
@@ -4744,6 +4811,15 @@ static unsigned long deferred_split_scan(struct shrinker *shrink,
 		list_lru_unlock_irq(lru);
 		rcu_read_unlock();
 
+		if (!folio_test_partially_mapped(folio) &&
+		    folio_needs_mapped_check(folio)) {
+			if (!folio_trylock(folio))
+				goto requeue;
+			/* Sets PG_partially_mapped and puts it back on the queue. */
+			if (folio_mapped_shared_partially(folio))
+				deferred_split_folio(folio, true);
+			folio_unlock(folio);
+		}
 		if (!folio_test_partially_mapped(folio)) {
 			/*
 			 * See try_to_map_unused_to_zeropage(): we cannot
@@ -4752,7 +4828,7 @@ static unsigned long deferred_split_scan(struct shrinker *shrink,
 			 */
 			if (folio_test_mlocked(folio))
 				goto next;
-			underused = thp_underused(folio);
+			underused = split_underused_thp && thp_underused(folio);
 			if (!underused)
 				goto next;
 		}
diff --git a/mm/rmap.c b/mm/rmap.c
index 0faef46806e5..95444ed6189a 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1803,6 +1803,7 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 	atomic_t *mapped = &folio->_nr_pages_mapped;
 	int last = 0, nr = 0, nr_pmdmapped = 0;
 	bool partially_mapped = false;
+	bool maybe_partially_mapped = false;
 
 	__folio_rmap_sanity_checks(folio, page, nr_pages, level);
 
@@ -1821,6 +1822,16 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 			} else {
 				partially_mapped = nr < folio_large_nr_pages(folio) &&
 						   !folio_entire_mapcount(folio);
+				/*
+				 * With other MMs mapping the folio, the total
+				 * mapcount cannot tell whether the pages we
+				 * just unmapped are still mapped by anyone.
+				 * Let the shrinker find out.
+				 */
+				maybe_partially_mapped = !partially_mapped &&
+					nr_pages < folio_large_nr_pages(folio) &&
+					!folio_entire_mapcount(folio) &&
+					folio_maybe_mapped_shared(folio);
 				nr = 0;
 			}
 			break;
@@ -1850,6 +1861,10 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 			} else {
 				partially_mapped = last &&
 						   nr < folio_large_nr_pages(folio);
+				/* Only PTE mappings by several MMs are left. */
+				maybe_partially_mapped = last &&
+					!partially_mapped &&
+					folio_maybe_mapped_shared(folio);
 				nr = 0;
 			}
 			break;
@@ -1888,10 +1903,10 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 	 * Device private folios do not support deferred splitting and
 	 * shrinker based scanning of the folios to free.
 	 */
-	if (partially_mapped && folio_test_anon(folio) &&
-	    !folio_test_partially_mapped(folio) &&
+	if ((partially_mapped || maybe_partially_mapped) &&
+	    folio_test_anon(folio) && !folio_test_partially_mapped(folio) &&
 	    !folio_is_device_private(folio))
-		deferred_split_folio(folio, true);
+		deferred_split_folio(folio, partially_mapped);
 
 	__folio_mod_stat(folio, -nr, -nr_pmdmapped);
 
-- 
2.56.0