[PATCH v2 0/4] mm/ksm: use per-VMA locking to reduce mmap_lock contention

[PATCH v2 0/4] mm/ksm: use per-VMA locking to reduce mmap_lock contention

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

KSM scans VM_MERGEABLE VMAs and currently protects each scan with
mmap_read_lock(). Per-VMA locking allows KSM to read-lock only the VMA
it is actually interested in, so that unrelated mmap()/munmap() activity
in the same mm no longer blocks ksmd.

This series is organized as follows:

  Patch 1 removes an unused 'vma' member from struct folio_walk.  It has
  never been used since its introduction and is pure cleanup.

  Patch 2 adds a 'walk_lock' member to struct folio_walk and extends
  folio_walk_start() to assert the required locking mode.  Existing
  callers are converted to pass PGWALK_RDLOCK, so there is no functional
  change.  This prepares folio_walk_start() for callers that hold a
  per-VMA read lock instead of mmap_read_lock(), which is needed by the
  Patch 4. No functional change.

  Patch 3 replaces the boolean 'lock_vma' argument of break_ksm() with a
  KSM-local enum ksm_walk_lock and builds the mm_walk_ops on the stack,
  so the function-pointer table stays const and there is no shared
  mutable state.  No functional change.

  Patch 4 introduces find_mergeable_vma_locked(), which uses the
  universal per-VMA locking helper vma_start_read_unlocked() to look up
  and read-lock a VM_MERGEABLE VMA without taking mmap_read_lock().  All
  KSM call sites that previously used find_mergeable_vma() under
  mmap_read_lock() are converted to the new helper, and the locking in
  get_mergeable_page() is switched to PGWALK_VMA_RDLOCK_VERIFY so that
  folio_walk_start() can verify the per-VMA lock is held.  Because ksmd
  no longer holds mmap_read_lock() while walking, __ksm_exit() must also
  write-lock every VMA (in addition to mmap_write_lock()) to keep
  serializing against an in-flight ksmd walk before the page tables are
  freed.

A microbenchmark was run to measure the time KSM takes to merge a
victim region under mmap_lock contention. Under interference from 4 churner
threads, the merge time of the per-VMA KSM-optimized kernel is
significantly reduced by 50%.

Changelog:
  v2:
  - Patch 1: rewrote the commit message to explain the introducing
    commit, how the VMA is actually used, and the audit performed;
    dropped the Fixes: tag.

  - Patch 3: 
    * replaced the boolean 'lock_vma' with a KSM-local enum
    ksm_walk_lock instead of exposing the internal enum page_walk_lock
    as a parameter;
    * build the mm_walk_ops on the stack rather than
    mutating a global struct, so the function-pointer table stays const
    and there is no shared mutable state; rewrote the commit message.

  - Patch 4:
    * Fixed a per-VMA read-lock leak in find_mergeable_vma_locked()
    when the VMA is not mergeable;
    made __ksm_exit() write-lock every VMA (vma_start_write()) in
    addition to mmap_write_lock(), so that it still serializes against a
    ksmd walk that holds only a per-VMA read lock.

Xu Xin (4):
  mm/pagewalk: delete the unused member
  mm: make folio_walk_start()'s locking asserts scalable
  mm/ksm: make break_ksm() more scalable
  ksm: add find_mergeable_vma_locked() to use per-VMA locking

 arch/s390/mm/fault.c     |   4 +-
 include/linux/pagewalk.h |   2 +-
 kernel/events/uprobes.c  |   4 +-
 mm/huge_memory.c         |   4 +-
 mm/ksm.c                 | 124 +++++++++++++++++++++++++++------------
 mm/migrate.c             |   8 ++-
 mm/pagewalk.c            |  10 +++-
 mm/rmap.c                |   4 +-
 8 files changed, 116 insertions(+), 44 deletions(-)

--
2.34.1

[PATCH v2 1/4] mm/pagewalk: delete the unused member

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

The 'vma' member of struct folio_walk has never been used since it was
introduced by commit aa39ca6940f1a ("mm/pagewalk: introduce
folio_walk_start() + folio_walk_end()").

The VMA is always passed to folio_walk_start() as a separate argument,
and neither folio_walk_start() nor folio_walk_end() ever reads or writes
fw->vma. I audited all the current users of struct folio_walk (the
callers of folio_walk_start()) and none of them assign to or read the
member, so it is safe to remove.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 include/linux/pagewalk.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
index b41d7265c01b..1c397be0d092 100644
--- a/include/linux/pagewalk.h
+++ b/include/linux/pagewalk.h
@@ -183,7 +183,6 @@ struct folio_walk {
 		pmd_t pmd;
 	};
 	/* private */
-	struct vm_area_struct *vma;
 	spinlock_t *ptl;
 };
 
-- 
2.34.1

Re: [PATCH v2 1/4] mm/pagewalk: delete the unused member

[sashiko-bot]

