[PATCH 4/6] mm/hugetlb: add memory shrinker for cache

[Sourav Panda <souravpanda@google.com>]

Register a memory shrinker for the HugeTLB dynamic cache to
reclaim cached pages under memory pressure.

Implement hugetlb_shrinker_count to report the number of cached pages
(currently limited to gigantic pages in the shrinker) and
hugetlb_shrinker_scan to free them back to the buddy allocator.

The shrinker scan uses LRU policy to reclaim the coldest pages first.

Signed-off-by: Sourav Panda <souravpanda@google.com>
---
 mm/hugetlb.c | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index f18a3123cbcb..1bc0198a695c 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -4399,6 +4399,68 @@ ssize_t __max_cached_huge_pages_store_common(bool obey_mempolicy,
 
 	return err ? err : len;
 }
+
+static unsigned long hugetlb_shrinker_count(struct shrinker *shrink,
+					    struct shrink_control *sc)
+{
+	struct hstate *h;
+	unsigned long count = 0;
+	int nid = sc->nid;
+
+	for_each_hstate(h) {
+		count += h->nr_cached_hugepages_node[nid];
+	}
+	return count;
+}
+
+static unsigned long hugetlb_shrinker_scan(struct shrinker *shrink,
+					   struct shrink_control *sc)
+{
+	struct hstate *h;
+	int nid = sc->nid;
+	unsigned long freed = 0;
+	unsigned long to_scan = sc->nr_to_scan;
+
+	for_each_hstate(h) {
+		while (freed < to_scan) {
+			struct folio *folio = NULL;
+			bool atomic = true;
+
+			/*
+			 * Only do synchronous vmemmap restoration if the caller
+			 * context allows blocking and FS reclaim (GFP_KERNEL compatible).
+			 * Otherwise, defer to workqueue to avoid lock inversion.
+			 */
+			if (gfpflags_allow_blocking(sc->gfp_mask) &&
+			    (sc->gfp_mask & __GFP_FS))
+				atomic = false;
+
+			spin_lock_irq(&hugetlb_lock);
+			if (list_empty(&h->hugepage_cache_lists[nid])) {
+				spin_unlock_irq(&hugetlb_lock);
+				break;
+			}
+
+			folio = list_first_entry(&h->hugepage_cache_lists[nid],
+						 struct folio, lru);
+
+			remove_hugetlb_folio(h, folio, false);
+			spin_unlock_irq(&hugetlb_lock);
+
+			update_and_free_hugetlb_folio(h, folio, atomic);
+
+			freed++;
+			if (current->reclaim_state)
+				current->reclaim_state->reclaimed += pages_per_huge_page(h);
+		}
+		if (freed >= to_scan)
+			break;
+	}
+
+	return freed;
+}
+
+static struct shrinker *hugetlb_cached_pages_shrinker;
 #endif /* CONFIG_HUGETLB_CACHE */
 
 ssize_t __nr_hugepages_store_common(bool obey_mempolicy,
@@ -4502,6 +4564,21 @@ static int __init hugetlb_init(void)
 
 	for (i = 0; i < num_fault_mutexes; i++)
 		mutex_init(&hugetlb_fault_mutex_table[i]);
+
+#ifdef CONFIG_HUGETLB_CACHE
+	hugetlb_cached_pages_shrinker = shrinker_alloc(SHRINKER_NUMA_AWARE |
+						       SHRINKER_NONSLAB,
+						       "hugetlb-cached-pages");
+	if (!hugetlb_cached_pages_shrinker) {
+		pr_warn("HugeTLB: Failed to allocate shrinker\n");
+	} else {
+		hugetlb_cached_pages_shrinker->count_objects = hugetlb_shrinker_count;
+		hugetlb_cached_pages_shrinker->scan_objects = hugetlb_shrinker_scan;
+		hugetlb_cached_pages_shrinker->batch = 1;
+		shrinker_register(hugetlb_cached_pages_shrinker);
+	}
+#endif
+
 	return 0;
 }
 subsys_initcall(hugetlb_init);
-- 
2.55.0.rc0.799.gd6f94ed593-goog