When approaching the end of an rmap_walk traversal, we may need to
perform some resource statistics or cleanup work, and these operations
need to be done under lock.

This patch adds a new exit hook to rmap_walk_control.

Note that currently, the exit hook only be used in the anon and file
rmap_walk implementations, as ksm no need it.

Signed-off-by: Huan Yang <link@vivo.com>
---
 include/linux/rmap.h | 2 ++
 mm/rmap.c            | 6 ++++++
 2 files changed, 8 insertions(+)

diff --git a/include/linux/rmap.h b/include/linux/rmap.h
index a305811d6310..594dc4df3bfb 100644
--- a/include/linux/rmap.h
+++ b/include/linux/rmap.h
@@ -999,6 +999,7 @@ void remove_migration_ptes(struct folio *src, struct folio *dst, int flags);
  * locked: already locked before invoke rmap_walk
  * rmap_one: executed on each vma where page is mapped
  * done: for checking traversing termination condition
+ * exit: do some clean work below lock before leave rmap_walk
  * anon_lock: for getting anon_lock by optimized way rather than default
  * invalid_vma: for skipping uninterested vma
  */
@@ -1014,6 +1015,7 @@ struct rmap_walk_control {
 	bool (*rmap_one)(struct folio *folio, struct vm_area_struct *vma,
 					unsigned long addr, void *arg);
 	int (*done)(struct folio *folio, struct rmap_walk_control *rwc);
+	void (*exit)(struct folio *folio, struct rmap_walk_control *rwc);
 	struct anon_vma *(*anon_lock)(const struct folio *folio,
 				      struct rmap_walk_control *rwc);
 	bool (*invalid_vma)(struct vm_area_struct *vma, void *arg);
diff --git a/mm/rmap.c b/mm/rmap.c
index e590f4071eca..66b48ab192f5 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2832,6 +2832,9 @@ static void rmap_walk_anon(struct folio *folio, struct rmap_walk_control *rwc)
 			break;
 	}
 
+	if (rwc->exit)
+		rwc->exit(folio, rwc);
+
 	if (!locked)
 		anon_vma_unlock_read(anon_vma);
 }
@@ -2895,6 +2898,9 @@ static void __rmap_walk_file(struct folio *folio, struct address_space *mapping,
 			goto done;
 	}
 done:
+	if (rwc->exit)
+		rwc->exit(folio, rwc);
+
 	if (!locked)
 		i_mmap_unlock_read(mapping);
 }
-- 
2.34.1