[RFC PATCH 5/6] mm, pgtable: Add the reference counter for COW PTE

[Chih-En Lin <shiyn.lin@xxxxxxxxx>]

Add the reference counter cow_pgtable_refcount to maintain the number
of process references to COW PTE. Before decreasing the reference
count, it will check whether the counter is one or not for reusing
COW PTE when the counter is one.

Signed-off-by: Chih-En Lin <shiyn.lin@xxxxxxxxx>
---
 include/linux/mm.h       |  1 +
 include/linux/mm_types.h |  1 +
 include/linux/pgtable.h  | 27 +++++++++++++++++++++++++++
 mm/memory.c              |  1 +
 4 files changed, 30 insertions(+)

diff --git a/include/linux/mm.h b/include/linux/mm.h
index 221926a3d818..e48bb3fbc33c 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -2329,6 +2329,7 @@ static inline bool pgtable_pte_page_ctor(struct page *page)
 	__SetPageTable(page);
 	inc_lruvec_page_state(page, NR_PAGETABLE);
 	page->cow_pte_owner = NULL;
+	atomic_set(&page->cow_pgtable_refcount, 1);
 	return true;
 }
 
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 5dcbd7f6c361..984d81e47d53 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -221,6 +221,7 @@ struct page {
 #ifdef LAST_CPUPID_NOT_IN_PAGE_FLAGS
 	int _last_cpupid;
 #endif
+	atomic_t cow_pgtable_refcount; /* COW page table */
 	pmd_t *cow_pte_owner; /* cow pte: pmd */
 } _struct_page_alignment;
 
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index faca57af332e..33c01fec7b92 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -604,6 +604,33 @@ static inline bool cow_pte_owner_is_same(pmd_t *pmd, pmd_t *owner)
 		true : false;
 }
 
+extern void cow_pte_fallback(struct vm_area_struct *vma, pmd_t *pmd,
+		unsigned long addr);
+
+static inline int pmd_get_pte(pmd_t *pmd)
+{
+	return atomic_inc_return(&pmd_page(*pmd)->cow_pgtable_refcount);
+}
+
+/* If the COW PTE page->cow_pgtable_refcount is 1, instead of decreasing the
+ * counter, clear write protection of the corresponding PMD entry and reset
+ * the COW PTE owner to reuse the table.
+ */
+static inline int pmd_put_pte(struct vm_area_struct *vma, pmd_t *pmd,
+		unsigned long addr)
+{
+	if (!atomic_add_unless(&pmd_page(*pmd)->cow_pgtable_refcount, -1, 1)) {
+		cow_pte_fallback(vma, pmd, addr);
+		return 1;
+	}
+	return 0;
+}
+
+static inline int cow_pte_refcount_read(pmd_t *pmd)
+{
+	return atomic_read(&pmd_page(*pmd)->cow_pgtable_refcount);
+}
+
 #ifndef pte_access_permitted
 #define pte_access_permitted(pte, write) \
 	(pte_present(pte) && (!(write) || pte_write(pte)))
diff --git a/mm/memory.c b/mm/memory.c
index dcb678cbb051..aa66af76e214 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -4597,6 +4597,7 @@ void cow_pte_fallback(struct vm_area_struct *vma, pmd_t *pmd,
 	pmd_t new;
 
 	BUG_ON(pmd_write(*pmd));
+	BUG_ON(cow_pte_refcount_read(pmd) != 1);
 
 	start = addr & PMD_MASK;
 	end = (addr + PMD_SIZE) & PMD_MASK;
-- 
2.36.1