[PATCH] mm/gup: factor out LRU cache draining for folio into lru_cache_drain_for_folio()

[David Hildenbrand (Arm)]

KVM with guest_memfd wants to remove any folio references due to LRU
caches, as it really must only allow to convert folios from shared to
private when there are no unexpected folio references (e.g., from GUP
references).

So, to drive the refcount down, it needs a way to flush the LRU caches.
Let's factor out what we have in lru_cache_drain_for_folio(),

Maybe there is a chance to avoid the draining entirely in the future,
by avoiding extra references from the LRU cache: Hugh thinks there might
be a way. But for the time being, this handling is unfortunately
required.

Make folio_may_be_lru_cached() accept a const pointer so
lru_cache_drain_for_folio() can accept a const pointer as well.

Signed-off-by: David Hildenbrand (Arm) <david@xxxxxxxxxx>
---
 include/linux/swap.h |  8 ++++++++
 mm/folio.c           | 46 ++++++++++++++++++++++++++++++++++++++++++++
 mm/gup.c             | 15 ++-------------
 mm/internal.h        |  2 +-
 4 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/include/linux/swap.h b/include/linux/swap.h
index 0544b2ec4c565..733b84e0bce43 100644
--- a/include/linux/swap.h
+++ b/include/linux/swap.h
@@ -298,6 +298,14 @@ void folio_add_lru(struct folio *folio);
 void folio_mark_accessed(struct folio *folio);
 void lru_add_drain_all(void);

+enum lru_cache_drained {
+	LRU_CACHE_NOT_DRAINED,
+	LRU_CACHE_DRAINED,
+	LRU_CACHE_DRAINED_ALL,
+};
+void lru_cache_drain_for_folio(const struct folio *folio,
+		unsigned int extra_refs, enum lru_cache_drained *drained);
+
 /* linux/mm/folio-compat.c */
 void mark_page_accessed(struct page *page);

diff --git a/mm/folio.c b/mm/folio.c
index d2937600cf726..a41d107b8b28a 100644
--- a/mm/folio.c
+++ b/mm/folio.c
@@ -948,6 +948,52 @@ void lru_add_drain_all(void)
 }
 #endif /* CONFIG_SMP */

+/**
+ * lru_cache_drain_for_folio() - drain LRU caches if the folio might have
+ *				 references from the caches
+ * @folio: The folio.
+ * @extra_refs: Extra folio references held by the caller.
+ * @drained: Drain status for batch folio processing.
+ *
+ * Drain LRU caches if the folio might be referenced by these caches, such
+ * that e.g., later page migration will not fail due to them. Start with a
+ * local LRU cache drain, to then drain LRU caches on all CPUs if local
+ * draining was insufficient.
+ *
+ * This function detects LRU cache references by comparing the folio refcount
+ * with the sum of the expected folio refcount + extra references held by the
+ * caller. Note that we cannot rely on PG_lru to reliably detect all LRU
+ * cache references.
+ *
+ * If @drained is not NULL, the function will avoid re-draining LRU caches
+ * when processing multiple folios in a row. In that case, the value
+ * @drained points to must be initialized to LRU_CACHE_NOT_DRAINED before
+ * the first invocation, to keep passing @drained to successive invocations.
+ */
+void lru_cache_drain_for_folio(const struct folio *folio,
+		unsigned int extra_refs, enum lru_cache_drained *drained)
+{
+	const int expected_refs = folio_expected_ref_count(folio) + extra_refs;
+
+	if (!folio_may_be_lru_cached(folio))
+		return;
+
+	if (!drained || *drained == LRU_CACHE_NOT_DRAINED) {
+		if (folio_ref_count(folio) == expected_refs)
+			return;
+		lru_add_drain();
+		if (drained)
+			*drained = LRU_CACHE_DRAINED;
+	}
+	if (!drained || *drained == LRU_CACHE_DRAINED) {
+		if (folio_ref_count(folio) == expected_refs)
+			return;
+		lru_add_drain_all();
+		if (drained)
+			*drained = LRU_CACHE_DRAINED_ALL;
+	}
+}
+
 atomic_t lru_disable_count = ATOMIC_INIT(0);

 /*
diff --git a/mm/gup.c b/mm/gup.c
index 41c3317e0f0f4..297e7de81c374 100644
--- a/mm/gup.c
+++ b/mm/gup.c
@@ -2266,9 +2266,9 @@ static unsigned long collect_longterm_unpinnable_folios(
 		struct list_head *movable_folio_list,
 		struct pages_or_folios *pofs)
 {
+	enum lru_cache_drained drained = LRU_CACHE_NOT_DRAINED;
 	unsigned long collected = 0;
 	struct folio *folio;
-	int drained = 0;
 	long i = 0;

 	for (folio = pofs_get_folio(pofs, i); folio;
@@ -2293,18 +2293,7 @@ static unsigned long collect_longterm_unpinnable_folios(
 		 * but also to remove any other folio references from LRU
 		 * caches.
 		 */
-		if (drained == 0 && folio_may_be_lru_cached(folio) &&
-				folio_ref_count(folio) !=
-				folio_expected_ref_count(folio) + pin_refs) {
-			lru_add_drain();
-			drained = 1;
-		}
-		if (drained == 1 && folio_may_be_lru_cached(folio) &&
-				folio_ref_count(folio) !=
-				folio_expected_ref_count(folio) + pin_refs) {
-			lru_add_drain_all();
-			drained = 2;
-		}
+		lru_cache_drain_for_folio(folio, pin_refs, &drained);

 		if (!folio_isolate_lru(folio))
 			continue;
diff --git a/mm/internal.h b/mm/internal.h
index f26423de4ca28..7aa787768353a 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -46,7 +46,7 @@ void lru_note_cost_unlock_irq(struct lruvec *lruvec, bool file,
 void lru_note_cost_refault(struct folio *folio);
 void folio_add_lru_vma(struct folio *folio, struct vm_area_struct *vma);

-static inline bool folio_may_be_lru_cached(struct folio *folio)
+static inline bool folio_may_be_lru_cached(const struct folio *folio)
 {
 	/*
 	 * Holding PMD-sized folios in per-CPU LRU cache unbalances accounting.
-- 
2.43.0


-- 
Cheers,

David