[PATCH v9 14/41] mm: swap: Introduce lru_add_drain_progressive()

[Ackerley Tng via B4 Relay]

From: Ackerley Tng <ackerleytng@xxxxxxxxxx>

Extract the progressive LRU drain retry logic from
collect_longterm_unpinnable_folios() into a reusable helper,
lru_add_drain_progressive().

When attempting to isolate folios that may still reside in per-CPU folio
batches, draining is escalated progressively:

1. State 0: Call lru_add_drain() to flush local CPU batches.
2. State 1: Call lru_add_drain_all() to flush all CPU batches.
3. State >= 2: Return false to stop retrying.

Refactor collect_longterm_unpinnable_folios() to use this new helper.

The helper will be used by KVM's guest_memfd in a later patch.

Signed-off-by: Ackerley Tng <ackerleytng@xxxxxxxxxx>
---
 include/linux/swap.h |  2 ++
 mm/gup.c             | 19 ++++++-------------
 mm/swap.c            | 15 +++++++++++++++
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/include/linux/swap.h b/include/linux/swap.h
index 8f0f68e245baa..cd54f73f34f39 100644
--- a/include/linux/swap.h
+++ b/include/linux/swap.h
@@ -344,6 +344,8 @@ extern void lru_add_drain(void);
 extern void lru_add_drain_cpu(int cpu);
 extern void lru_add_drain_cpu_zone(struct zone *zone);
 extern void lru_add_drain_all(void);
+bool lru_add_drain_progressive(int *drain_state);
+
 void folio_deactivate(struct folio *folio);
 void folio_mark_lazyfree(struct folio *folio);
 extern void swap_setup(void);
diff --git a/mm/gup.c b/mm/gup.c
index 0692119b79043..5f00435e2c635 100644
--- a/mm/gup.c
+++ b/mm/gup.c
@@ -2268,7 +2268,7 @@ static unsigned long collect_longterm_unpinnable_folios(
 {
 	unsigned long collected = 0;
 	struct folio *folio;
-	int drained = 0;
+	int drain_state = 0;
 	long i = 0;
 
 	for (folio = pofs_get_folio(pofs, i); folio;
@@ -2287,18 +2287,11 @@ static unsigned long collect_longterm_unpinnable_folios(
 			continue;
 		}
 
-		if (drained == 0 && folio_may_be_lru_cached(folio) &&
-				folio_ref_count(folio) !=
-				folio_expected_ref_count(folio) + 1) {
-			lru_add_drain();
-			drained = 1;
-		}
-		if (drained == 1 && folio_may_be_lru_cached(folio) &&
-				folio_ref_count(folio) !=
-				folio_expected_ref_count(folio) + 1) {
-			lru_add_drain_all();
-			drained = 2;
-		}
+		while (folio_may_be_lru_cached(folio) &&
+		       folio_ref_count(folio) !=
+		       folio_expected_ref_count(folio) + 1 &&
+		       lru_add_drain_progressive(&drain_state))
+			;
 
 		if (!folio_isolate_lru(folio))
 			continue;
diff --git a/mm/swap.c b/mm/swap.c
index 588f50d8f1a8c..0f9465d31fe52 100644
--- a/mm/swap.c
+++ b/mm/swap.c
@@ -950,6 +950,21 @@ void lru_add_drain_all(void)
 }
 #endif /* CONFIG_SMP */
 
+bool lru_add_drain_progressive(int *drain_state)
+{
+	if (*drain_state == 0) {
+		lru_add_drain();
+		*drain_state = 1;
+		return true;
+	}
+	if (*drain_state == 1) {
+		lru_add_drain_all();
+		*drain_state = 2;
+		return true;
+	}
+	return false;
+}
+
 atomic_t lru_disable_count = ATOMIC_INIT(0);
 
 /*

-- 
2.55.0.508.g3f0d502094-goog