[PATCH v2 0/8] mm/khugepaged: collapse_file() cleanups

[PATCH v2 0/8] mm/khugepaged: collapse_file() cleanups

[Pedro Falcato]

Here are a bunch of cleanups for collapse_file() that attempt to aid
code readability and understanding. This is part 1 out of N, and only
attempts to address the main collapse loop. The exit path is still a
mess but I figure it's much easier to review if I send this batch as-is,
versus sending the whole thing.

Tested using the khugepaged shmem and file selftests.

Based on top of latest mm-core/for-next.

Note that some review comments from v1 went unaddressed
(Nico's larger restructuring, some small Lorenzo suggestions) for a couple
of small reasons. The main one being that I want to get this v2 out before
going insane :)) I agreed to work with Nico, and I'll either send out some
of his patches as a followup, or he will.

In any case, I'm tired of this function and I hope this kicks off some
discussion :)

Cc: Baolin Wang <baolin.wang@linux.alibaba.com>
Cc: "Liam R. Howlett" <liam@infradead.org>
Cc: Nico Pache <nico.pache@linux.dev>
Cc: Ryan Roberts <ryan.roberts@arm.com>
Cc: Dev Jain <dev.jain@arm.com>
Cc: Barry Song <baohua@kernel.org>
Cc: Lance Yang <lance.yang@linux.dev>
Cc: Usama Arif <usama.arif@linux.dev>
Cc: Lorenzo Stoakes <ljs@kernel.org>
Cc: Kiryl Shutsemau <kas@kernel.org>
Cc: linux-mm@kvack.org
Cc: linux-kernel@vger.kernel.org

Changes in v2:
- Rebased on current mm-core/for-next, which contains a number of fixes to
  the amazingly simple and awesome collapse_file() function.
- Fix the "bad unlock balance" reported by syzbot ci (syzbot ci, Lorenzo)
- 1/8:
  - Fold in v1 5/8 (kerneldoc); new helpers get their kerneldoc in
    the patch that introduces them. (Lorenzo)
  - Group the input fields of collapse_file_state and make end and
    is_shmem const; is_shmem is now a bool. (Lorenzo)
  - Wrap the prepare_collapse_file_folio() declaration to fit in 80
    columns. (Lorenzo)
  - Keep the original LRU drain behaviour (drain only after a
    successful shmem_get_folio()) (Lorenzo)
- 2/8: New patch. Turn value entries into NULL before calling
  prepare_collapse_file_folio().
- 3/8 (v1 2/8):
  - Rename collapse_read_folio() to maybe_read_folio(). (Lorenzo)
  - Replace the compound condition with an early-return guard.
    (Lorenzo)
  - Comment on why the xa_lock is dropped. (Lorenzo)
  - Expand the commit message. (Lorenzo)
- 5/8 (v1 4/8): Drop a duplicate blank line.
- 6/8: Add a commit message, noting that errors before try_to_unmap()
  no longer call try_to_unmap_flush() themselves.
- 7/8: Explain in the commit message why the refcount and dirty checks
  don't need the i_pages lock. The flush added by 854e5ff8b3ce now
  happens at out_putback.
- 8/8: Retitle (drop "fix"), since 854e5ff8b3ce already reworded the
  comment; drop a stray blank line.

  Link to v1: https://lore.kernel.org/all/20260720142913.846902-1-pfalcato@suse.de/

Pedro Falcato (8):
  mm/khugepaged: separate out windy folio logic from collapse_file
  mm/khugepaged: sanitize the folio passed to file folio collapse logic
  mm/khugepaged: factor out page cache folio reading
  mm/khugepaged: factor out and simplify dirty/writeback handling
  mm/khugepaged: simplify prepare folio locking and exit paths
  mm/khugepaged: hoist isolation into collapse_isolate_folio()
  mm/khugepaged: hoist more code into collapse_isolate_folio()
  mm/khugepaged: flesh out try_to_unmap_flush() comment

 mm/khugepaged.c | 461 ++++++++++++++++++++++++++++++------------------
 1 file changed, 286 insertions(+), 175 deletions(-)


base-commit: 8d21e3d3faf19f461f6d0380fe8a1f5bdba97f25
-- 
2.56.0

[PATCH v2 1/8] mm/khugepaged: separate out windy folio logic from collapse_file

[Pedro Falcato]

Separate out complex folio-related logic from the main collapse_file()
loop, and introduce a new helper struct to help marshal arguments back
and forth.