> From: Xu Xin <xu.xin@linux.dev>
> 
> The 'vma' member of struct folio_walk has never been used since it was
> introduced by commit aa39ca6940f1a ("mm/pagewalk: introduce
> folio_walk_start() + folio_walk_end()").
> 
> The VMA is always passed to folio_walk_start() as a separate argument,
> and neither folio_walk_start() nor folio_walk_end() ever reads or writes
> fw->vma. I audited all the current users of struct folio_walk (the
> callers of folio_walk_start()) and none of them assign to or read the
> member, so it is safe to remove.
> 
> Signed-off-by: Xu Xin <xu.xin@linux.dev>

Sashiko has reviewed this patch and found no issues.

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261102132416.300084-1-xu.xin@linux.dev?part=1

[PATCH v2 2/4] mm: make folio_walk_start()'s locking asserts scalable

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

Add an additional member 'walk_lock' to folio_walk to indicate
locking requirements for the walk. Similar to commit 49b0638502da0
("mm: enable page walking API to lock vmas during the walk").
But no change is made on any existing locking behavior,
all existing folio_walk_start() callers are currently still
under mmap_read_lock() protection.

This change is prepared for the later patch to enable VMA locking
asserts. folio_walk_start now operate under write-locked mmap_lock.
With introduction of vma locks at the next patch, the vmas have to
be locked as well during such walks to prevent concurrent page
faults in these areas.

No functional change intended.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 arch/s390/mm/fault.c     |  4 +++-
 include/linux/pagewalk.h |  1 +
 kernel/events/uprobes.c  |  4 +++-
 mm/huge_memory.c         |  4 +++-
 mm/ksm.c                 |  4 +++-
 mm/migrate.c             |  8 ++++++--
 mm/pagewalk.c            | 10 +++++++++-
 mm/rmap.c                |  4 +++-
 8 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/arch/s390/mm/fault.c b/arch/s390/mm/fault.c
index 46d828926009..968f128a5232 100644
--- a/arch/s390/mm/fault.c
+++ b/arch/s390/mm/fault.c
@@ -412,7 +412,9 @@ __context_unsafe(/* folio_walk_end() not instrumented */)
 	unsigned long addr = get_fault_address(regs);
 	struct mm_struct *mm = current->mm;
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	int rc;
 
diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
index 1c397be0d092..455428eaf1fc 100644
--- a/include/linux/pagewalk.h
+++ b/include/linux/pagewalk.h
@@ -184,6 +184,7 @@ struct folio_walk {
 	};
 	/* private */
 	spinlock_t *ptl;
+	enum page_walk_lock walk_lock;
 };
 
 struct folio *folio_walk_start(struct folio_walk *fw,
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index caf444762a27..3e3927109ee9 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -507,7 +507,9 @@ int uprobe_write(struct arch_uprobe *auprobe, struct vm_area_struct *vma,
 	int ret, ref_ctr_updated = 0;
 	unsigned int gup_flags = FOLL_FORCE;
 	struct mmu_notifier_range range;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	struct page *page;
 
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index bd44d58e8684..62ec81396c9d 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -4816,7 +4816,9 @@ static int split_huge_pages_pid(int pid, unsigned long vaddr_start,
 	 */
 	for (addr = vaddr_start; addr < vaddr_end; addr += PAGE_SIZE) {
 		struct vm_area_struct *vma = vma_lookup(mm, addr);
-		struct folio_walk fw;
+		struct folio_walk fw = {
+			.walk_lock  = PGWALK_RDLOCK,
+		};
 		struct folio *folio;
 		struct address_space *mapping;
 		unsigned int target_order = new_order;
diff --git a/mm/ksm.c b/mm/ksm.c
index d6f736a1aa2f..d4fc0b0965aa 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -819,8 +819,10 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 	unsigned long addr = rmap_item->address;
 	struct vm_area_struct *vma;
 	struct page *page = NULL;
-	struct folio_walk fw;
 	struct folio *folio;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 
 	mmap_read_lock(mm);
 	vma = find_mergeable_vma(mm, addr);
diff --git a/mm/migrate.c b/mm/migrate.c
index 875079464930..89662d5bedf4 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -2301,7 +2301,9 @@ static int add_folio_for_migration(struct mm_struct *mm, const void __user *p,
 		int node, struct list_head *pagelist, bool migrate_all)
 {
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	unsigned long addr;
 	int err = -EFAULT;
@@ -2463,7 +2465,9 @@ static void do_pages_stat_array(struct mm_struct *mm, unsigned long nr_pages,
 	for (i = 0; i < nr_pages; i++) {
 		unsigned long addr = (unsigned long)untagged_addr_remote(mm, *pages);
 		struct vm_area_struct *vma;
-		struct folio_walk fw;
+		struct folio_walk fw = {
+			.walk_lock  = PGWALK_RDLOCK,
+		};
 		struct folio *folio;
 		int err = -EFAULT;
 
diff --git a/mm/pagewalk.c b/mm/pagewalk.c
index e6493bbe6919..7d85289fcb87 100644
--- a/mm/pagewalk.c
+++ b/mm/pagewalk.c
@@ -910,7 +910,15 @@ struct folio *folio_walk_start(struct folio_walk *fw,
 	pgd_t *pgdp;
 	p4d_t *p4dp;
 
-	mmap_assert_locked(vma->vm_mm);
+	/*
+	 * Other locking modes except for mmap or vma read locking are not
+	 * expected.
+	 */
+	if (fw->walk_lock != PGWALK_RDLOCK && fw->walk_lock != PGWALK_VMA_RDLOCK_VERIFY)
+		WARN_ONCE(1, "walk_lock is not expected!\n");
+	process_mm_walk_lock(vma->vm_mm, fw->walk_lock);
+	process_vma_walk_lock(vma, fw->walk_lock);
+
 	vma_pgtable_walk_begin(vma);
 
 	if (WARN_ON_ONCE(addr < vma->vm_start || addr >= vma->vm_end))
diff --git a/mm/rmap.c b/mm/rmap.c
index 1131b76bbbc2..6e3bc2acd1d7 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2862,7 +2862,9 @@ struct page *make_device_exclusive(struct mm_struct *mm, unsigned long addr,
 	struct mmu_notifier_range range;
 	struct folio *folio, *fw_folio;
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct page *page;
 	swp_entry_t entry;
 	pte_t swp_pte;
-- 
2.34.1

Re: [PATCH v2 2/4] mm: make folio_walk_start()'s locking asserts scalable

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 3 potential issue(s) to consider:
- [Medium] mm: inaccurate commit message on write-locking in folio_walk_start()
- [Low] mm: folio_walk_start() duplicates vma_assert_stabilised() via walk_lock
- [Low] mm: missing kernel-doc and private placement for folio_walk::walk_lock
--

commit fbf030fed81b9a0de749bef1b3fa6307e8400b93
Author: Xu Xin <xu.xin@linux.dev>

mm: make folio_walk_start()'s locking asserts scalable

This commit adds a walk_lock member to struct folio_walk and updates
folio_walk_start() and its callers to specify locking assertion modes for
mmap_lock and per-VMA locks.

>     mm: make folio_walk_start()'s locking asserts scalable
>     
>     Add an additional member 'walk_lock' to folio_walk to indicate
>     locking requirements for the walk. Similar to commit 49b0638502da0
>     ("mm: enable page walking API to lock vmas during the walk").
>     But no change is made on any existing locking behavior,
>     all existing folio_walk_start() callers are currently still
>     under mmap_read_lock() protection.
>     
>     This change is prepared for the later patch to enable VMA locking
>     asserts. folio_walk_start now operate under write-locked mmap_lock.
>     With introduction of vma locks at the next patch, the vmas have to
>     be locked as well during such walks to prevent concurrent page
>     faults in these areas.
>     
>     No functional change intended.
>     
>     Signed-off-by: Xu Xin <xu.xin@linux.dev>

[Severity: Medium]
The second paragraph states that "folio_walk_start now operate under
write-locked mmap_lock" and that VMAs have to be locked during such walks.
Is this description accurate?

In mm/pagewalk.c, folio_walk_start() only expects PGWALK_RDLOCK or
PGWALK_VMA_RDLOCK_VERIFY and emits a warning on write-lock modes:

	if (fw->walk_lock != PGWALK_RDLOCK && fw->walk_lock != PGWALK_VMA_RDLOCK_VERIFY)
		WARN_ONCE(1, "walk_lock is not expected!\n");

Could the commit message be updated to reflect that folio_walk_start() only
verifies read-stabilized mmap or VMA locks?

[ ... ]

> diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
> index 1c397be0d092a..455428eaf1fc0 100644
> --- a/include/linux/pagewalk.h
> +++ b/include/linux/pagewalk.h
> @@ -184,6 +184,7 @@ struct folio_walk {
>  	};
>  	/* private */
>  	spinlock_t *ptl;
> +	enum page_walk_lock walk_lock;
>  };

[Severity: Low]
This isn't a bug, but adding walk_lock to struct folio_walk without a
@walk_lock entry in the kernel-doc comment above struct folio_walk triggers a
scripts/kernel-doc warning ("struct member 'walk_lock' not described in
'folio_walk'"), because the /* private */ comment does not use the
/* private: */ syntax recognized by kernel-doc.

Also, since callers of folio_walk_start() are required to initialize
fw.walk_lock before the call, should walk_lock be placed in the public
section of struct folio_walk rather than under /* private */, and should the
kernel-doc for folio_walk_start() in mm/pagewalk.c (which currently states
"@fw: filled with information on success." and "The mmap lock must be held in
read mode.") be updated to match?

[ ... ]

> diff --git a/mm/pagewalk.c b/mm/pagewalk.c
> index e6493bbe6919e..7d85289fcb875 100644
> --- a/mm/pagewalk.c
> +++ b/mm/pagewalk.c
> @@ -910,7 +910,15 @@ struct folio *folio_walk_start(struct folio_walk *fw,
>  	pgd_t *pgdp;
>  	p4d_t *p4dp;
>  
> -	mmap_assert_locked(vma->vm_mm);
> +	/*
> +	 * Other locking modes except for mmap or vma read locking are not
> +	 * expected.
> +	 */
> +	if (fw->walk_lock != PGWALK_RDLOCK && fw->walk_lock != PGWALK_VMA_RDLOCK_VERIFY)
> +		WARN_ONCE(1, "walk_lock is not expected!\n");
> +	process_mm_walk_lock(vma->vm_mm, fw->walk_lock);
> +	process_vma_walk_lock(vma, fw->walk_lock);

[Severity: Low]
If a caller passes an unexpected walk_lock value (such as PGWALK_WRLOCK) or
leaves struct folio_walk uninitialized, folio_walk_start() triggers
WARN_ONCE() here, but then still calls process_mm_walk_lock() and
process_vma_walk_lock(), which for PGWALK_WRLOCK will call
vma_start_write(vma).

More broadly, since folio_walk_start() only permits mmap or VMA read-lock
verification, could folio_walk_start() simply call vma_assert_stabilised(vma)
from include/linux/mmap_lock.h instead of adding walk_lock to struct
folio_walk and requiring every caller to initialize it:

static inline void vma_assert_stabilised(const struct vm_area_struct *vma)
{
	if (IS_ENABLED(CONFIG_LOCKDEP)) {
		if (lockdep_is_held(&vma->vm_mm->mmap_lock))
			return;
	} else {
		if (rwsem_is_locked(&vma->vm_mm->mmap_lock))
			return;
	}

	vma_assert_locked(vma);
}

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261102132416.300084-1-xu.xin@linux.dev?part=2

[PATCH v2 3/4] mm/ksm: make break_ksm() more scalable

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

break_ksm() currently takes a boolean 'lock_vma' to select between two
static mm_walk_ops tables.  This is fragile: the ops struct is full of
function pointers and making it mutable (to switch walk_lock at runtime)
turns it into a juicy target for attackers.

Replace the boolean with a KSM-local enum ksm_walk_lock and build the
mm_walk_ops on the stack, so the table stays const and there is no
shared mutable state.  This also prepares for the later patch that
switches break_ksm() to per-VMA read-locking, where a third locking
mode will be needed.

No functional change intended.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 mm/ksm.c | 42 +++++++++++++++++++++++++++---------------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/mm/ksm.c b/mm/ksm.c
index d4fc0b0965aa..84d363dabdc8 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -661,15 +661,25 @@ static int break_ksm_pmd_entry(pmd_t *pmdp, unsigned long addr, unsigned long en
 	return found;
 }
 
-static const struct mm_walk_ops break_ksm_ops = {
-	.pmd_entry = break_ksm_pmd_entry,
-	.walk_lock = PGWALK_RDLOCK,
+/*
+ * Locking contexts in which break_ksm() may walk the page tables.
+ */
+enum ksm_walk_lock {
+	/* walk under mmap_read_lock() */
+	KSM_WALK_MMAP_RDLOCK,
+	/* walk under mmap_write_lock() */
+	KSM_WALK_MMAP_WRLOCK,
 };
 
-static const struct mm_walk_ops break_ksm_lock_vma_ops = {
-	.pmd_entry = break_ksm_pmd_entry,
-	.walk_lock = PGWALK_WRLOCK,
-};
+static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
+{
+	switch (lock) {
+	case KSM_WALK_MMAP_WRLOCK:
+		return PGWALK_WRLOCK;
+	default:
+		return PGWALK_RDLOCK;
+	}
+}
 
 /*
  * Though it's very tempting to unmerge rmap_items from stable tree rather
@@ -697,17 +707,19 @@ static const struct mm_walk_ops break_ksm_lock_vma_ops = {
  * protection keys here anyway.
  */
 static int break_ksm(struct vm_area_struct *vma, unsigned long addr,
-		unsigned long end, bool lock_vma)
+		unsigned long end, enum ksm_walk_lock walk_lock)
 {
 	vm_fault_t ret = 0;
-	const struct mm_walk_ops *ops = lock_vma ?
-				&break_ksm_lock_vma_ops : &break_ksm_ops;
+	struct mm_walk_ops ops = {
+		.pmd_entry = break_ksm_pmd_entry,
+		.walk_lock = ksm_walk_lock_to_pgwalk(walk_lock),
+	};
 
 	do {
 		int ksm_page;
 
 		cond_resched();
-		ksm_page = walk_page_range_vma(vma, addr, end, ops, &addr);
+		ksm_page = walk_page_range_vma(vma, addr, end, &ops, &addr);
 		if (ksm_page <= 0)
 			return ksm_page;
 		ret = handle_mm_fault(vma, addr,
@@ -809,7 +821,7 @@ static void break_cow(struct ksm_rmap_item *rmap_item)
 	mmap_read_lock(mm);
 	vma = find_mergeable_vma(mm, addr);
 	if (vma)
-		break_ksm(vma, addr, addr + PAGE_SIZE, false);
+		break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_MMAP_RDLOCK);
 	mmap_read_unlock(mm);
 }
 
@@ -1247,7 +1259,7 @@ static int unmerge_and_remove_all_rmap_items(void)
 		for_each_vma(vmi, vma) {
 			if (!(vma->vm_flags & VM_MERGEABLE) || !vma_has_anon_rmap(vma))
 				continue;
-			err = break_ksm(vma, vma->vm_start, vma->vm_end, false);
+			err = break_ksm(vma, vma->vm_start, vma->vm_end, KSM_WALK_MMAP_RDLOCK);
 			if (err)
 				goto error;
 		}
@@ -2896,7 +2908,7 @@ static int __ksm_del_vma(struct vm_area_struct *vma)
 		return 0;
 
 	if (vma_has_anon_rmap(vma)) {
-		err = break_ksm(vma, vma->vm_start, vma->vm_end, true);
+		err = break_ksm(vma, vma->vm_start, vma->vm_end, KSM_WALK_MMAP_WRLOCK);
 		if (err)
 			return err;
 	}
@@ -3048,7 +3060,7 @@ int ksm_madvise(struct vm_area_struct *vma, unsigned long start,
 			return 0;		/* just ignore the advice */
 
 		if (vma_has_anon_rmap(vma)) {
-			err = break_ksm(vma, start, end, true);
+			err = break_ksm(vma, start, end, KSM_WALK_MMAP_WRLOCK);
 			if (err)
 				return err;
 		}
-- 
2.34.1

Re: [PATCH v2 3/4] mm/ksm: make break_ksm() more scalable

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [Low] mm/ksm: redundant ksm_walk_lock enum and helper in break_ksm()
- [Low] mm/ksm: non-const stack mm_walk_ops in break_ksm() contradicts commit log
--

commit 05d53e6a6817e77860f887304a4f4cd4936f2c61
Author: Xu Xin <xu.xin@linux.dev>

mm/ksm: make break_ksm() more scalable

This commit replaces the boolean lock_vma parameter of break_ksm() with a
KSM-local enum ksm_walk_lock and constructs struct mm_walk_ops on the stack
instead of selecting between two static const tables.

> diff --git a/mm/ksm.c b/mm/ksm.c
> index d4fc0b0965aa1..84d363dabdc8a 100644
> --- a/mm/ksm.c
> +++ b/mm/ksm.c
> @@ -661,15 +661,25 @@ static int break_ksm_pmd_entry(pmd_t *pmdp, unsigned long addr, unsigned long en
>  	return found;
>  }
>  
> -static const struct mm_walk_ops break_ksm_ops = {
> -	.pmd_entry = break_ksm_pmd_entry,
> -	.walk_lock = PGWALK_RDLOCK,
> +/*
> + * Locking contexts in which break_ksm() may walk the page tables.
> + */
> +enum ksm_walk_lock {
> +	/* walk under mmap_read_lock() */
> +	KSM_WALK_MMAP_RDLOCK,
> +	/* walk under mmap_write_lock() */
> +	KSM_WALK_MMAP_WRLOCK,
>  };
>  
> -static const struct mm_walk_ops break_ksm_lock_vma_ops = {
> -	.pmd_entry = break_ksm_pmd_entry,
> -	.walk_lock = PGWALK_WRLOCK,
> -};
> +static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
> +{
> +	switch (lock) {
> +	case KSM_WALK_MMAP_WRLOCK:
> +		return PGWALK_WRLOCK;
> +	default:
> +		return PGWALK_RDLOCK;
> +	}
> +}

[Severity: Low]
This isn't a bug, but is enum ksm_walk_lock and ksm_walk_lock_to_pgwalk()
needed here?

include/linux/pagewalk.h is already included in mm/ksm.c and defines enum
page_walk_lock (PGWALK_RDLOCK, PGWALK_WRLOCK, PGWALK_VMA_RDLOCK_VERIFY),
which is already used directly elsewhere in mm/ksm.c. Since break_ksm() only
uses walk_lock to populate ops.walk_lock, could break_ksm() take enum
page_walk_lock directly instead of introducing a 1:1 wrapper enum and
conversion helper?

Also, the comment on KSM_WALK_MMAP_WRLOCK says "walk under mmap_write_lock()",
whereas mapping it to PGWALK_WRLOCK causes process_vma_walk_lock() to acquire
the VMA write lock via vma_start_write(vma) rather than only asserting that
mmap_write_lock() is held.

>  /*
>   * Though it's very tempting to unmerge rmap_items from stable tree rather
> @@ -697,17 +707,19 @@ static const struct mm_walk_ops break_ksm_lock_vma_ops = {
>   * protection keys here anyway.
>   */
>  static int break_ksm(struct vm_area_struct *vma, unsigned long addr,
> -		unsigned long end, bool lock_vma)
> +		unsigned long end, enum ksm_walk_lock walk_lock)
>  {
>  	vm_fault_t ret = 0;
> -	const struct mm_walk_ops *ops = lock_vma ?
> -				&break_ksm_lock_vma_ops : &break_ksm_ops;
> +	struct mm_walk_ops ops = {
> +		.pmd_entry = break_ksm_pmd_entry,
> +		.walk_lock = ksm_walk_lock_to_pgwalk(walk_lock),
> +	};

[Severity: Low]
The commit message states:

  break_ksm() currently takes a boolean 'lock_vma' to select between two
  static mm_walk_ops tables.  This is fragile: the ops struct is full of
  function pointers and making it mutable (to switch walk_lock at runtime)
  turns it into a juicy target for attackers.

  Replace the boolean with a KSM-local enum ksm_walk_lock and build the
  mm_walk_ops on the stack, so the table stays const and there is no
  shared mutable state.

Previously, break_ksm_ops and break_ksm_lock_vma_ops were declared as static
const struct mm_walk_ops in .rodata, whereas this change declares struct
mm_walk_ops ops on the writable stack without const and reinitializes it on
every break_ksm() call.

Does moving the function pointer table from static const storage in .rodata
onto the stack without const match the security rationale in the commit
message, or should break_ksm() either keep selecting between static const
struct mm_walk_ops tables (as done in fs/proc/task_mmu.c) or at least mark
the stack struct const?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261102132416.300084-1-xu.xin@linux.dev?part=3

[PATCH v2 4/4] ksm: add find_mergeable_vma_locked() to use per-VMA locking

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

Longlong Xia reported that get_mergeable_page() takes mmap_read_lock()
on every merge, which becomes a contention point under load [1]. David
Hildenbrand suggested switching to per-VMA locking to avoid the
process-wide mmap lock [2].

Add find_mergeable_vma_locked(), which is similar to find_tcp_vma(),
using the universal per-VMA locking helper vma_start_read_unlocked() to
look up and read-lock a VM_MERGEABLE VMA without taking
mmap_read_lock().  All KSM call sites that previously used
find_mergeable_vma() under mmap_read_lock() are converted to the new
helper, and the locking in get_mergeable_page() is switched to
PGWALK_VMA_RDLOCK_VERIFY so that folio_walk_start() can verify the
per-VMA lock is held.

Switching to per-VMA locking also changes how ksmd is synchronized with
a process that is exiting.  __ksm_exit() previously relied on a plain
mmap_write_lock() to serialize against ksmd, but that no longer blocks
a ksmd walk that holds only a per-VMA read lock.  So __ksm_exit() now
also write-locks every VMA (vma_start_write()) while holding
mmap_write_lock(): this blocks until any in-flight ksmd walk drops its
per-VMA read lock, and prevents new ones from starting, since
vma_start_read_unlocked() backs out when a VMA is write-locked.  The
per-VMA write locks are released by mmap_write_unlock() via
vma_end_write_all().  This is safe because __mmput() runs ksm_exit()
after mm_users has gone to zero and before exit_mmap() frees the page
tables, so any new ksmd walk backs out via ksm_test_exit().

Performance
===========
As a test, I construct a testcase which follows the approach: Create
one victim and several churner threads sharing one mm_struct [3]; The
victim registers a 32 MiB anonymous VM_MERGEABLE region containing
8192 pages totally:churners hammer mmap_lock via mmap/munmap
repeatedly; ksmd merges the victim's pages; Compare latency baseline
VS this per-VMA patch.

                 Before patched         After Patched
0 churner:     1.627 seconds          1.426 seconds
4 churners:    72.45 seconds          36.61 seconds

In conclusion, when no mmap_lock contention (0 churner), there is little
difference between the baseline kernel and the per-VMA optimized kernel;
But under interference from 4 churner threads, the merge time of the
per-VMA KSM-optimized kernel is significantly reduced by 50%.

[1] https://lore.kernel.org/all/20260909023507.1408925-1-xialonglong2025@163.com/
[2] https://lore.kernel.org/all/91cf8129-d705-4ebd-b420-129bab7ed261@kernel.org/
[3] https://lore.kernel.org/all/20260911162242418ODBbr9wspxxIXDS8uvpeZ@zte.com.cn/

Reported-by: Longlong Xia <xialonglong@kylinos.cn>
Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 mm/ksm.c | 82 +++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 60 insertions(+), 22 deletions(-)

diff --git a/mm/ksm.c b/mm/ksm.c
index 84d363dabdc8..64f9e803b99e 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -669,6 +669,8 @@ enum ksm_walk_lock {
 	KSM_WALK_MMAP_RDLOCK,
 	/* walk under mmap_write_lock() */
 	KSM_WALK_MMAP_WRLOCK,
+	/* walk with the vma already read-locked (per-VMA locking) */
+	KSM_WALK_VMA_RDLOCK,
 };
 
 static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
@@ -676,6 +678,8 @@ static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
 	switch (lock) {
 	case KSM_WALK_MMAP_WRLOCK:
 		return PGWALK_WRLOCK;
+	case KSM_WALK_VMA_RDLOCK:
+		return PGWALK_VMA_RDLOCK_VERIFY;
 	default:
 		return PGWALK_RDLOCK;
 	}
@@ -784,15 +788,36 @@ static bool vma_ksm_compatible(struct vm_area_struct *vma)
 	return ksm_compatible(vma->vm_file, vma->flags);
 }
 
-static struct vm_area_struct *find_mergeable_vma(struct mm_struct *mm,
-		unsigned long addr)
+/**
+ * find_mergeable_vma_locked() - Find the VMA covering 'address' which is
+ * VM_MERGEABLE and read-lock it by per-VMA locks. Please use vma_end_read()
+ * to unlock vma after finishing reading the VMA (non-NULL).
+ *
+ * Return: If a VMA exists which spans @address, return that VMA, read-locked.
+ * If no VMA is mapped there or, very unlikely, a reference count overflow
+ * occurred, return NULL, and no read-locked.
+ *
+ * IMPORTANT: If a VMA exists but is not VM_MERGEABLE or has no anon_vma,
+ * this function releases the per-VMA read lock before returning NULL.
+ * Callers must NOT call vma_end_read() on a NULL return value.
+ */
+static struct vm_area_struct *find_mergeable_vma_locked(struct mm_struct *mm,
+		unsigned long address)
 {
 	struct vm_area_struct *vma;
+
 	if (ksm_test_exit(mm))
 		return NULL;
-	vma = vma_lookup(mm, addr);
-	if (!vma || !(vma->vm_flags & VM_MERGEABLE) || !vma_has_anon_rmap(vma))
+
+	vma = vma_start_read_unlocked(mm, address);
+	if (!vma)
 		return NULL;
+
+	if (!(vma->vm_flags & VM_MERGEABLE) || !vma->anon_vma) {
+		vma_end_read(vma);
+		return NULL;
+	}
+
 	return vma;
 }
 
@@ -818,11 +843,12 @@ static void break_cow(struct ksm_rmap_item *rmap_item)
 	 */
 	rmap_item->linear_page_index = 0;
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, addr);
-	if (vma)
-		break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_MMAP_RDLOCK);
-	mmap_read_unlock(mm);
+	vma = find_mergeable_vma_locked(mm, addr);
+	if (!vma)
+		return;
+
+	break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_VMA_RDLOCK);
+	vma_end_read(vma);
 }
 
 static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
@@ -833,13 +859,12 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 	struct page *page = NULL;
 	struct folio *folio;
 	struct folio_walk fw = {
-		.walk_lock  = PGWALK_RDLOCK,
+		.walk_lock  = PGWALK_VMA_RDLOCK_VERIFY,
 	};
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, addr);
+	vma = find_mergeable_vma_locked(mm, addr);
 	if (!vma)
-		goto out;
+		return NULL;
 
 	folio = folio_walk_start(&fw, vma, addr, 0);
 	if (folio) {
@@ -850,12 +875,12 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 		}
 		folio_walk_end(&fw, vma);
 	}
-out:
+
 	if (page) {
 		flush_anon_page(vma, page, addr);
 		flush_dcache_page(page);
 	}
-	mmap_read_unlock(mm);
+	vma_end_read(vma);
 	return page;
 }
 
@@ -1596,14 +1621,14 @@ static int try_to_merge_with_zero_page(struct ksm_rmap_item *rmap_item,
 	if (ksm_use_zero_pages && (rmap_item->oldchecksum == zero_checksum)) {
 		struct vm_area_struct *vma;
 
-		mmap_read_lock(mm);
-		vma = find_mergeable_vma(mm, rmap_item->address);
+		vma = find_mergeable_vma_locked(mm, rmap_item->address);
 		if (vma) {
 			err = try_to_merge_one_page(vma, page,
 					ZERO_PAGE(rmap_item->address));
 			trace_ksm_merge_one_page(
 				page_to_pfn(ZERO_PAGE(rmap_item->address)),
 				rmap_item, mm, err);
+			vma_end_read(vma);
 		} else {
 			/*
 			 * If the vma is out of date, we do not need to
@@ -1611,7 +1636,6 @@ static int try_to_merge_with_zero_page(struct ksm_rmap_item *rmap_item,
 			 */
 			err = 0;
 		}
-		mmap_read_unlock(mm);
 	}
 
 	return err;
@@ -1630,10 +1654,9 @@ static int try_to_merge_with_ksm_page(struct ksm_rmap_item *rmap_item,
 	struct vm_area_struct *vma;
 	int err = -EFAULT;
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, rmap_item->address);
+	vma = find_mergeable_vma_locked(mm, rmap_item->address);
 	if (!vma)
-		goto out;
+		goto out_trace;
 
 	err = try_to_merge_one_page(vma, page, kpage);
 	if (err)
@@ -1653,7 +1676,8 @@ static int try_to_merge_with_ksm_page(struct ksm_rmap_item *rmap_item,
 	rmap_item->linear_page_index = linear_anon_page_index(vma, rmap_item->address);
 	get_anon_vma(vma->anon_vma);
 out:
-	mmap_read_unlock(mm);
+	vma_end_read(vma);
+out_trace:
 	trace_ksm_merge_with_ksm_page(kpage, page_to_pfn(kpage ? kpage : page),
 				rmap_item, mm, err);
 	return err;
@@ -3119,6 +3143,7 @@ void __ksm_exit(struct mm_struct *mm)
 {
 	struct ksm_mm_slot *mm_slot = NULL;
 	struct mm_slot *slot;
+	struct vm_area_struct *vma;
 	int easy_to_free = 0;
 
 	/*
@@ -3153,7 +3178,20 @@ void __ksm_exit(struct mm_struct *mm)
 		mm_flags_clear(MMF_VM_MERGEABLE, mm);
 		mmdrop(mm);
 	} else if (mm_slot) {
+		/*
+		 * ksmd may be walking this mm's page tables under a per-VMA
+		 * read lock (see find_mergeable_vma_locked()).  A plain
+		 * mmap_write_lock() does not exclude such readers, so also
+		 * write-lock every VMA: this blocks until any in-flight ksmd
+		 * walk drops its per-VMA read lock, and prevents new ones
+		 * from starting (vma_start_read_unlocked() backs out when a
+		 * VMA is write-locked).  The per-VMA write locks are released
+		 * by mmap_write_unlock() via vma_end_write_all().
+		 */
 		mmap_write_lock(mm);
+		VMA_ITERATOR(vmi, mm, 0);
+		for_each_vma(vmi, vma)
+			vma_start_write(vma);
 		mmap_write_unlock(mm);
 	}
 
-- 
2.34.1

Re: [PATCH v2 4/4] ksm: add find_mergeable_vma_locked() to use per-VMA locking

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 4 potential issue(s) to consider:
- [Critical] mm: ksm: race with exit_mmap() in find_mergeable_vma_locked() and __ksm_exit()
- [High] mm: ksm: missing FAULT_FLAG_VMA_LOCK in break_ksm() under per-VMA read lock
- [Low] mm: ksm: unmarked vma->anon_vma read in find_mergeable_vma_locked()
- [Low] mm: ksm: missing kernel-doc params for find_mergeable_vma_locked()
--

commit 984120154ab3ef7af44dc3dbb32b8522a336aefe
Author: Xu Xin <xu.xin@linux.dev>

ksm: add find_mergeable_vma_locked() to use per-VMA locking

This commit introduces find_mergeable_vma_locked() to look up and read-lock
VM_MERGEABLE VMAs using per-VMA locks instead of mmap_read_lock() across KSM
page-scanning and merging paths. It also updates __ksm_exit() to write-lock
each VMA under mmap_write_lock() to synchronize with in-flight per-VMA readers
when a process exits.

Link: https://lore.kernel.org/all/20260909023507.1408925-1-xialonglong2025@163.com/
Link: https://lore.kernel.org/all/91cf8129-d705-4ebd-b420-129bab7ed261@kernel.org/
Link: https://lore.kernel.org/all/20260911162242418ODBbr9wspxxIXDS8uvpeZ@zte.com.cn/

> diff --git a/mm/ksm.c b/mm/ksm.c
> index 84d363dabdc8a..64f9e803b99e2 100644
> --- a/mm/ksm.c
> +++ b/mm/ksm.c
> @@ -669,6 +669,8 @@ enum ksm_walk_lock {
>  	KSM_WALK_MMAP_RDLOCK,
>  	/* walk under mmap_write_lock() */
>  	KSM_WALK_MMAP_WRLOCK,
> +	/* walk with the vma already read-locked (per-VMA locking) */
> +	KSM_WALK_VMA_RDLOCK,
>  };
>  
>  static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
> @@ -676,6 +678,8 @@ static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
>  	switch (lock) {
>  	case KSM_WALK_MMAP_WRLOCK:
>  		return PGWALK_WRLOCK;
> +	case KSM_WALK_VMA_RDLOCK:
> +		return PGWALK_VMA_RDLOCK_VERIFY;
>  	default:
>  		return PGWALK_RDLOCK;
>  	}
> @@ -784,15 +788,36 @@ static bool vma_ksm_compatible(struct vm_area_struct *vma)
>  	return ksm_compatible(vma->vm_file, vma->flags);
>  }
>  
> -static struct vm_area_struct *find_mergeable_vma(struct mm_struct *mm,
> -		unsigned long addr)
> +/**
> + * find_mergeable_vma_locked() - Find the VMA covering 'address' which is
> + * VM_MERGEABLE and read-lock it by per-VMA locks. Please use vma_end_read()
> + * to unlock vma after finishing reading the VMA (non-NULL).
> + *
> + * Return: If a VMA exists which spans @address, return that VMA, read-locked.
> + * If no VMA is mapped there or, very unlikely, a reference count overflow
> + * occurred, return NULL, and no read-locked.
> + *
> + * IMPORTANT: If a VMA exists but is not VM_MERGEABLE or has no anon_vma,
> + * this function releases the per-VMA read lock before returning NULL.
> + * Callers must NOT call vma_end_read() on a NULL return value.
> + */
> +static struct vm_area_struct *find_mergeable_vma_locked(struct mm_struct *mm,
> +		unsigned long address)

[Severity: Low]
This isn't a bug, but since this comment above find_mergeable_vma_locked() in
mm/ksm.c starts with the /** kernel-doc marker, won't scripts/kernel-doc (when
building with W=1) warn about the missing @mm and @address parameter
descriptions?

>  {
>  	struct vm_area_struct *vma;
> +
>  	if (ksm_test_exit(mm))
>  		return NULL;
> -	vma = vma_lookup(mm, addr);
> -	if (!vma || !(vma->vm_flags & VM_MERGEABLE) || !vma_has_anon_rmap(vma))
> +
> +	vma = vma_start_read_unlocked(mm, address);
> +	if (!vma)
>  		return NULL;
> +
> +	if (!(vma->vm_flags & VM_MERGEABLE) || !vma->anon_vma) {

[Severity: Low]
In find_mergeable_vma_locked() in mm/ksm.c, replacing !vma_has_anon_rmap(vma)
with a plain !vma->anon_vma read under the per-VMA read lock drops the
data_race() annotation and vma_assert_stabilised() check from
vma_has_anon_rmap() in mm/vma.h:

mm/vma.h:vma_has_anon_rmap() {
	if (vma_is_attached(vma))
		vma_assert_stabilised(vma);
	/* KCSAN gets confused about the optimistic check. Silence it. */
	return data_race(vma->anon_vma);
}

When another thread concurrently faults in the first anonymous page on the
same VMA and initializes vma->anon_vma in __anon_vma_prepare() under
mm->page_table_lock, won't this trigger a KCSAN data-race warning? Should
find_mergeable_vma_locked() continue to use !vma_has_anon_rmap(vma) here?

> +		vma_end_read(vma);
> +		return NULL;
> +	}
> +
>  	return vma;
>  }
>  
> @@ -818,11 +843,12 @@ static void break_cow(struct ksm_rmap_item *rmap_item)
>  	 */
>  	rmap_item->linear_page_index = 0;
>  
> -	mmap_read_lock(mm);
> -	vma = find_mergeable_vma(mm, addr);
> -	if (vma)
> -		break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_MMAP_RDLOCK);
> -	mmap_read_unlock(mm);
> +	vma = find_mergeable_vma_locked(mm, addr);
> +	if (!vma)
> +		return;
> +
> +	break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_VMA_RDLOCK);
> +	vma_end_read(vma);
>  }

[Severity: High]
When break_cow() in mm/ksm.c calls break_ksm() with KSM_WALK_VMA_RDLOCK, only
the per-VMA read lock is held, without mm->mmap_lock. However, break_ksm()
still calls handle_mm_fault() without FAULT_FLAG_VMA_LOCK:

mm/ksm.c:break_ksm() {
	...
	do {
		int ksm_page;

		cond_resched();
		ksm_page = walk_page_range_vma(vma, addr, end, &ops, &addr);
		if (ksm_page <= 0)
			return ksm_page;
		ret = handle_mm_fault(vma, addr,
				      FAULT_FLAG_UNSHARE | FAULT_FLAG_REMOTE,
				      NULL);
	} while (!(ret & (VM_FAULT_SIGBUS | VM_FAULT_SIGSEGV | VM_FAULT_OOM)));
}

Because break_ksm_pmd_entry() drops the PTE lock before returning, the PTE at
addr can be modified (for example, zapped by concurrent MADV_DONTNEED or
truncation, or turned into a device-private swap entry) before
handle_pte_fault() inspects it.

Without FAULT_FLAG_VMA_LOCK set in vmf->flags, handle_mm_fault() and its
callees assume mm->mmap_lock is held:

1. In mm/memory.c, vmf_can_call_fault() checks !(vmf->flags &
   FAULT_FLAG_VMA_LOCK) and allows do_read_fault() to invoke
   vma->vm_ops->fault() on a ksm_compatible() MAP_PRIVATE file-backed VMA
   without holding mm->mmap_lock.

2. In include/linux/mm.h, assert_fault_locked() checks
   mmap_assert_locked(vmf->vma->vm_mm) when FAULT_FLAG_VMA_LOCK is not set,
   and release_fault_lock() calls mmap_read_unlock() on the unheld
   mm->mmap_lock (e.g. if handle_userfault() or migration entry waiting is
   reached).

3. If FAULT_FLAG_VMA_LOCK were passed, any path returning VM_FAULT_RETRY
   would release the per-VMA read lock via vma_end_read(vma), which neither
   the loop in break_ksm() nor the unconditional vma_end_read(vma) in
   break_cow() accounts for.

Can break_ksm() safely call handle_mm_fault() under only a per-VMA read lock?

[ ... ]

> @@ -3153,7 +3178,20 @@ void __ksm_exit(struct mm_struct *mm)
>  		mm_flags_clear(MMF_VM_MERGEABLE, mm);
>  		mmdrop(mm);
>  	} else if (mm_slot) {
> +		/*
> +		 * ksmd may be walking this mm's page tables under a per-VMA
> +		 * read lock (see find_mergeable_vma_locked()).  A plain
> +		 * mmap_write_lock() does not exclude such readers, so also
> +		 * write-lock every VMA: this blocks until any in-flight ksmd
> +		 * walk drops its per-VMA read lock, and prevents new ones
> +		 * from starting (vma_start_read_unlocked() backs out when a
> +		 * VMA is write-locked).  The per-VMA write locks are released
> +		 * by mmap_write_unlock() via vma_end_write_all().
> +		 */
>  		mmap_write_lock(mm);
> +		VMA_ITERATOR(vmi, mm, 0);
> +		for_each_vma(vmi, vma)
> +			vma_start_write(vma);
>  		mmap_write_unlock(mm);
>  	}

