1 files changed, 31 insertions, 7 deletions

diff --git a/mm/migrate.c b/mm/migrate.c
index bb940045fe85..2cabbd5fa5bf 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -1722,6 +1722,7 @@ int migrate_misplaced_transhuge_page(struct mm_struct *mm,
         struct page *new_page = NULL;
         struct mem_cgroup *memcg = NULL;
         int page_lru = page_is_file_cache(page);
+        pmd_t orig_entry;
 
         /*
          * Rate-limit the amount of data that is being migrated to a node.
@@ -1756,7 +1757,8 @@ int migrate_misplaced_transhuge_page(struct mm_struct *mm,
 
         /* Recheck the target PMD */
         ptl = pmd_lock(mm, pmd);
-        if (unlikely(!pmd_same(*pmd, entry))) {
+        if (unlikely(!pmd_same(*pmd, entry) || page_count(page) != 2)) {
+fail_putback:
                 spin_unlock(ptl);
 
                 /* Reverse changes made by migrate_page_copy() */
@@ -1786,16 +1788,34 @@ int migrate_misplaced_transhuge_page(struct mm_struct *mm,
          */
         mem_cgroup_prepare_migration(page, new_page, &memcg);
 
+        orig_entry = *pmd;
         entry = mk_pmd(new_page, vma->vm_page_prot);
-        entry = pmd_mknonnuma(entry);
-        entry = maybe_pmd_mkwrite(pmd_mkdirty(entry), vma);
         entry = pmd_mkhuge(entry);
+        entry = maybe_pmd_mkwrite(pmd_mkdirty(entry), vma);
 
+        /*
+         * Clear the old entry under pagetable lock and establish the new PTE.
+         * Any parallel GUP will either observe the old page blocking on the
+         * page lock, block on the page table lock or observe the new page.
+         * The SetPageUptodate on the new page and page_add_new_anon_rmap
+         * guarantee the copy is visible before the pagetable update.
+         */
+        flush_cache_range(vma, haddr, haddr + HPAGE_PMD_SIZE);
+        page_add_new_anon_rmap(new_page, vma, haddr);
         pmdp_clear_flush(vma, haddr, pmd);
         set_pmd_at(mm, haddr, pmd, entry);
-        page_add_new_anon_rmap(new_page, vma, haddr);
         update_mmu_cache_pmd(vma, address, &entry);
+
+        if (page_count(page) != 2) {
+                set_pmd_at(mm, haddr, pmd, orig_entry);
+                flush_tlb_range(vma, haddr, haddr + HPAGE_PMD_SIZE);
+                update_mmu_cache_pmd(vma, address, &entry);
+                page_remove_rmap(new_page);
+                goto fail_putback;
+        }
+
         page_remove_rmap(page);
+
         /*
          * Finish the charge transaction under the page table lock to
          * prevent split_huge_page() from dividing up the charge
@@ -1820,9 +1840,13 @@ int migrate_misplaced_transhuge_page(struct mm_struct *mm,
 out_fail:
         count_vm_events(PGMIGRATE_FAIL, HPAGE_PMD_NR);
 out_dropref:
-        entry = pmd_mknonnuma(entry);
-        set_pmd_at(mm, haddr, pmd, entry);
-        update_mmu_cache_pmd(vma, address, &entry);
+        ptl = pmd_lock(mm, pmd);
+        if (pmd_same(*pmd, entry)) {
+                entry = pmd_mknonnuma(entry);
+                set_pmd_at(mm, haddr, pmd, entry);
+                update_mmu_cache_pmd(vma, address, &entry);
+        }
+        spin_unlock(ptl);
 
         unlock_page(page);
         put_page(page);