2 files changed, 3 insertions, 4 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 6424ba0fce83..8ab841031436 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2896,7 +2896,7 @@ int mem_cgroup_prepare_migration(struct page *page,
 
 /* remove redundant charge if migration failed*/
 void mem_cgroup_end_migration(struct mem_cgroup *mem,
-        struct page *oldpage, struct page *newpage)
+        struct page *oldpage, struct page *newpage, bool migration_ok)
 {
         struct page *used, *unused;
         struct page_cgroup *pc;
@@ -2905,8 +2905,7 @@ void mem_cgroup_end_migration(struct mem_cgroup *mem,
                 return;
         /* blocks rmdir() */
         cgroup_exclude_rmdir(&mem->css);
-        /* at migration success, oldpage->mapping is NULL. */
-        if (oldpage->mapping) {
+        if (!migration_ok) {
                 used = oldpage;
                 unused = newpage;
         } else {
diff --git a/mm/migrate.c b/mm/migrate.c
index 5b7d1fd29621..46fe8cc13d67 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -768,7 +768,7 @@ skip_unmap:
 
 uncharge:
         if (!charge)
-                mem_cgroup_end_migration(mem, page, newpage);
+                mem_cgroup_end_migration(mem, page, newpage, rc == 0);
 unlock:
         unlock_page(page);