[Severity: Critical]
Does this synchronization in __ksm_exit() in mm/ksm.c actually prevent ksmd
from acquiring a per-VMA read lock in find_mergeable_vma_locked() and racing
with exit_mmap()?

Unlike normal callers of vma_start_read_unlocked(), ksmd only holds mm_count
(via mmgrab()), not mm_users. Previously, callers took mmap_read_lock(mm)
first and then checked ksm_test_exit(mm) inside find_mergeable_vma() under
mmap_read_lock(mm), which serialized against mmap_write_lock(mm) in
__ksm_exit().

With this patch, find_mergeable_vma_locked() in mm/ksm.c checks
ksm_test_exit(mm) locklessly before calling vma_start_read_unlocked(mm,
address), and never re-checks ksm_test_exit(mm) after acquiring the per-VMA
read lock:

mm/ksm.c:find_mergeable_vma_locked() {
	if (ksm_test_exit(mm))
		return NULL;

	vma = vma_start_read_unlocked(mm, address);
	if (!vma)
		return NULL;
	...
}

Contrary to the comment in __ksm_exit(), vma_start_read_unlocked() in
mm/mmap_lock.c does not back out when a VMA is write-locked; when
lock_vma_under_rcu() returns NULL, it enters the slow path and blocks on
mmap_read_lock(mm):

