[PATCH 4/4] mm/mglru: invalidate empty-walk skip on page fault and migration

[Baoquan He <baoquan.he@linux.dev>]

empty_map skips an mm on a node for up to K generations after an empty
walk, leaving a window where pages that appear on that node are not aged
until the forced rescan. Notify MGLRU when a page of the mm appears:
set the node's bitmap bit and clear its empty_map bit. Fault and
migration are both software events, so the invalidation is complete - a
node marked empty has no pages there, so a page can only appear via a
fault or a migration.

- mm/memory.c do_anonymous_page()/finish_fault(): mark the folio's node
  after the page is allocated (anon and file/COW faults).
- mm/migrate.c remove_migration_pte(): mark the folio's destination node
  (also covers NUMA-balancing migration).

The bitmap bit is mostly redundant with schedule-time marking, but is
needed when migration targets an idle mm whose bits were already
cleared by a previous walk.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 include/linux/mm_types.h | 18 ++++++++++++++++++
 mm/memory.c              |  7 +++++++
 mm/migrate.c             |  3 +++
 3 files changed, 28 insertions(+)

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 68ec8bb2ab71..89b723483675 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -1531,6 +1531,20 @@ static inline void lru_gen_use_mm(struct mm_struct *mm)
 	WRITE_ONCE(mm->lru_gen.bitmap, -1);
 }
 
+/*
+ * A page of this mm appeared on (or was accessed on) node @nid — e.g. a page
+ * fault or a migration. Set that node's bitmap bit so the aging walker walks
+ * the mm, and clear the empty-walk skip so a page that just appeared on a node
+ * previously marked empty is not ignored for up to K generations.
+ */
+static inline void lru_gen_mm_accessed(struct mm_struct *mm, int nid)
+{
+	unsigned long key = nid % BITS_PER_TYPE(mm->lru_gen.bitmap);
+
+	set_bit(key, &mm->lru_gen.bitmap);
+	clear_bit(key, &mm->lru_gen.empty_map);
+}
+
 #else /* !CONFIG_LRU_GEN_WALKS_MMU */
 
 static inline void lru_gen_add_mm(struct mm_struct *mm)
@@ -1553,6 +1567,10 @@ static inline void lru_gen_use_mm(struct mm_struct *mm)
 {
 }
 
+static inline void lru_gen_mm_accessed(struct mm_struct *mm, int nid)
+{
+}
+
 #endif /* CONFIG_LRU_GEN_WALKS_MMU */
 
 struct vma_iterator {
diff --git a/mm/memory.c b/mm/memory.c
index 428eb555ecb7..e3d7c8f7ca7d 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -5512,6 +5512,9 @@ static vm_fault_t do_anonymous_page(struct vm_fault *vmf)
 		folio_put(folio);
 		return handle_userfault(vmf, VM_UFFD_MISSING);
 	}
+	/* a new page of this mm lands on this node: invalidate any empty skip */
+	lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	map_anon_folio_pte_pf(folio, vmf->pte, vma, addr,
 			      vmf_orig_pte_uffd_wp(vmf));
 unlock:
@@ -5772,6 +5775,10 @@ vm_fault_t finish_fault(struct vm_fault *vmf)
 		page = vmf->page;
 
 	folio = page_folio(page);
+
+	/* mapping a page of this mm on this node: invalidate any empty skip */
+	lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	/*
 	 * check even for read faults because we might have lost our CoWed
 	 * page
diff --git a/mm/migrate.c b/mm/migrate.c
index b937cbd76480..2e0674e89ba8 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -354,6 +354,9 @@ static bool remove_migration_pte(struct folio *folio,
 	struct rmap_walk_arg *rmap_walk_arg = arg;
 	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
 
+	/* the folio ends up on folio_nid(): notify MGLRU for this mm */
+	lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	while (page_vma_mapped_walk(&pvmw)) {
 		rmap_t rmap_flags = RMAP_NONE;
 		pte_t old_pte;
-- 
2.54.0