[RFC PATCH 5/7] 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 track_pte_clear() is used to track the clearing of the
PTE page table entry, we update the pte_refcount of the PTE page in
these two functions.

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

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

diff --git a/include/linux/pte_ref.h b/include/linux/pte_ref.h
index db14e03e1dff..ab49c7fac120 100644
--- a/include/linux/pte_ref.h
+++ b/include/linux/pte_ref.h
@@ -12,12 +12,25 @@
 
 void pte_ref_init(pgtable_t pte);
 
+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 */
 
 static inline void pte_ref_init(pgtable_t pte)
 {
 }
 
+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 12b27646e88c..818821d068af 100644
--- a/mm/pte_ref.c
+++ b/mm/pte_ref.c
@@ -69,4 +69,40 @@ void pte_ref_init(pgtable_t pte)
 	pte->pte_refcount = 0;
 }
 
+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);
+	if (pte_none(*ptep) && !pte_none(pte)) {
+		pte_refcount_add(mm, page, PTE_MAPPED_OFFSET);
+		if (is_zero_pfn(pte_pfn(pte)))
+			pte_refcount_add(mm, page, PTE_ZERO_OFFSET);
+	} else if (is_zero_pfn(pte_pfn(*ptep)) && !is_zero_pfn(pte_pfn(pte))) {
+		pte_refcount_sub(mm, page, PTE_ZERO_OFFSET);
+	}
+}
+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);
+	if (!pte_none(pte)) {
+		pte_refcount_sub(mm, page, PTE_MAPPED_OFFSET);
+		if (is_zero_pfn(pte_pfn(pte)))
+			pte_refcount_sub(mm, page, PTE_ZERO_OFFSET);
+	}
+}
+EXPORT_SYMBOL(track_pte_clear);
+
 #endif /* CONFIG_FREE_USER_PTE */
-- 
2.20.1