mm/mmap_lock.c:vma_start_read_unlocked() {
	vma = lock_vma_under_rcu(mm, address);
	if (vma)
		return vma;

	/* Slow path: preclude VMA writers by temporarily getting mmap read lock. */
	mmap_read_lock(mm);
	vma = vma_lookup(mm, address);
	if (vma && !vma_start_read_locked(vma))
		vma = NULL;
	mmap_read_unlock(mm);

	return vma;
}

And in __ksm_exit(), mmap_write_unlock(mm) calls vma_end_write_all(mm), which
immediately releases all of the per-VMA write locks just acquired.

If mm->mm_users drops to 0 right after find_mergeable_vma_locked() passes
ksm_test_exit(mm) (in get_mergeable_page(), try_to_merge_with_zero_page(),
try_to_merge_with_ksm_page(), or break_cow()):

1. If ksmd calls vma_start_read_unlocked() while __ksm_exit() holds
   mmap_write_lock(mm), lock_vma_under_rcu() fails and ksmd waits on
   mmap_read_lock(mm) in the slow path. As soon as __ksm_exit() calls
   mmap_write_unlock(mm) (which also calls vma_end_write_all(mm)), ksmd
   acquires mmap_read_lock(mm), looks up the VMA, and acquires the per-VMA
   read lock via vma_start_read_locked(vma). Or if ksmd enters
   vma_start_read_unlocked() right after __ksm_exit() calls
   mmap_write_unlock(mm), lock_vma_under_rcu() succeeds directly because the
   per-VMA write locks were already released.

   Because find_mergeable_vma_locked() does not re-check ksm_test_exit(mm),
   ksmd proceeds to walk and modify page tables (and invoke MMU notifiers)
   under only the per-VMA read lock after __ksm_exit() and
   mmu_notifier_release() have finished, racing with exit_mmap() running
   unmap_vmas() and free_pgtables() (which frees page tables before
   tear_down_vmas() detaches the VMAs).

