1 files changed, 15 insertions, 7 deletions

diff --git a/mm/migrate.c b/mm/migrate.c
index 376cceba82f9..f6d7f8efd1a8 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -358,6 +358,10 @@ static int migrate_page_move_mapping(struct address_space *mapping,
         __inc_zone_page_state(newpage, NR_FILE_PAGES);
 
         write_unlock_irq(&mapping->tree_lock);
+        if (!PageSwapCache(newpage)) {
+                mem_cgroup_uncharge_page(page);
+                mem_cgroup_getref(newpage);
+        }
 
         return 0;
 }
@@ -611,7 +615,6 @@ static int move_to_new_page(struct page *newpage, struct page *page)
                 rc = fallback_migrate_page(mapping, newpage, page);
 
         if (!rc) {
-                mem_cgroup_page_migration(page, newpage);
                 remove_migration_ptes(page, newpage);
         } else
                 newpage->mapping = NULL;
@@ -641,6 +644,14 @@ static int unmap_and_move(new_page_t get_new_page, unsigned long private,
                 /* page was freed from under us. So we are done. */
                 goto move_newpage;
 
+        charge = mem_cgroup_prepare_migration(page, newpage);
+        if (charge == -ENOMEM) {
+                rc = -ENOMEM;
+                goto move_newpage;
+        }
+        /* prepare cgroup just returns 0 or -ENOMEM */
+        BUG_ON(charge);
+
         rc = -EAGAIN;
         if (TestSetPageLocked(page)) {
                 if (!force)
@@ -692,19 +703,14 @@ static int unmap_and_move(new_page_t get_new_page, unsigned long private,
                 goto rcu_unlock;
         }
 
-        charge = mem_cgroup_prepare_migration(page);
         /* Establish migration ptes or remove ptes */
         try_to_unmap(page, 1);
 
         if (!page_mapped(page))
                 rc = move_to_new_page(newpage, page);
 
-        if (rc) {
+        if (rc)
                 remove_migration_ptes(page, page);
-                if (charge)
-                        mem_cgroup_end_migration(page);
-        } else if (charge)
-                mem_cgroup_end_migration(newpage);
 rcu_unlock:
         if (rcu_locked)
                 rcu_read_unlock();
@@ -725,6 +731,8 @@ unlock:
         }
 
 move_newpage:
+        if (!charge)
+                mem_cgroup_end_migration(newpage);
         /*
          * Move the new page to the LRU. If migration was not successful
          * then this will free the page.