[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@xxxxxxx>
---
 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