[RFC PATCH 47/57] mm/madvise: collapse under a per-VMA read lock

[Kiryl Shutsemau <kirill@shutemov.name>]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

MADV_COLLAPSE arrives under mmap_lock, gives it up, and takes it again per
PMD to scan and hand a table to the collapse.  Every writer to the address
space waits behind each of those, and the range can be as large as the
caller asked for.

Take a read lock on the VMA instead: MADVISE_VMA_READ_LOCK in the lock
mode, lock_vma_under_rcu() per PMD, and the walk told not to release what
the behaviour already let go.

A scan that finds nothing keeps the lock, so a range that is already
collapsed walks it without relocking.  A collapse gives the lock up and
looks the VMA up again afterwards: it can shrink while nothing is held,
which the scan reports as a refused range like any other.

Remote madvise is the exception.  process_madvise() has to untag the range
with untagged_addr_remote() before any VMA is looked at.  That reads mm
state mmap_lock protects, so remote MADV_COLLAPSE keeps the mmap_read it
has.

A VMA that cannot be locked is reported as SCAN_VMA_LOCK, which reaches
the caller as -EAGAIN.  lock_vma_under_rcu() also fails on a VMA being
written to, and reporting that as a range which shrank would tell the
caller a collapse succeeded where none was attempted.  With that, nothing
produces SCAN_VMA_NULL any more, and the test for it goes.

Both callers hold a read lock on the VMA now.  That is what the changes
outside madvise.c are for: the engine's interface documented mmap_lock as
its precondition, and only here does that stop being true of every caller.
The scan asserts the lock again too, which was not possible while the
callers disagreed.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 24 +++++++------
 mm/collapse.h |  8 ++---
 mm/madvise.c  | 98 +++++++++++++++++++++++++++++++++++++++------------
 3 files changed, 93 insertions(+), 37 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index b3343595bdf2..1e3b2d202ffe 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -3685,8 +3685,8 @@ static enum scan_result collapse_scan_file_pmd(struct vm_area_struct *vma,
 
 	/*
 	 * A PMD that is huge already has nothing left to collapse, and skipping
-	 * it here is what keeps mmap_lock out of a collapse that would find
-	 * nothing.  Everything else is worth the page cache scan, pmd_none()
+	 * it here is what keeps a collapse that would find nothing from being
+	 * run at all.  Everything else is worth the page cache scan, pmd_none()
 	 * included: a file range can be collapsed out of the cache without being
 	 * mapped first, which is why this is not the test the anonymous side
 	 * makes.
@@ -3703,8 +3703,8 @@ static enum scan_result collapse_scan_file_pmd(struct vm_area_struct *vma,
 
 /*
  * Build a PMD over what the page cache holds, and map it over the range if a huge
- * folio is already there but mapped by PTEs.  Runs with no mmap_lock, which the
- * caller gave up, and takes it again only for that last step.
+ * folio is already there but mapped by PTEs.  Runs with no lock on the VMA,
+ * which the caller gave up, and takes mmap_lock only for that last step.
  */
 static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 		unsigned long addr, struct collapse_control *cc)
@@ -3746,9 +3746,9 @@ static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 
 /*
  * Scan one table's worth of @vma and decide whether there is anything to collapse
- * in it.  The caller holds a read lock and still holds it when this returns:
- * what is looked at is either the VMA or a page table that the lock keeps in
- * place.
+ * in it.  The caller holds a read lock on @vma and still holds it when this
+ * returns: what is looked at is either the VMA or a page table that the lock
+ * keeps in place.
  *
  * Returns whether collapse_run_pmd() has anything to do, and a scan that found
  * something has to be run: the file side takes a reference on the file while it
@@ -3761,6 +3761,8 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 {
 	struct mm_struct *mm = vma->vm_mm;
 
+	vma_assert_locked(vma);
+
 	/*
 	 * What the scan answers with, so cleared before it runs.
 	 * collapse_anon_scan_init() clears the orders too, but only once the
@@ -3789,10 +3791,10 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 }
 
 /*
- * Collapse what the scan selected.  Called with no mmap_lock: the caller gives it
- * up first, because a collapse takes it again for each round and revalidates
- * under it, and holding it across the whole collapse would keep a writer to the
- * address space waiting for it.
+ * Collapse what the scan selected.  Called with no lock on the VMA the scan
+ * looked at: the caller gives that up first, because a collapse takes its own
+ * for each round and revalidates under it, and holding one across the whole
+ * collapse would keep a writer to the VMA waiting for it.
  */
 enum scan_result collapse_run_pmd(struct mm_struct *mm, unsigned long addr,
 		unsigned long end, struct collapse_control *cc)