No functional change intended.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 202 ++++++++++++++++++++++++++++++------------------
 1 file changed, 128 insertions(+), 74 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 24eb8281ea3c..0fe19acf470b 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2211,6 +2211,105 @@ static void retract_page_tables(struct address_space *mapping, pgoff_t pgoff)
 	i_mmap_unlock_read(mapping);
 }
 
+struct collapse_file_state {
+	/* Inputs */
+	struct address_space *mapping;
+	struct file *file;
+	struct xa_state *xas;
+	/* folio is in-out; preparing might modify it */
+	struct folio *folio;
+	/* collapse end index */
+	const pgoff_t end;
+	const bool is_shmem : 1;
+	bool do_readahead : 1;
+	bool readahead_done : 1;
+};
+
+/**
+ * prepare_collapse_file_folio() - Prepare the collapsing of a single file folio
+ * @index: folio index in the page cache
+ * @state: helper struct of state being passed back and forth.
+ *
+ * prepare_collapse_file_folio() helps prepare the THP collapsing of file folios
+ * (both shmem and !shmem), including the reading of folios, initial tests for
+ * dirty/writeback, locking, etc.
+ *
+ * Context: Expects @state->mapping->i_pages.xa_lock to be held. Releases the same xa_lock.
+ *          On success, returns a refcounted and locked folio in @state->folio (that may
+ *          not be the same that was passed in).
+ *          If readahead needs to be (and can be) done, sets state->do_readahead.
+ * Return: Appropriate scan_result value (on success, SCAN_SUCCEED).
+ */
+static enum scan_result
+prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
+{
+	struct address_space *mapping = state->mapping;
+	enum scan_result result = SCAN_SUCCEED;
+	const bool is_shmem = state->is_shmem;
+	struct folio *folio = state->folio;
+
+	if (is_shmem) {
+		if (xa_is_value(folio) || !folio_test_uptodate(folio)) {
+			xas_unlock_irq(state->xas);
+			/* swap in or instantiate fallocated page */
+			if (shmem_get_folio(mapping->host, index, 0,
+					&folio, SGP_NOALLOC)) {
+				result = SCAN_FAIL;
+				goto xa_unlocked;
+			}
+			/* drain lru cache to help folio_isolate_lru() */
+			lru_add_drain();
+			goto xa_unlocked;
+		} else if (folio_trylock(folio)) {
+			folio_get(folio);
+		} else {
+			result = SCAN_PAGE_LOCK;
+			goto xa_locked;
+		}
+	} else {	/* !is_shmem */
+		if (!folio || xa_is_value(folio)) {
+			state->do_readahead = !state->readahead_done;
+			result = SCAN_FAIL;
+			goto xa_unlocked;
+		} else if (folio_test_dirty(folio)) {
+			/*
+			 * This page is dirty because it hasn't
+			 * been flushed since first write.
+			 *
+			 * Trigger async flush for read-only files and
+			 * hope the writeback is done when khugepaged
+			 * revisits this page. Writable files can have
+			 * their folios dirty at any time; blindly
+			 * flushing them would cause undesirable
+			 * system-wide writeback.
+			 *
+			 * This is a one-off situation. We are not
+			 * forcing writeback in loop.
+			 */
+			xas_unlock_irq(state->xas);
+			if (!inode_is_open_for_write(mapping->host))
+				filemap_flush(mapping);
+			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+			goto xa_unlocked;
+		} else if (folio_test_writeback(folio)) {
+			xas_unlock_irq(state->xas);
+			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+			goto xa_unlocked;
+		} else if (folio_trylock(folio)) {
+			folio_get(folio);
+		} else {
+			result = SCAN_PAGE_LOCK;
+			goto xa_locked;
+		}
+	}
+
+xa_locked:
+	xas_unlock_irq(state->xas);
+xa_unlocked:
+	state->folio = folio;
+	return result;
+}
+
 /**
  * collapse_file - collapse filemap/tmpfs/shmem pages into huge one.
  *
@@ -2249,8 +2348,13 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 	XA_STATE_ORDER(xas, &mapping->i_pages, start, HPAGE_PMD_ORDER);
 	enum scan_result result = SCAN_SUCCEED;
 	int nr_none = 0;
-	bool do_readahead = false;
-	bool readahead_done = false;
+	struct collapse_file_state state = {
+		.is_shmem = is_shmem,
+		.xas = &xas,
+		.mapping = mapping,
+		.file = file,
+		.end = end,
+	};
 
 	/*
 	 * MADV_COLLAPSE ignores shmem huge config, so do not check shmem
@@ -2270,7 +2374,7 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 
 retry:
 	result = SCAN_SUCCEED;
-	do_readahead = false;
+	state.do_readahead = false;
 	xas_set_order(&xas, start, HPAGE_PMD_ORDER);
 	__folio_set_locked(new_folio);
 	if (is_shmem)
@@ -2299,79 +2403,29 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 		folio = xas_load(&xas);
 
 		VM_BUG_ON(index != xas.xa_index);
-		if (is_shmem) {
-			if (!folio) {
-				/*
-				 * Stop if extent has been truncated or
-				 * hole-punched, and is now completely
-				 * empty.
-				 */
-				if (index == start) {
-					if (!xas_next_entry(&xas, end - 1)) {
-						result = SCAN_TRUNCATED;
-						goto xa_locked;
-					}
-				}
-				nr_none++;
-				index++;
-				continue;
-			}
-
-			if (xa_is_value(folio) || !folio_test_uptodate(folio)) {
-				xas_unlock_irq(&xas);
-				/* swap in or instantiate fallocated page */
-				if (shmem_get_folio(mapping->host, index, 0,
-						&folio, SGP_NOALLOC)) {
-					result = SCAN_FAIL;
-					goto xa_unlocked;
+		if (is_shmem && !folio) {
+			/*
+			 * Stop if extent has been truncated or
+			 * hole-punched, and is now completely
+			 * empty.
+			 */
+			if (index == start) {
+				if (!xas_next_entry(&xas, end - 1)) {
+					result = SCAN_TRUNCATED;
+					goto xa_locked;
 				}
-				/* drain lru cache to help folio_isolate_lru() */
-				lru_add_drain();
-			} else if (folio_trylock(folio)) {
-				folio_get(folio);
-				xas_unlock_irq(&xas);
-			} else {
-				result = SCAN_PAGE_LOCK;
-				goto xa_locked;
-			}
-		} else {	/* !is_shmem */
-			if (!folio || xa_is_value(folio)) {
-				do_readahead = !readahead_done;
-				result = SCAN_FAIL;
-				goto xa_locked;
-			} else if (folio_test_dirty(folio)) {
-				/*
-				 * This page is dirty because it hasn't
-				 * been flushed since first write.
-				 *
-				 * Trigger async flush for read-only files and
-				 * hope the writeback is done when khugepaged
-				 * revisits this page. Writable files can have
-				 * their folios dirty at any time; blindly
-				 * flushing them would cause undesirable
-				 * system-wide writeback.
-				 *
-				 * This is a one-off situation. We are not
-				 * forcing writeback in loop.
-				 */
-				xas_unlock_irq(&xas);
-				if (!inode_is_open_for_write(mapping->host))
-					filemap_flush(mapping);
-				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-				goto xa_unlocked;
-			} else if (folio_test_writeback(folio)) {
-				xas_unlock_irq(&xas);
-				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-				goto xa_unlocked;
-			} else if (folio_trylock(folio)) {
-				folio_get(folio);
-				xas_unlock_irq(&xas);
-			} else {
-				result = SCAN_PAGE_LOCK;
-				goto xa_locked;
 			}
+			nr_none++;
+			index++;
+			continue;
 		}
 
+		/* At this point folio can be NULL, or a value. */
+		state.folio = folio;
+		result = prepare_collapse_file_folio(index, &state);
+		folio = state.folio;
+		if (result != SCAN_SUCCEED)
+			goto xa_unlocked;
 		/*
 		 * The folio must be locked, so we can drop the i_pages lock
 		 * without racing with truncate.
@@ -2667,7 +2721,7 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 	folio_unlock(new_folio);
 
 	/* Readahead may take invalidate_lock; source folios are unlocked. */
-	if (do_readahead) {
+	if (state.do_readahead) {
 		for (index = start; index < end;) {
 			folio = filemap_lock_folio(mapping, index);
 			if (IS_ERR(folio)) {
@@ -2685,7 +2739,7 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 		}
 		/* Drain the local LRU batch before retrying isolation. */
 		lru_add_drain();
-		readahead_done = true;
+		state.readahead_done = true;
 		goto retry;
 	}
 out_put:
-- 
2.56.0

[PATCH v2 2/8] mm/khugepaged: sanitize the folio passed to file folio collapse logic

[Pedro Falcato]

At that point, the folio (which is taken from the i_pages xarray
directly) can assume 3 different kinds of values:

1) NULL. This is the complete absence of value from the xarray.
2) An xa_value. This is generally a shadow entry in the page cache, or a
   swap entry in shmem page cache.
3) A real struct folio*.

This is confusing. Collapse NULL and xa_value into NULL right at the top
of the loop. Collapse logic generally doesn't care whether something is
fully absent, or whether there's some sort of underlying shadow entry.

This simplifies a bit of reading logic in prepare_collapse_file().
No functional change intended.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 0fe19acf470b..ac148db2a4a0 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2225,6 +2225,19 @@ struct collapse_file_state {
 	bool readahead_done : 1;
 };
 
+/**
+ * real_folio() - Convert a page cache folio * into a real folio *
+ * @folio: struct folio * to convert.
+ *
+ * Return: NULL if the passed folio is invalid (XA value, NULL), otherwise @folio.
+ */
+static inline struct folio *real_folio(struct folio *folio)
+{
+	if (!folio || xa_is_value(folio))
+		return NULL;
+	return folio;
+}
+
 /**
  * prepare_collapse_file_folio() - Prepare the collapsing of a single file folio
  * @index: folio index in the page cache
@@ -2249,7 +2262,7 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 	struct folio *folio = state->folio;
 
 	if (is_shmem) {
-		if (xa_is_value(folio) || !folio_test_uptodate(folio)) {
+		if (!folio || !folio_test_uptodate(folio)) {
 			xas_unlock_irq(state->xas);
 			/* swap in or instantiate fallocated page */
 			if (shmem_get_folio(mapping->host, index, 0,
@@ -2267,7 +2280,7 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 			goto xa_locked;
 		}
 	} else {	/* !is_shmem */
-		if (!folio || xa_is_value(folio)) {
+		if (!folio) {
 			state->do_readahead = !state->readahead_done;
 			result = SCAN_FAIL;
 			goto xa_unlocked;
@@ -2420,8 +2433,12 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 			continue;
 		}
 
-		/* At this point folio can be NULL, or a value. */
-		state.folio = folio;
+		/*
+		 * At this point folio can be a value. Make values NULL. The
+		 * value itself isn't interesting here. Collapse only cares
+		 * about the absence of a folio.
+		 */
+		state.folio = real_folio(folio);
 		result = prepare_collapse_file_folio(index, &state);
 		folio = state.folio;
 		if (result != SCAN_SUCCEED)
-- 
2.56.0

[PATCH v2 3/8] mm/khugepaged: factor out page cache folio reading

[Pedro Falcato]

The two reading mechanisms (shmem, filemap) have similar purposes but
vary subtly. There is no need to have two variants.

Factor out the do_readahead toggling (for regular files) and
shmem_get_folio (for shmem) into its own separate function,
maybe_read_folio(), shared by both paths.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 77 +++++++++++++++++++++++++++++++++++++------------
 1 file changed, 59 insertions(+), 18 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index ac148db2a4a0..d5647a3ac095 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2238,6 +2238,46 @@ static inline struct folio *real_folio(struct folio *folio)
 	return folio;
 }
 
+/**
+ * maybe_read_folio() - Ensure a folio is uptodate
+ * @index: folio index in the page cache
+ * @state: collapse_file state helper struct
+ * @new: if a new folio is returned, *new = true.
+ *
+ * maybe_read_folio() ensures a folio is uptodate and, if not, either reads it
+ * in synchronously (for shmem), or schedules readahead (for after the collapse
+ * logic untangles itself from all the folio locks it's holding).
+ *
+ * Context: Expects @state->mapping->i_pages.xa_lock to be held. Releases that
+ * same lock on error, or if *@new = true.
+ * Return: Valid folio (locked and uptodate iff *@new), or NULL on error.
+ */
+static struct folio *
+maybe_read_folio(pgoff_t index, struct collapse_file_state *state, bool *new)
+{
+	struct folio *folio = real_folio(state->folio);
+	const bool may_bring_uptodate = state->is_shmem;
+
+	/* Perfectly cromulent, uptodate folio. Proceed. */
+	if (folio && (!may_bring_uptodate || folio_test_uptodate(folio)))
+		return folio;
+
+	/* We're either returning NULL, or starting IO. Drop the lock now. */
+	xas_unlock_irq(state->xas);
+	if (!state->is_shmem) {
+		/* Folio is missing? Let's start readahead, if we haven't already. */
+		state->do_readahead = !state->readahead_done;
+		return NULL;
+	}
+
+	*new = true;
+	/* swap in or instantiate fallocated page */
+	if (shmem_get_folio(state->mapping->host, index, 0,
+				 &folio, SGP_NOALLOC))
+		return NULL;
+	return folio;
+}
+
 /**
  * prepare_collapse_file_folio() - Prepare the collapsing of a single file folio
  * @index: folio index in the page cache
@@ -2259,32 +2299,33 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 	struct address_space *mapping = state->mapping;
 	enum scan_result result = SCAN_SUCCEED;
 	const bool is_shmem = state->is_shmem;
-	struct folio *folio = state->folio;
+	struct folio *folio;
+	bool new = false;
+
+	folio = maybe_read_folio(index, state, &new);
+	if (!folio)
+		return SCAN_FAIL;
+	if (new) {
+		/*
+		 * maybe_read_folio() got a new folio. This folio is locked
+		 * and ref'd up. Nothing tricky (dirty, writeback, etc) is likely to
+		 * have happened, so bail now. But before that, drain the local LRU
+		 * add batch. Otherwise, it's very possible folio_isolate_lru()
+		 * will not succeed.
+		 */
+		lru_add_drain();
+		goto xa_unlocked;
+	}
 
 	if (is_shmem) {
-		if (!folio || !folio_test_uptodate(folio)) {
-			xas_unlock_irq(state->xas);
-			/* swap in or instantiate fallocated page */
-			if (shmem_get_folio(mapping->host, index, 0,
-					&folio, SGP_NOALLOC)) {
-				result = SCAN_FAIL;
-				goto xa_unlocked;
-			}
-			/* drain lru cache to help folio_isolate_lru() */
-			lru_add_drain();
-			goto xa_unlocked;
-		} else if (folio_trylock(folio)) {
+		if (folio_trylock(folio)) {
 			folio_get(folio);
 		} else {
 			result = SCAN_PAGE_LOCK;
 			goto xa_locked;
 		}
 	} else {	/* !is_shmem */
-		if (!folio) {
-			state->do_readahead = !state->readahead_done;
-			result = SCAN_FAIL;
-			goto xa_unlocked;
-		} else if (folio_test_dirty(folio)) {
+		if (folio_test_dirty(folio)) {
 			/*
 			 * This page is dirty because it hasn't
 			 * been flushed since first write.
-- 
2.56.0

[PATCH v2 4/8] mm/khugepaged: factor out and simplify dirty/writeback handling

[Pedro Falcato]

Dirty/writeback folios can easily share most of the branch. These checks
are still gated under !is_shmem due to shmem not really having issues
with dirty folios.

While here, reflow the comment.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 47 ++++++++++++++++++++++++-----------------------
 1 file changed, 24 insertions(+), 23 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index d5647a3ac095..0eee2b6bbe64 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2299,8 +2299,8 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 	struct address_space *mapping = state->mapping;
 	enum scan_result result = SCAN_SUCCEED;
 	const bool is_shmem = state->is_shmem;
+	bool dirty, new = false;
 	struct folio *folio;
-	bool new = false;
 
 	folio = maybe_read_folio(index, state, &new);
 	if (!folio)
@@ -2317,39 +2317,40 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 		goto xa_unlocked;
 	}
 
-	if (is_shmem) {
-		if (folio_trylock(folio)) {
-			folio_get(folio);
-		} else {
-			result = SCAN_PAGE_LOCK;
-			goto xa_locked;
-		}
-	} else {	/* !is_shmem */
-		if (folio_test_dirty(folio)) {
+	if (!is_shmem) {
+		dirty = folio_test_dirty(folio);
+		if (dirty || folio_test_writeback(folio)) {
 			/*
-			 * This page is dirty because it hasn't
-			 * been flushed since first write.
+			 * This folio is either dirty or under writeback.
+			 * khugepaged cannot operate on such folios.
 			 *
-			 * Trigger async flush for read-only files and
-			 * hope the writeback is done when khugepaged
-			 * revisits this page. Writable files can have
-			 * their folios dirty at any time; blindly
-			 * flushing them would cause undesirable
+			 * For dirty folios, trigger async flush for
+			 * read-only files and hope the writeback is done
+			 * when khugepaged revisits this page. Writable
+			 * files can have their folios dirty at any time;
+			 * blindly flushing them would cause undesirable
 			 * system-wide writeback.
 			 *
 			 * This is a one-off situation. We are not
 			 * forcing writeback in loop.
 			 */
 			xas_unlock_irq(state->xas);
-			if (!inode_is_open_for_write(mapping->host))
+			if (dirty && !inode_is_open_for_write(mapping->host))
 				filemap_flush(mapping);
 			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
 			goto xa_unlocked;
-		} else if (folio_test_writeback(folio)) {
-			xas_unlock_irq(state->xas);
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			goto xa_unlocked;
-		} else if (folio_trylock(folio)) {
+		}
+	}
+
+	if (is_shmem) {
+		if (folio_trylock(folio)) {
+			folio_get(folio);
+		} else {
+			result = SCAN_PAGE_LOCK;
+			goto xa_locked;
+		}
+	} else {	/* !is_shmem */
+		if (folio_trylock(folio)) {
 			folio_get(folio);
 		} else {
 			result = SCAN_PAGE_LOCK;
-- 
2.56.0

[PATCH v2 5/8] mm/khugepaged: simplify prepare folio locking and exit paths

[Pedro Falcato]

folio_trylock() + folio_get() do the exact same thing on both is_shmem
and !is_shmem paths. Unify them into one path.

At the end of the day, there are two large invariants for the function:
 - xarray lock is not held at exit. Either it was dropped on error, or
   dropped on success/IO.
 - success always means a locked-and-ref'd-up folio.

After this transformation, it's trivial to notice how simple the error
paths are, so do away with the gotos and hoist some code.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 33 +++++++++++----------------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 0eee2b6bbe64..236df9be7e27 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2314,7 +2314,8 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 		 * will not succeed.
 		 */
 		lru_add_drain();
-		goto xa_unlocked;
+		state->folio = folio;
+		return SCAN_SUCCEED;
 	}
 
 	if (!is_shmem) {
@@ -2337,31 +2338,19 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 			xas_unlock_irq(state->xas);
 			if (dirty && !inode_is_open_for_write(mapping->host))
 				filemap_flush(mapping);
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			goto xa_unlocked;
-		}
-	}
-
-	if (is_shmem) {
-		if (folio_trylock(folio)) {
-			folio_get(folio);
-		} else {
-			result = SCAN_PAGE_LOCK;
-			goto xa_locked;
-		}
-	} else {	/* !is_shmem */
-		if (folio_trylock(folio)) {
-			folio_get(folio);
-		} else {
-			result = SCAN_PAGE_LOCK;
-			goto xa_locked;
+			return SCAN_PAGE_DIRTY_OR_WRITEBACK;
 		}
 	}
 
-xa_locked:
+	/*
+	 * Note: trylock + simple folio_get() is safe here due to the i_pages
+	 * lock being held; this excludes truncation happening in parallel.
+	 */
+	if (!folio_trylock(folio))
+		result = SCAN_PAGE_LOCK;
+	else
+		folio_get(folio);
 	xas_unlock_irq(state->xas);
-xa_unlocked:
-	state->folio = folio;
 	return result;
 }
 
-- 
2.56.0

[PATCH v2 6/8] mm/khugepaged: hoist isolation into collapse_isolate_folio()

[Pedro Falcato]

Move unwindy folio isolation logic from the main loop to
collapse_isolate_folio().

As a small side effect, try_to_unmap_flush() is no longer called for
errors prior to try_to_unmap(). This should be fine as there are other
mechanisms in place to flush it before references/locks are dropped.

Aside from that, no functional changes intended.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 118 +++++++++++++++++++++++++++---------------------
 1 file changed, 66 insertions(+), 52 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 236df9be7e27..1e58ae5b3e74 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2354,6 +2354,69 @@ prepare_collapse_file_folio(pgoff_t index, struct collapse_file_state *state)
 	return result;
 }
 
+static enum scan_result collapse_isolate_folio(struct collapse_file_state *state)
+{
+	struct folio *folio = state->folio;
+	enum scan_result result;
+
+	/*
+	 * The folio must be locked, so we can drop the i_pages lock
+	 * without racing with truncate.
+	 */
+	VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
+
+	/* make sure the folio is up to date */
+	if (unlikely(!folio_test_uptodate(folio))) {
+		result = SCAN_FAIL;
+		goto out_unlock;
+	}
+
+	/*
+	 * If file was truncated then extended, or hole-punched, before
+	 * we locked the first folio, then a THP might be there already.
+	 * This will be discovered on the first iteration.
+	 */
+	if (is_pmd_order(folio_order(folio))) {
+		result = SCAN_PTE_MAPPED_HUGEPAGE;
+		goto out_unlock;
+	}
+
+	if (folio_mapping(folio) != state->mapping) {
+		result = SCAN_TRUNCATED;
+		goto out_unlock;
+	}
+
+	if (!state->is_shmem && (folio_test_dirty(folio) ||
+				 folio_test_writeback(folio))) {
+		/*
+		 * khugepaged only works on clean file-backed folios,
+		 * so this folio is dirty because it hasn't been flushed
+		 * since first write.
+		 */
+		result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+		goto out_unlock;
+	}
+
+	if (!folio_isolate_lru(folio)) {
+		result = SCAN_DEL_PAGE_LRU;
+		goto out_unlock;
+	}
+
+	if (!filemap_release_folio(folio, GFP_KERNEL)) {
+		result = SCAN_PAGE_HAS_PRIVATE;
+		folio_putback_lru(folio);
+		goto out_unlock;
+	}
+
+	if (folio_mapped(folio))
+		try_to_unmap(folio, TTU_IGNORE_MLOCK | TTU_BATCH_FLUSH);
+	return SCAN_SUCCEED;
+out_unlock:
+	folio_unlock(folio);
+	folio_put(folio);
+	return result;
+}
+
 /**
  * collapse_file - collapse filemap/tmpfs/shmem pages into huge one.
  *
@@ -2474,61 +2537,12 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 		folio = state.folio;
 		if (result != SCAN_SUCCEED)
 			goto xa_unlocked;
-		/*
-		 * The folio must be locked, so we can drop the i_pages lock
-		 * without racing with truncate.
-		 */
-		VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
-
-		/* make sure the folio is up to date */
-		if (unlikely(!folio_test_uptodate(folio))) {
-			result = SCAN_FAIL;
-			goto out_unlock;
-		}
 
-		/*
-		 * If file was truncated then extended, or hole-punched, before
-		 * we locked the first folio, then a THP might be there already.
-		 * This will be discovered on the first iteration.
-		 */
-		if (is_pmd_order(folio_order(folio))) {
-			result = SCAN_PTE_MAPPED_HUGEPAGE;
-			goto out_unlock;
-		}
-
-		if (folio_mapping(folio) != mapping) {
-			result = SCAN_TRUNCATED;
-			goto out_unlock;
-		}
-
-		if (!is_shmem && (folio_test_dirty(folio) ||
-				  folio_test_writeback(folio))) {
-			/*
-			 * khugepaged only works on clean file-backed folios,
-			 * so this folio is dirty because it hasn't been flushed
-			 * since first write.
-			 */
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			goto out_unlock;
-		}
-
-		if (!folio_isolate_lru(folio)) {
-			result = SCAN_DEL_PAGE_LRU;
-			goto out_unlock;
-		}
-
-		if (!filemap_release_folio(folio, GFP_KERNEL)) {
-			result = SCAN_PAGE_HAS_PRIVATE;
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
-		if (folio_mapped(folio))
-			try_to_unmap(folio,
-					TTU_IGNORE_MLOCK | TTU_BATCH_FLUSH);
+		result = collapse_isolate_folio(&state);
+		if (result != SCAN_SUCCEED)
+			goto xa_unlocked;
 
 		xas_lock_irq(&xas);
-
 		VM_BUG_ON_FOLIO(folio != xa_load(xas.xa, index), folio);
 
 		/*
-- 
2.56.0

[PATCH v2 7/8] mm/khugepaged: hoist more code into collapse_isolate_folio()

[Pedro Falcato]

Checking expected refcounts or dirtiness has no need for the i_pages
xarray lock.

Page cache lookups are generally done locklessly, only requiring the RCU
lock and folio_try_get(). This makes it so refcount stability isn't
guaranteed under the i_pages lock.

Testing for dirtiness is in no way excluding anything with the i_pages
lock. Rather, the folio lock plays a more important role here.

Thus, move this logic into collapse_isolate_folio(), without the lock.

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 93 ++++++++++++++++++++++---------------------------
 1 file changed, 42 insertions(+), 51 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 1e58ae5b3e74..accd8b4895d3 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2404,13 +2404,53 @@ static enum scan_result collapse_isolate_folio(struct collapse_file_state *state
 
 	if (!filemap_release_folio(folio, GFP_KERNEL)) {
 		result = SCAN_PAGE_HAS_PRIVATE;
-		folio_putback_lru(folio);
-		goto out_unlock;
+		goto out_putback;
 	}
 
 	if (folio_mapped(folio))
 		try_to_unmap(folio, TTU_IGNORE_MLOCK | TTU_BATCH_FLUSH);
+	/*
+	 * We control 2 + nr_pages references to the folio:
+	 *  - we hold a pin on it;
+	 *  - nr_pages reference from page cache;
+	 *  - one from lru_isolate_folio;
+	 * If those are the only references, then any new usage
+	 * of the folio will have to fetch it from the page
+	 * cache. That requires locking the folio to handle
+	 * truncate, so any new usage will be blocked until we
+	 * unlock folio after collapse/during rollback.
+	 */
+	if (folio_ref_count(folio) != 2 + folio_nr_pages(folio)) {
+		result = SCAN_PAGE_COUNT;
+		goto out_putback;
+	}
+
+	/*
+	 * At this point, the folio is locked and unmapped. If the PTE
+	 * was dirty, try_to_unmap() has transferred the dirty bit to
+	 * the folio and we must not collapse it into a clean
+	 * file-backed folio.
+	 *
+	 * If the folio is clean here, no one can write it until we
+	 * drop the folio lock. A write through a stale TLB entry came
+	 * from a clean PTE and must fault because the PTE has been
+	 * cleared; the fault path has to take the folio lock before
+	 * installing a writable mapping. Buffered write paths also
+	 * have to take the folio lock before modifying file contents
+	 * without a mapping, typically via write_begin_get_folio().
+	 */
+	if (!state->is_shmem && folio_test_dirty(folio)) {
+		result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+		goto out_putback;
+	}
 	return SCAN_SUCCEED;
+out_putback:
+	/*
+	 * The folio may have been unmapped with TTU_BATCH_FLUSH.
+	 * Flush before releasing the lock and our last reference.
+	 */
+	try_to_unmap_flush();
+	folio_putback_lru(folio);
 out_unlock:
 	folio_unlock(folio);
 	folio_put(folio);
@@ -2545,60 +2585,11 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 		xas_lock_irq(&xas);
 		VM_BUG_ON_FOLIO(folio != xa_load(xas.xa, index), folio);
 
-		/*
-		 * We control 2 + nr_pages references to the folio:
-		 *  - we hold a pin on it;
-		 *  - nr_pages reference from page cache;
-		 *  - one from lru_isolate_folio;
-		 * If those are the only references, then any new usage
-		 * of the folio will have to fetch it from the page
-		 * cache. That requires locking the folio to handle
-		 * truncate, so any new usage will be blocked until we
-		 * unlock folio after collapse/during rollback.
-		 */
-		if (folio_ref_count(folio) != 2 + folio_nr_pages(folio)) {
-			result = SCAN_PAGE_COUNT;
-			xas_unlock_irq(&xas);
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
-		/*
-		 * At this point, the folio is locked and unmapped. If the PTE
-		 * was dirty, try_to_unmap() has transferred the dirty bit to
-		 * the folio and we must not collapse it into a clean
-		 * file-backed folio.
-		 *
-		 * If the folio is clean here, no one can write it until we
-		 * drop the folio lock. A write through a stale TLB entry came
-		 * from a clean PTE and must fault because the PTE has been
-		 * cleared; the fault path has to take the folio lock before
-		 * installing a writable mapping. Buffered write paths also
-		 * have to take the folio lock before modifying file contents
-		 * without a mapping, typically via write_begin_get_folio().
-		 */
-		if (!is_shmem && folio_test_dirty(folio)) {
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			xas_unlock_irq(&xas);
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
 		/*
 		 * Accumulate the folios that are being collapsed.
 		 */
 		list_add_tail(&folio->lru, &pagelist);
 		index += folio_nr_pages(folio);
-		continue;
-out_unlock:
-		/*
-		 * The folio may have been unmapped with TTU_BATCH_FLUSH.
-		 * Flush before releasing the lock and our last reference.
-		 */
-		try_to_unmap_flush();
-		folio_unlock(folio);
-		folio_put(folio);
-		goto xa_unlocked;
 	}
 
 xa_locked:
-- 
2.56.0

[PATCH v2 8/8] mm/khugepaged: flesh out try_to_unmap_flush() comment

[Pedro Falcato]

Flesh out any correctness concerns with try_to_unmap batching, flushing
and its interaction with the rest of collapse_file().

Signed-off-by: Pedro Falcato <pfalcato@suse.de>
---
 mm/khugepaged.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index accd8b4895d3..fd947e698ed0 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2597,8 +2597,12 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 xa_unlocked:
 
 	/*
-	 * Flush before copying the folios, or releasing them in rollback.
-	 * This is a no-op if out_unlock already flushed the batch.
+	 * try_to_unmap() flush must be done now before copying, regardless
+	 * of success or not. In case of success, folios are about to be
+	 * copied and collapsed onto a single large folio. In that case,
+	 * stale TLB entries need to be flushed out, so no racing write may
+	 * get lost. In case of failure, stale TLB entries need to be flushed
+	 * out before putting the folio (which can possibly free it).
 	 */
 	try_to_unmap_flush();
 
-- 
2.56.0