2. Because mm_users is not pinned across vma_start_read_unlocked(),
   exit_mmap() in mm/mmap.c can run concurrently with
   vma_start_read_unlocked():

mm/mmap.c:exit_mmap() {
	...
	mmap_write_lock(mm);
	unmap.mm_wr_locked = true;
	mt_clear_in_rcu(&mm->mm_mt);
	unmap_pgtable_init(&unmap, &vmi);
	free_pgtables(&tlb, &unmap);
	tlb_finish_mmu(&tlb);
	nr_accounted = tear_down_vmas(mm, &vmi, vma, ULONG_MAX);
destroy:
	__mt_destroy(&mm->mm_mt);
	mmap_write_unlock(mm);
	...
}

   In exit_mmap(), mt_clear_in_rcu(&mm->mm_mt) clears MT_FLAGS_USE_RCU and
   __mt_destroy(&mm->mm_mt) frees maple tree nodes synchronously without RCU
   callbacks, while lock_vma_under_rcu() -> mas_walk() in ksmd may still be
   walking mm->mm_mt under rcu_read_lock(), leading to a use-after-free on
   maple tree nodes (or a maple tree lockdep warning in the slow path's
   vma_lookup() after __mt_destroy() clears MT_FLAGS_LOCK_EXTERN).

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261102132416.300084-1-xu.xin@linux.dev?part=4