[PATCH RESEND v5 3/4] mm/mlock: check for THP missing the mlock in try_to_unmap_one()

[Lance Yang]

The TTU_SPLIT_HUGE_PMD will no longer perform immediately, so we might
encounter a PMD-mapped THP missing the mlock in the VM_LOCKED range
during the pagewalk. It's likely necessary to mlock this THP to prevent
it from being picked up during page reclaim.

Suggested-by: Baolin Wang <baolin.wang@xxxxxxxxxxxxxxxxx>
Signed-off-by: Lance Yang <ioworker0@xxxxxxxxx>
---
 mm/rmap.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index 4c4d14325f2e..08a93347f283 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1671,7 +1671,8 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		if (!(flags & TTU_IGNORE_MLOCK) &&
 		    (vma->vm_flags & VM_LOCKED)) {
 			/* Restore the mlock which got missed */
-			if (!folio_test_large(folio))
+			if (!folio_test_large(folio) ||
+			    (!pvmw.pte && (flags & TTU_SPLIT_HUGE_PMD)))
 				mlock_vma_folio(folio, vma);
 			goto walk_done_err;
 		}
-- 
2.33.1