[RFC PATCH 4/4] mm: adjust free_pages to make __zone_watermark_ok() mTHP-aware

[Bo Zhang]

To improve the mTHP allocation success rate and reduce
fragmentation over time, introduce
zone_effective_free_pages() for mTHP-aware free page accounting.

Refactor the free page counting logic used by pgdat_balanced() and
compaction_suit_allocation_order() into a shared helper function
zone_effective_free_pages().

The function selects the appropriate free page metric based on context:
- defrag_mode: use NR_FREE_PAGES_BLOCKS (whole pageblock accounting)
- mTHP always-enabled: count only free pages in buddy blocks >= the
  minimum always-enabled mTHP order, since smaller fragments cannot
  satisfy mTHP allocations
- otherwise: use NR_FREE_PAGES

Signed-off-by: Bo Zhang <zhangbo56@xxxxxxxxxx>
---
 mm/compaction.c | 38 ++++++++++++++++++++++++++++++++++----
 mm/internal.h   |  3 +++
 mm/vmscan.c     | 23 +++++------------------
 3 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/mm/compaction.c b/mm/compaction.c
index 15b92475562a..29be72597415 100644
--- a/mm/compaction.c
+++ b/mm/compaction.c
@@ -2528,6 +2528,38 @@ bool compaction_zonelist_suitable(struct alloc_context *ac, int order,
 	return false;
 }
 
+/**
+ * zone_effective_free_pages - get free pages relevant to allocation order
+ * @zone:       target zone
+ * @order:      allocation order
+ * @use_blocks: if true, use NR_FREE_PAGES_BLOCKS
+ *
+ * In defrag_mode, watermarks must be met in whole blocks to avoid
+ * polluting allocator fallbacks. kswapd usually cannot accomplish
+ * this on its own and needs kcompactd support.
+ *
+ * When mTHP always-enabled orders are configured, count only free pages
+ * in blocks >= min mTHP order, as smaller fragments cannot satisfy mTHP
+ * allocations.
+ */
+unsigned long zone_effective_free_pages(struct zone *zone,
+					unsigned int order,
+					bool use_blocks)
+{
+	if (use_blocks)
+		return zone_page_state(zone, NR_FREE_PAGES_BLOCKS);
+
+	if (READ_ONCE(huge_anon_orders_always) && order == compact_hpage_order()) {
+		unsigned long free_pages = 0;
+
+		for (int o = order; o < NR_PAGE_ORDERS; o++)
+			free_pages += zone->free_area[o].nr_free << o;
+		return free_pages;
+	}
+
+	return zone_page_state(zone, NR_FREE_PAGES);
+}
+
 /*
  * Should we do compaction for target allocation order.
  * Return COMPACT_SUCCESS if allocation for target order can be already
@@ -2543,10 +2575,8 @@ compaction_suit_allocation_order(struct zone *zone, unsigned int order,
 	unsigned long free_pages;
 	unsigned long watermark;
 
-	if (kcompactd && defrag_mode)
-		free_pages = zone_page_state(zone, NR_FREE_PAGES_BLOCKS);
-	else
-		free_pages = zone_page_state(zone, NR_FREE_PAGES);
+	free_pages = zone_effective_free_pages(zone, order,
+					       kcompactd && defrag_mode);
 
 	watermark = wmark_pages(zone, alloc_flags & ALLOC_WMARK_MASK);
 	if (__zone_watermark_ok(zone, order, watermark, highest_zoneidx,
diff --git a/mm/internal.h b/mm/internal.h
index 38b1165212c9..14bb9543879b 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1654,4 +1654,7 @@ static inline bool can_spin_trylock(void)
 	return true;
 }
 
+unsigned long zone_effective_free_pages(struct zone *zone,
+					unsigned int order,
+					bool use_blocks);
 #endif	/* __MM_INTERNAL_H */
diff --git a/mm/vmscan.c b/mm/vmscan.c
index c1404a59523d..a419a2c2fca4 100644
--- a/mm/vmscan.c
+++ b/mm/vmscan.c
@@ -6966,7 +6966,6 @@ static bool pgdat_balanced(pg_data_t *pgdat, int order, int highest_zoneidx)
 	 * meet watermarks.
 	 */
 	for_each_managed_zone_pgdat(zone, pgdat, i, highest_zoneidx) {
-		enum zone_stat_item item;
 		unsigned long free_pages;
 
 		if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
@@ -6974,21 +6973,6 @@ static bool pgdat_balanced(pg_data_t *pgdat, int order, int highest_zoneidx)
 		else
 			mark = high_wmark_pages(zone);
 
-		/*
-		 * In defrag_mode, watermarks must be met in whole
-		 * blocks to avoid polluting allocator fallbacks.
-		 *
-		 * However, kswapd usually cannot accomplish this on
-		 * its own and needs kcompactd support. Once it's
-		 * reclaimed a compaction gap, and kswapd_shrink_node
-		 * has dropped order, simply ensure there are enough
-		 * base pages for compaction, wake kcompactd & sleep.
-		 */
-		if (defrag_mode && order)
-			item = NR_FREE_PAGES_BLOCKS;
-		else
-			item = NR_FREE_PAGES;
-
 		/*
 		 * When there is a high number of CPUs in the system,
 		 * the cumulative error from the vmstat per-cpu cache
@@ -7001,9 +6985,12 @@ static bool pgdat_balanced(pg_data_t *pgdat, int order, int highest_zoneidx)
 		 * counter won't actually be per-cpu cached. But keep
 		 * things simple for now; revisit when somebody cares.
 		 */
-		free_pages = zone_page_state(zone, item);
+		free_pages = zone_effective_free_pages(zone, order,
+						       defrag_mode & order);
 		if (zone->percpu_drift_mark && free_pages < zone->percpu_drift_mark)
-			free_pages = zone_page_state_snapshot(zone, item);
+			free_pages = zone_page_state_snapshot(zone,
+					defrag_mode & order ?
+					NR_FREE_PAGES_BLOCKS : NR_FREE_PAGES);
 
 		if (__zone_watermark_ok(zone, order, mark, highest_zoneidx,
 					0, free_pages))
-- 
2.34.1