[RFC PATCH 7/7] mm: add proc interface to free user PTE page table pages

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

Add /proc/sys/vm/free_ptes file to procfs, when pid is written
to the file, we will traverse its process address space, find
and free empty PTE pages or zero PTE pages.

Signed-off-by: Qi Zheng <zhengqi.arch@bytedance.com>
---
 include/linux/pte_ref.h |   5 ++
 kernel/sysctl.c         |  12 ++++
 mm/pte_ref.c            | 126 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 143 insertions(+)

diff --git a/include/linux/pte_ref.h b/include/linux/pte_ref.h
index ab49c7fac120..f7e244129291 100644
--- a/include/linux/pte_ref.h
+++ b/include/linux/pte_ref.h
@@ -16,6 +16,11 @@ 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);
+
+int free_ptes_sysctl_handler(struct ctl_table *table, int write,
+		void *buffer, size_t *length, loff_t *ppos);
+extern int sysctl_free_ptes_pid;
+
 #else /* !CONFIG_FREE_USER_PTE */
 
 static inline void pte_ref_init(pgtable_t pte)
diff --git a/kernel/sysctl.c b/kernel/sysctl.c
index 35d034219513..14e1a9841cb8 100644
--- a/kernel/sysctl.c
+++ b/kernel/sysctl.c
@@ -64,6 +64,7 @@
 #include <linux/mount.h>
 #include <linux/userfaultfd_k.h>
 #include <linux/pid.h>
+#include <linux/pte_ref.h>
 
 #include "../lib/kstrtox.h"
 
@@ -2153,6 +2154,17 @@ static struct ctl_table vm_table[] = {
 		.extra1		= SYSCTL_ONE,
 		.extra2		= SYSCTL_FOUR,
 	},
+#ifdef CONFIG_FREE_USER_PTE
+	{
+		.procname	= "free_ptes",
+		.data		= &sysctl_free_ptes_pid,
+		.maxlen		= sizeof(int),
+		.mode		= 0200,
+		.proc_handler	= free_ptes_sysctl_handler,
+		.extra1		= SYSCTL_ZERO,
+		.extra2		= SYSCTL_INT_MAX,
+	},
+#endif
 #ifdef CONFIG_COMPACTION
 	{
 		.procname	= "compact_memory",
diff --git a/mm/pte_ref.c b/mm/pte_ref.c
index 818821d068af..e7080a3100a6 100644
--- a/mm/pte_ref.c
+++ b/mm/pte_ref.c
@@ -6,6 +6,14 @@
  */
 #include <linux/pgtable.h>
 #include <linux/pte_ref.h>
+#include <linux/mm.h>
+#include <linux/pagewalk.h>
+#include <linux/sched/mm.h>
+#include <linux/jump_label.h>
+#include <linux/hugetlb.h>
+#include <asm/tlbflush.h>
+
+#include "internal.h"
 
 #ifdef CONFIG_FREE_USER_PTE
 
@@ -105,4 +113,122 @@ void track_pte_clear(struct mm_struct *mm, unsigned long addr, pte_t *ptep,
 }
 EXPORT_SYMBOL(track_pte_clear);
 
+#ifdef CONFIG_DEBUG_VM
+void pte_free_debug(pmd_t pmd)
+{
+	pte_t *ptep = (pte_t *)pmd_page_vaddr(pmd);
+	int i = 0;
+
+	for (i = 0; i < PTRS_PER_PTE; i++, ptep++) {
+		pte_t pte = *ptep;
+		BUG_ON(!(pte_none(pte) || is_zero_pfn(pte_pfn(pte))));
+	}
+}
+#else
+static inline void pte_free_debug(pmd_t pmd)
+{
+}
+#endif
+
+
+static int kfreeptd_pmd_entry(pmd_t *pmd, unsigned long addr,
+			      unsigned long next, struct mm_walk *walk)
+{
+	pmd_t pmdval;
+	pgtable_t page;
+	struct mm_struct *mm = walk->mm;
+	struct vm_area_struct vma = TLB_FLUSH_VMA(mm, 0);
+	spinlock_t *ptl;
+	bool free = false;
+	unsigned long haddr = addr & PMD_MASK;
+
+	if (pmd_trans_unstable(pmd))
+		goto out;
+
+	mmap_read_unlock(mm);
+	mmap_write_lock(mm);
+
+	if (mm_find_pmd(mm, addr) != pmd)
+		goto unlock_out;
+
+	ptl = pmd_lock(mm, pmd);
+	pmdval = *pmd;
+	if (pmd_none(pmdval) || pmd_leaf(pmdval)) {
+		spin_unlock(ptl);
+		goto unlock_out;
+	}
+	page = pmd_pgtable(pmdval);
+	if (!pte_mapped_count(page) || pte_zero_count(page) == PTRS_PER_PTE) {
+		pmd_clear(pmd);
+		flush_tlb_range(&vma, haddr, haddr + PMD_SIZE);
+		free = true;
+	}
+	spin_unlock(ptl);
+
+unlock_out:
+	mmap_write_unlock(mm);
+	mmap_read_lock(mm);
+
+	if (free) {
+		pte_free_debug(pmdval);
+		mm_dec_nr_ptes(mm);
+		pgtable_pte_page_dtor(page);
+		__free_page(page);
+	}
+
+out:
+	cond_resched();
+	return 0;
+}
+
+static const struct mm_walk_ops kfreeptd_walk_ops = {
+	.pmd_entry		= kfreeptd_pmd_entry,
+};
+
+int sysctl_free_ptes_pid;
+int free_ptes_sysctl_handler(struct ctl_table *table, int write,
+		void *buffer, size_t *length, loff_t *ppos)
+{
+	int ret;
+
+	ret = proc_dointvec_minmax(table, write, buffer, length, ppos);
+	if (ret)
+		return ret;
+	if (write) {
+		struct task_struct *task;
+		struct mm_struct *mm;
+
+		rcu_read_lock();
+		task = find_task_by_vpid(sysctl_free_ptes_pid);
+		if (!task) {
+			rcu_read_unlock();
+			return -ESRCH;
+		}
+		mm = get_task_mm(task);
+		rcu_read_unlock();
+
+		if (!mm) {
+			mmput(mm);
+			return -ESRCH;
+		}
+
+		do {
+			ret = -EBUSY;
+
+			if (mmap_read_trylock(mm)) {
+				ret = walk_page_range(mm, FIRST_USER_ADDRESS,
+						      ULONG_MAX,
+						      &kfreeptd_walk_ops, NULL);
+
+				mmap_read_unlock(mm);
+			}
+
+			cond_resched();
+		} while (ret == -EAGAIN);
+
+		mmput(mm);
+	}
+	return ret;
+}
+
 #endif /* CONFIG_FREE_USER_PTE */
-- 
2.20.1