1 files changed, 40 insertions, 24 deletions

diff --git a/mm/migrate.c b/mm/migrate.c
index bd3fdc202e8b..c74412b381ff 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -915,6 +915,19 @@ out_unlock:
                 put_anon_vma(anon_vma);
         unlock_page(page);
 out:
+        /*
+         * If migration is successful, decrease refcount of the newpage
+         * which will not free the page because new page owner increased
+         * refcounter. As well, if it is LRU page, add the page to LRU
+         * list in here.
+         */
+        if (rc == MIGRATEPAGE_SUCCESS) {
+                if (unlikely(__is_movable_balloon_page(newpage)))
+                        put_page(newpage);
+                else
+                        putback_lru_page(newpage);
+        }
+
         return rc;
 }
 
@@ -948,6 +961,12 @@ static ICE_noinline int unmap_and_move(new_page_t get_new_page,
 
         if (page_count(page) == 1) {
                 /* page was freed from under us. So we are done. */
+                ClearPageActive(page);
+                ClearPageUnevictable(page);
+                if (put_new_page)
+                        put_new_page(newpage, private);
+                else
+                        put_page(newpage);
                 goto out;
         }
 
@@ -960,10 +979,8 @@ static ICE_noinline int unmap_and_move(new_page_t get_new_page,
         }
 
         rc = __unmap_and_move(page, newpage, force, mode);
-        if (rc == MIGRATEPAGE_SUCCESS) {
-                put_new_page = NULL;
+        if (rc == MIGRATEPAGE_SUCCESS)
                 set_page_owner_migrate_reason(newpage, reason);
-        }
 
 out:
         if (rc != -EAGAIN) {
@@ -976,34 +993,33 @@ out:
                 list_del(&page->lru);
                 dec_zone_page_state(page, NR_ISOLATED_ANON +
                                 page_is_file_cache(page));
-                /* Soft-offlined page shouldn't go through lru cache list */
-                if (reason == MR_MEMORY_FAILURE && rc == MIGRATEPAGE_SUCCESS) {
+        }
+
+        /*
+         * If migration is successful, releases reference grabbed during
+         * isolation. Otherwise, restore the page to right list unless
+         * we want to retry.
+         */
+        if (rc == MIGRATEPAGE_SUCCESS) {
+                put_page(page);
+                if (reason == MR_MEMORY_FAILURE) {
                         /*
-                         * With this release, we free successfully migrated
-                         * page and set PG_HWPoison on just freed page
-                         * intentionally. Although it's rather weird, it's how
-                         * HWPoison flag works at the moment.
+                         * Set PG_HWPoison on just freed page
+                         * intentionally. Although it's rather weird,
+                         * it's how HWPoison flag works at the moment.
                          */
-                        put_page(page);
                         if (!test_set_page_hwpoison(page))
                                 num_poisoned_pages_inc();
-                } else
+                }
+        } else {
+                if (rc != -EAGAIN)
                         putback_lru_page(page);
+                if (put_new_page)
+                        put_new_page(newpage, private);
+                else
+                        put_page(newpage);
         }
 
-        /*
-         * If migration was not successful and there's a freeing callback, use
-         * it.  Otherwise, putback_lru_page() will drop the reference grabbed
-         * during isolation.
-         */
-        if (put_new_page)
-                put_new_page(newpage, private);
-        else if (unlikely(__is_movable_balloon_page(newpage))) {
-                /* drop our reference, page already in the balloon */
-                put_page(newpage);
-        } else
-                putback_lru_page(newpage);
-
         if (result) {
                 if (rc)
                         *result = rc;