2 files changed, 53 insertions, 3 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 3270ce7375db..128f45c16fa6 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -492,6 +492,49 @@ void mem_cgroup_uncharge(struct page_cgroup *pc)
                 }
         }
 }
+/*
+ * Returns non-zero if a page (under migration) has valid page_cgroup member.
+ * Refcnt of page_cgroup is incremented.
+ */
+
+int mem_cgroup_prepare_migration(struct page *page)
+{
+        struct page_cgroup *pc;
+        int ret = 0;
+        lock_page_cgroup(page);
+        pc = page_get_page_cgroup(page);
+        if (pc && atomic_inc_not_zero(&pc->ref_cnt))
+                ret = 1;
+        unlock_page_cgroup(page);
+        return ret;
+}
+
+void mem_cgroup_end_migration(struct page *page)
+{
+        struct page_cgroup *pc = page_get_page_cgroup(page);
+        mem_cgroup_uncharge(pc);
+}
+/*
+ * We know both *page* and *newpage* are now not-on-LRU and Pg_locked.
+ * And no race with uncharge() routines because page_cgroup for *page*
+ * has extra one reference by mem_cgroup_prepare_migration.
+ */
+
+void mem_cgroup_page_migration(struct page *page, struct page *newpage)
+{
+        struct page_cgroup *pc;
+retry:
+        pc = page_get_page_cgroup(page);
+        if (!pc)
+                return;
+        if (clear_page_cgroup(page, pc) != pc)
+                goto retry;
+        pc->page = newpage;
+        lock_page_cgroup(newpage);
+        page_assign_page_cgroup(newpage, pc);
+        unlock_page_cgroup(newpage);
+        return;
+}
 
 int mem_cgroup_write_strategy(char *buf, unsigned long long *tmp)
 {
diff --git a/mm/migrate.c b/mm/migrate.c
index 763794144697..a73504ff5ab9 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -593,9 +593,10 @@ static int move_to_new_page(struct page *newpage, struct page *page)
         else
                 rc = fallback_migrate_page(mapping, newpage, page);
 
-        if (!rc)
+        if (!rc) {
+                mem_cgroup_page_migration(page, newpage);
                 remove_migration_ptes(page, newpage);
-        else
+        } else
                 newpage->mapping = NULL;
 
         unlock_page(newpage);
@@ -614,6 +615,7 @@ static int unmap_and_move(new_page_t get_new_page, unsigned long private,
         int *result = NULL;
         struct page *newpage = get_new_page(page, private, &result);
         int rcu_locked = 0;
+        int charge = 0;
 
         if (!newpage)
                 return -ENOMEM;
@@ -673,14 +675,19 @@ 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();