[RFC PATCH 16/18] pte_ref: add track_pte_{set, clear}() helper

[Qi Zheng <zhengqi.arch@bytedance.com>]

The track_pte_set() is used to track the setting of the PTE page table
entry, and the percpu_ref of the PTE page table page will be incremented
when the entry changes from pte_none() to !pte_none().

The track_pte_clear() is used to track the clearing of the PTE page
table entry, and the percpu_ref of the PTE page table page will be
decremented when the entry changes from !pte_none() to pte_none().

In this way, the usage of the PTE page table page can be tracked by
its percpu_ref.

Signed-off-by: Qi Zheng <zhengqi.arch@bytedance.com>
---
 include/linux/pte_ref.h | 14 ++++++++++++++
 mm/pte_ref.c            | 30 ++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/include/linux/pte_ref.h b/include/linux/pte_ref.h
index 379c3b45a6ab..6ab740e1b989 100644
--- a/include/linux/pte_ref.h
+++ b/include/linux/pte_ref.h
@@ -18,6 +18,10 @@ void __pte_put(pgtable_t page);
 void pte_put(pte_t *ptep);
 void try_to_free_user_pte(struct mm_struct *mm, pmd_t *pmd, unsigned long addr,
 			  bool switch_back);
+void track_pte_set(struct mm_struct *mm, unsigned long addr, pte_t *ptep,
+		   pte_t pte);
+void track_pte_clear(struct mm_struct *mm, unsigned long addr, pte_t *ptep,
+		     pte_t pte);
 
 #else /* !CONFIG_FREE_USER_PTE */
 
@@ -54,6 +58,16 @@ static inline void try_to_free_user_pte(struct mm_struct *mm, pmd_t *pmd,
 {
 }
 
+static inline void track_pte_set(struct mm_struct *mm, unsigned long addr,
+				 pte_t *ptep, pte_t pte)
+{
+}
+
+static inline void track_pte_clear(struct mm_struct *mm, unsigned long addr,
+				   pte_t *ptep, pte_t pte)
+{
+}
+
 #endif /* CONFIG_FREE_USER_PTE */
 
 #endif /* _LINUX_PTE_REF_H */
diff --git a/mm/pte_ref.c b/mm/pte_ref.c
index bf9629272c71..e92510deda0b 100644
--- a/mm/pte_ref.c
+++ b/mm/pte_ref.c
@@ -197,4 +197,34 @@ void try_to_free_user_pte(struct mm_struct *mm, pmd_t *pmd, unsigned long addr,
 	}
 }
 
+void track_pte_set(struct mm_struct *mm, unsigned long addr, pte_t *ptep,
+		   pte_t pte)
+{
+	pgtable_t page;
+
+	if (&init_mm == mm || pte_huge(pte))
+		return;
+
+	page = pte_to_page(ptep);
+	BUG_ON(percpu_ref_is_zero(page->pte_ref));
+	if (pte_none(*ptep) && !pte_none(pte))
+		percpu_ref_get(page->pte_ref);
+}
+EXPORT_SYMBOL(track_pte_set);
+
+void track_pte_clear(struct mm_struct *mm, unsigned long addr, pte_t *ptep,
+		     pte_t pte)
+{
+	pgtable_t page;
+
+	if (&init_mm == mm || pte_huge(pte))
+		return;
+
+	page = pte_to_page(ptep);
+	BUG_ON(percpu_ref_is_zero(page->pte_ref));
+	if (!pte_none(pte))
+		percpu_ref_put(page->pte_ref);
+}
+EXPORT_SYMBOL(track_pte_clear);
+
 #endif /* CONFIG_FREE_USER_PTE */
-- 
2.20.1