[PATCH RFC 3/3] mm/ksm: scan VMAs with per-VMA locks

[Longlong Xia]

From: Longlong Xia <xialonglong@xxxxxxxxxx>

KSM currently holds mm->mmap_lock while scanning the VMAs of an mm. Use
lock_next_vma() and walk_page_range_vma() to acquire each VMA's read lock
and release it after finding a candidate page, allowing unrelated VMA
updates to proceed while ksmd scans.

Pin mm_users for the duration of the lockless VMA traversal to keep page
tables alive, and retain the mmap-lock scan as a fallback when a VMA lock
cannot be acquired or per-VMA locking is disabled. Allocation failures
preserve the existing early-stop semantics. This reduces mmap_lock
contention for workloads with concurrent address-space updates while
preserving existing scan semantics.

When the per-VMA scan completes without finding a candidate, remember the
complete traversal and revalidate the mmap-lock sequence before entering
the fallback. If the address space is unchanged, skip the duplicate VMA
walk; if a VMA writer raced with the traversal, restart the fallback walk.

Assisted-by: Zcode:GLM-5.3
Signed-off-by: Longlong Xia <xialonglong@xxxxxxxxxx>
---
 mm/ksm.c | 107 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 107 insertions(+)

diff --git a/mm/ksm.c b/mm/ksm.c
index aee1a1b49b1b..2c0b2adb83d4 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -2647,6 +2647,98 @@ static struct mm_walk_ops ksm_next_page_ops = {
 	.walk_lock = PGWALK_RDLOCK,
 };
 
+#ifdef CONFIG_PER_VMA_LOCK
+static const struct mm_walk_ops ksm_next_page_vma_ops = {
+	.pmd_entry = ksm_next_page_pmd_entry,
+	.walk_lock = PGWALK_VMA_RDLOCK_VERIFY,
+};
+
+/* Return true when a candidate was processed, including allocation failure. */
+static bool scan_get_next_rmap_item_vma(struct ksm_mm_slot *mm_slot,
+		struct page **page, struct ksm_rmap_item **result,
+		bool *complete, unsigned int *mm_wr_seq)
+{
+	struct mm_struct *mm = mm_slot->slot.mm;
+	unsigned long address = ksm_scan.address;
+	struct vm_area_struct *vma;
+	struct vma_iterator vmi;
+	bool done = false;
+	bool mmap_unlocked;
+
+	if (!mmget_not_zero(mm))
+		return false;
+
+	mmap_unlocked = mmap_lock_speculate_try_begin(mm, mm_wr_seq);
+
+	for (;;) {
+		rcu_read_lock();
+		vma_iter_init(&vmi, mm, address);
+		vma = lock_next_vma(mm, &vmi, address);
+		rcu_read_unlock();
+		if (IS_ERR_OR_NULL(vma))
+			break;
+
+		address = vma->vm_end;
+		if (!(vma->vm_flags & VM_MERGEABLE))
+			goto next_vma;
+		if (ksm_scan.address < vma->vm_start)
+			ksm_scan.address = vma->vm_start;
+		if (!vma->anon_vma)
+			ksm_scan.address = vma->vm_end;
+
+		while (ksm_scan.address < vma->vm_end) {
+			struct ksm_next_page_arg arg;
+			struct ksm_rmap_item *rmap_item;
+			int found;
+
+			found = walk_page_range_vma(vma, ksm_scan.address,
+					vma->vm_end, &ksm_next_page_vma_ops, &arg);
+			if (found <= 0) {
+				VM_WARN_ON_ONCE(found < 0);
+				ksm_scan.address = vma->vm_end;
+				break;
+			}
+
+			ksm_scan.address = arg.addr;
+			flush_anon_page(vma, arg.page, arg.addr);
+			flush_dcache_page(arg.page);
+			rmap_item = get_next_rmap_item(mm_slot,
+					ksm_scan.rmap_list, arg.addr);
+			if (rmap_item) {
+				ksm_scan.rmap_list = &rmap_item->rmap_list;
+				if (should_skip_rmap_item(arg.folio, rmap_item)) {
+					folio_put(arg.folio);
+					ksm_scan.address += PAGE_SIZE;
+					cond_resched();
+					continue;
+				}
+				ksm_scan.address += PAGE_SIZE;
+				*page = arg.page;
+			} else {
+				folio_put(arg.folio);
+			}
+			*result = rmap_item;
+			done = true;
+			vma_end_read(vma);
+			goto out;
+		}
+next_vma:
+		/*
+		 * Don't advance ksm_scan.address for VMAs the mmap-lock loop
+		 * skips with a plain continue: the cursor has to stay 0 when
+		 * this mm holds no VM_MERGEABLE vma, so the fallback walk can
+		 * remove the mm from the scan list at the end of the pass.
+		 */
+		vma_end_read(vma);
+		cond_resched();
+	}
+	*complete = mmap_unlocked && !vma;
+out:
+	mmput_async(mm);
+	return done;
+}
+#endif
+
 static struct ksm_rmap_item *scan_get_next_rmap_item(struct page **page)
 {
 	struct mm_struct *mm;
@@ -2655,6 +2747,9 @@ static struct ksm_rmap_item *scan_get_next_rmap_item(struct page **page)
 	struct vm_area_struct *vma;
 	struct ksm_rmap_item *rmap_item;
 	struct vma_iterator vmi;
+	bool skip_vma_scan = false;
+	bool vma_scan_complete = false;
+	unsigned int mm_wr_seq;
 	int nid;
 
 	if (list_empty(&ksm_mm_head.slot.mm_node))
@@ -2719,12 +2814,23 @@ static struct ksm_rmap_item *scan_get_next_rmap_item(struct page **page)
 
 	slot = &mm_slot->slot;
 	mm = slot->mm;
+#ifdef CONFIG_PER_VMA_LOCK
+	rmap_item = NULL;
+	if (scan_get_next_rmap_item_vma(mm_slot, page, &rmap_item,
+			&vma_scan_complete, &mm_wr_seq))
+		return rmap_item;
+#endif
+	/* Recheck the end of the scan under mmap_lock before removing the mm. */
 	vma_iter_init(&vmi, mm, ksm_scan.address);
 
 	mmap_read_lock(mm);
+	if (vma_scan_complete && !mmap_lock_speculate_retry(mm, mm_wr_seq))
+		skip_vma_scan = true;
 	if (ksm_test_exit(mm))
 		goto no_vmas;
 
+	if (skip_vma_scan)
+		goto scan_cleanup;
 	for_each_vma(vmi, vma) {
 		if (!(vma->vm_flags & VM_MERGEABLE))
 			continue;
@@ -2785,6 +2891,7 @@ static struct ksm_rmap_item *scan_get_next_rmap_item(struct page **page)
 		}
 	}
 
+scan_cleanup:
 	if (ksm_test_exit(mm)) {
 no_vmas:
 		ksm_scan.address = 0;
-- 
2.43.0