diff --git a/mm/collapse.h b/mm/collapse.h
index e5a0ffab049c..aeaca305e71a 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -207,8 +207,8 @@ static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
  *	collapse_run_pmd(mm, addr, end, cc);	when the scan found work
  *	collapse_control_release(cc);
  *
- * The caller holds mmap_lock for reading and passes a range within one PTE table
- * of @vma.  A range the VMA does not cover is refused, which is also how a caller
+ * The caller holds a read lock on @vma and passes a range within one PTE table
+ * of it.  A range the VMA does not cover is refused, which is also how a caller
  * learns that its own range shrank.
  *
  * A scan returns with that lock still held: it only reads, and almost every table
@@ -218,8 +218,8 @@ static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
  * A collapse is called without it: the caller gives the lock up first, and with it
  * @vma and anything derived under it, so a caller carrying on has to look up
  * again.  What the collapse does -- allocate, quiesce, copy, flush -- is slow
- * enough that a writer would wait behind it, so it takes the lock again per round
- * instead, and revalidates rather than trusting what the scan saw.
+ * enough that a writer to the VMA would wait behind it, so it takes its own lock
+ * per round instead, and revalidates rather than trusting what the scan saw.
  *
  * A scan that found something has to be run: the file side takes a reference on
  * the file while it still has the VMA to take it from, and the run is what gives
diff --git a/mm/madvise.c b/mm/madvise.c
index bd9123ee3cb1..5f6d815d70ad 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -276,6 +276,18 @@ static void mark_mmap_lock_dropped(struct madvise_behavior *madv_behavior)
 	madv_behavior->lock_dropped = true;
 }
 
+/*
+ * The VMA-lock counterpart, for a behaviour that releases the VMA it was handed
+ * and locks what it needs for itself.  The walk has nothing left to release,
+ * and unlike the mmap_lock case it has nothing to carry on with either: the VMA
+ * fast path applies to one VMA and returns.
+ */
+static void mark_vma_lock_dropped(struct madvise_behavior *madv_behavior)
+{
+	VM_WARN_ON_ONCE(madv_behavior->lock_mode != MADVISE_VMA_READ_LOCK);
+	madv_behavior->lock_dropped = true;
+}
+
 /*
  * Schedule all required I/O operations.  Do not wait for completion.
  */
@@ -948,6 +960,8 @@ static int madvise_collapse_errno(enum scan_result r)
 static int madvise_collapse(struct madvise_behavior *madv_behavior)
 {
 	struct madvise_behavior_range *range = &madv_behavior->range;
+	const bool vma_locked =
+		madv_behavior->lock_mode == MADVISE_VMA_READ_LOCK;
 	struct vm_area_struct *vma = madv_behavior->vma;
 	struct mm_struct *mm = madv_behavior->mm;
 	unsigned long hstart, hend, addr;
@@ -981,13 +995,22 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 	}
 
 	/*
-	 * Nothing below wants the lock the VMA walk left held, and
-	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
-	 * walk carries on under mmap_lock and its own caller is what drops it,
-	 * so reporting this only tells the walk that its VMA is now stale.
+	 * Give up whatever the caller locked for us.  lru_add_drain_all() below
+	 * must not run under a lock, and the loop locks what it works on for
+	 * itself, one VMA at a time, so the caller's VMA is of no use past here.
+	 *
+	 * Which lock that is depends on how we were reached.  A range inside one
+	 * VMA arrives with that VMA read-locked and nothing else; a range that
+	 * spans VMAs arrives under mmap_lock, because try_vma_read_lock() took
+	 * it and turned the walk generic.
 	 */
-	mmap_read_unlock(mm);
-	mark_mmap_lock_dropped(madv_behavior);
+	if (vma_locked) {
+		vma_end_read(vma);
+		mark_vma_lock_dropped(madv_behavior);
+	} else {
+		mmap_read_unlock(mm);
+		mark_mmap_lock_dropped(madv_behavior);
+	}
 	vma = NULL;
 	vma_orders = 0;
 	lru_add_drain_all();
