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

[Pedro Falcato <pfalcato@suse.de>]

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