3 files changed, 106 insertions, 60 deletions

diff --git a/include/linux/hugetlb.h b/include/linux/hugetlb.h
index feaf0c7fb7d8..c8958f563116 100644
--- a/include/linux/hugetlb.h
+++ b/include/linux/hugetlb.h
@@ -31,6 +31,7 @@ struct hugepage_subpool *hugepage_new_subpool(long nr_blocks);
 void hugepage_put_subpool(struct hugepage_subpool *spool);
 
 int PageHuge(struct page *page);
+int PageHeadHuge(struct page *page_head);
 
 void reset_vma_resv_huge_pages(struct vm_area_struct *vma);
 int hugetlb_sysctl_handler(struct ctl_table *, int, void __user *, size_t *, loff_t *);
@@ -98,6 +99,11 @@ static inline int PageHuge(struct page *page)
         return 0;
 }
 
+static inline int PageHeadHuge(struct page *page_head)
+{
+        return 0;
+}
+
 static inline void reset_vma_resv_huge_pages(struct vm_area_struct *vma)
 {
 }
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 7c5eb85ec645..40ad2c6e0ca9 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -690,6 +690,23 @@ int PageHuge(struct page *page)
 }
 EXPORT_SYMBOL_GPL(PageHuge);
 
+/*
+ * PageHeadHuge() only returns true for hugetlbfs head page, but not for
+ * normal or transparent huge pages.
+ */
+int PageHeadHuge(struct page *page_head)
+{
+        compound_page_dtor *dtor;
+
+        if (!PageHead(page_head))
+                return 0;
+
+        dtor = get_compound_page_dtor(page_head);
+
+        return dtor == free_huge_page;
+}
+EXPORT_SYMBOL_GPL(PageHeadHuge);
+
 pgoff_t __basepage_index(struct page *page)
 {
         struct page *page_head = compound_head(page);
diff --git a/mm/swap.c b/mm/swap.c
index 9f2225f2b5b0..ea58dbde788e 100644
--- a/mm/swap.c
+++ b/mm/swap.c
@@ -79,19 +79,6 @@ static void __put_compound_page(struct page *page)
 
 static void put_compound_page(struct page *page)
 {
-        /*
-         * hugetlbfs pages cannot be split from under us.  If this is a
-         * hugetlbfs page, check refcount on head page and release the page if
-         * the refcount becomes zero.
-         */
-        if (PageHuge(page)) {
-                page = compound_head(page);
-                if (put_page_testzero(page))
-                        __put_compound_page(page);
-
-                return;
-        }
-
         if (unlikely(PageTail(page))) {
                 /* __split_huge_page_refcount can run under us */
                 struct page *page_head = compound_trans_head(page);
@@ -108,14 +95,31 @@ static void put_compound_page(struct page *page)
                          * still hot on arches that do not support
                          * this_cpu_cmpxchg_double().
                          */
-                        if (PageSlab(page_head)) {
-                                if (PageTail(page)) {
+                        if (PageSlab(page_head) || PageHeadHuge(page_head)) {
+                                if (likely(PageTail(page))) {
+                                        /*
+                                         * __split_huge_page_refcount
+                                         * cannot race here.
+                                         */
+                                        VM_BUG_ON(!PageHead(page_head));
+                                        atomic_dec(&page->_mapcount);
                                         if (put_page_testzero(page_head))
                                                 VM_BUG_ON(1);
-
-                                        atomic_dec(&page->_mapcount);
-                                        goto skip_lock_tail;
+                                        if (put_page_testzero(page_head))
+                                                __put_compound_page(page_head);
+                                        return;
                                 } else
+                                        /*
+                                         * __split_huge_page_refcount
+                                         * run before us, "page" was a
+                                         * THP tail. The split
+                                         * page_head has been freed
+                                         * and reallocated as slab or
+                                         * hugetlbfs page of smaller
+                                         * order (only possible if
+                                         * reallocated as slab on
+                                         * x86).
+                                         */
                                         goto skip_lock;
                         }
                         /*
@@ -129,8 +133,27 @@ static void put_compound_page(struct page *page)
                                 /* __split_huge_page_refcount run before us */
                                 compound_unlock_irqrestore(page_head, flags);
 skip_lock:
-                                if (put_page_testzero(page_head))
-                                        __put_single_page(page_head);
+                                if (put_page_testzero(page_head)) {
+                                        /*
+                                         * The head page may have been
+                                         * freed and reallocated as a
+                                         * compound page of smaller
+                                         * order and then freed again.
+                                         * All we know is that it
+                                         * cannot have become: a THP
+                                         * page, a compound page of
+                                         * higher order, a tail page.
+                                         * That is because we still
+                                         * hold the refcount of the
+                                         * split THP tail and
+                                         * page_head was the THP head
+                                         * before the split.
+                                         */
+                                        if (PageHead(page_head))
+                                                __put_compound_page(page_head);
+                                        else
+                                                __put_single_page(page_head);
+                                }
 out_put_single:
                                 if (put_page_testzero(page))
                                         __put_single_page(page);
@@ -152,7 +175,6 @@ out_put_single:
                         VM_BUG_ON(atomic_read(&page->_count) != 0);
                         compound_unlock_irqrestore(page_head, flags);
 
-skip_lock_tail:
                         if (put_page_testzero(page_head)) {
                                 if (PageHead(page_head))
                                         __put_compound_page(page_head);
@@ -195,51 +217,52 @@ bool __get_page_tail(struct page *page)
          * proper PT lock that already serializes against
          * split_huge_page().
          */
+        unsigned long flags;
         bool got = false;
-        struct page *page_head;
-
-        /*
-         * If this is a hugetlbfs page it cannot be split under us.  Simply
-         * increment refcount for the head page.
-         */
-        if (PageHuge(page)) {
-                page_head = compound_head(page);
-                atomic_inc(&page_head->_count);
-                got = true;
-        } else {
-                unsigned long flags;
+        struct page *page_head = compound_trans_head(page);
 
-                page_head = compound_trans_head(page);
-                if (likely(page != page_head &&
-                                        get_page_unless_zero(page_head))) {
-
-                        /* Ref to put_compound_page() comment. */
-                        if (PageSlab(page_head)) {
-                                if (likely(PageTail(page))) {
-                                        __get_page_tail_foll(page, false);
-                                        return true;
-                                } else {
-                                        put_page(page_head);
-                                        return false;
-                                }
-                        }
-
-                        /*
-                         * page_head wasn't a dangling pointer but it
-                         * may not be a head page anymore by the time
-                         * we obtain the lock. That is ok as long as it
-                         * can't be freed from under us.
-                         */
-                        flags = compound_lock_irqsave(page_head);
-                        /* here __split_huge_page_refcount won't run anymore */
+        if (likely(page != page_head && get_page_unless_zero(page_head))) {
+                /* Ref to put_compound_page() comment. */
+                if (PageSlab(page_head) || PageHeadHuge(page_head)) {
                         if (likely(PageTail(page))) {
+                                /*
+                                 * This is a hugetlbfs page or a slab
+                                 * page. __split_huge_page_refcount
+                                 * cannot race here.
+                                 */
+                                VM_BUG_ON(!PageHead(page_head));
                                 __get_page_tail_foll(page, false);
-                                got = true;
-                        }
-                        compound_unlock_irqrestore(page_head, flags);
-                        if (unlikely(!got))
+                                return true;
+                        } else {
+                                /*
+                                 * __split_huge_page_refcount run
+                                 * before us, "page" was a THP
+                                 * tail. The split page_head has been
+                                 * freed and reallocated as slab or
+                                 * hugetlbfs page of smaller order
+                                 * (only possible if reallocated as
+                                 * slab on x86).
+                                 */
                                 put_page(page_head);
+                                return false;
+                        }
+                }
+
+                /*
+                 * page_head wasn't a dangling pointer but it
+                 * may not be a head page anymore by the time
+                 * we obtain the lock. That is ok as long as it
+                 * can't be freed from under us.
+                 */
+                flags = compound_lock_irqsave(page_head);
+                /* here __split_huge_page_refcount won't run anymore */
+                if (likely(PageTail(page))) {
+                        __get_page_tail_foll(page, false);
+                        got = true;
                 }
+                compound_unlock_irqrestore(page_head, flags);
+                if (unlikely(!got))
+                        put_page(page_head);
         }
         return got;
 }