@@ -996,22 +1019,36 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 		enum scan_result result;
 
 		/*
-		 * A collapse gives the lock up, and the VMA has to be found
-		 * again after one: it can shrink while nothing is held.  A scan
-		 * that finds nothing to collapse leaves the lock alone, so a
-		 * range that is already collapsed walks it without relocking.
+		 * On another process, the reference this call holds is what
+		 * keeps the address space from being torn down -- so if it is
+		 * the only one left, the owner has gone and every page of it is
+		 * waiting on us to stop.  Nothing else here would notice: the
+		 * range is the caller's, and it can be enormous.
+		 */
+		if (mm != current->mm && collapse_test_exit_mmref(mm)) {
+			hend = addr;
+			break;
+		}
+
+		/*
+		 * A collapse gives the VMA read lock up, and the VMA has to be
+		 * found again after one: it can shrink while nothing is held.
 		 *
 		 * Reschedule only here, where nothing is held: a preemption
 		 * point under a lock is a writer waiting longer.
 		 */
 		if (!vma) {
 			cond_resched();
-			mmap_read_lock(mm);
-			vma = vma_lookup(mm, addr);
-			if (!vma) {
-				mmap_read_unlock(mm);
-				hend = addr;
-				break;
+			vma = lock_vma_under_rcu(mm, addr);
+			if (IS_ERR_OR_NULL(vma)) {
+				/*
+				 * Not only a VMA that has gone: this also fails
+				 * on one being written to right now.  Say what
+				 * is true of both -- try again.
+				 */
+				vma = NULL;
+				last_fail = SCAN_VMA_LOCK;
+				goto out;
 			}
 			vma_orders = collapse_possible_orders(vma,
 					vma->vm_flags, TVA_FORCED_COLLAPSE);
@@ -1023,7 +1060,7 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 			result = cc->scan_refusal;
 		} else {
 			/* collapse_run_pmd() takes its own locks, so give this up */
-			mmap_read_unlock(mm);
+			vma_end_read(vma);
 			vma = NULL;
 			/* The mask belonged to that lock, not to this range */
 			vma_orders = 0;
@@ -1036,7 +1073,7 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 		 * The VMA shrank under us, so the rest of the range was never
 		 * ours to collapse: stop, and expect only what came before.
 		 */
-		if (result == SCAN_VMA_NULL || result == SCAN_ADDRESS_RANGE) {
+		if (result == SCAN_ADDRESS_RANGE) {
 			hend = addr;
 			break;
 		}
@@ -1067,8 +1104,14 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 	}
 
 out:
-	/* The VMA walk this returns to expects the lock it was holding */
-	if (!vma)
+	if (vma)
+		vma_end_read(vma);
+	/*
+	 * Hand mmap_lock back only to a caller that is going to carry on with
+	 * it: the generic walk finds the next VMA under it.  The VMA fast path
+	 * applies to one VMA and returns, so it wants nothing back.
+	 */
+	if (!vma_locked)
 		mmap_read_lock(mm);
 	collapse_control_release(cc);
 	kfree(cc);
@@ -1866,7 +1909,10 @@ int madvise_walk_vmas(struct madvise_behavior *madv_behavior)
 	if (madv_behavior->lock_mode == MADVISE_VMA_READ_LOCK &&
 	    try_vma_read_lock(madv_behavior)) {
 		error = madvise_vma_behavior(madv_behavior);
-		vma_end_read(madv_behavior->vma);
+		/* A behaviour that let the VMA go has nothing left to release */
+		if (!madv_behavior->lock_dropped)
+			vma_end_read(madv_behavior->vma);
+		madv_behavior->lock_dropped = false;
 		return error;
 	}
 
@@ -1941,8 +1987,16 @@ static enum madvise_lock_mode get_lock_mode(struct madvise_behavior *madv_behavi
 	case MADV_PAGEOUT:
 	case MADV_POPULATE_READ:
 	case MADV_POPULATE_WRITE:
-	case MADV_COLLAPSE:
 		return MADVISE_MMAP_READ_LOCK;
+	case MADV_COLLAPSE:
+		/*
+		 * Only for this process.  On another one the range has to be
+		 * untagged with untagged_addr_remote(), which reads mm state
+		 * that mmap_lock protects, before any VMA is looked at.
+		 */
+		if (madv_behavior->mm != current->mm)
+			return MADVISE_MMAP_READ_LOCK;
+		fallthrough;
 	case MADV_GUARD_INSTALL:
 	case MADV_GUARD_REMOVE:
 	case MADV_DONTNEED:
-- 
2.54.0