1 files changed, 11 insertions, 1 deletions

diff --git a/mm/z3fold.c b/mm/z3fold.c
index 985732c8b025..dfcd69d08c1e 100644
--- a/mm/z3fold.c
+++ b/mm/z3fold.c
@@ -924,7 +924,16 @@ retry:
                 set_bit(PAGE_HEADLESS, &page->private);
                 goto headless;
         }
-        __SetPageMovable(page, pool->inode->i_mapping);
+        if (can_sleep) {
+                lock_page(page);
+                __SetPageMovable(page, pool->inode->i_mapping);
+                unlock_page(page);
+        } else {
+                if (trylock_page(page)) {
+                        __SetPageMovable(page, pool->inode->i_mapping);
+                        unlock_page(page);
+                }
+        }
         z3fold_page_lock(zhdr);
 
 found:
@@ -1331,6 +1340,7 @@ static int z3fold_page_migrate(struct address_space *mapping, struct page *newpa
 
         VM_BUG_ON_PAGE(!PageMovable(page), page);
         VM_BUG_ON_PAGE(!PageIsolated(page), page);
+        VM_BUG_ON_PAGE(!PageLocked(newpage), newpage);
 
         zhdr = page_address(page);
         pool = zhdr_to_pool(zhdr);