Linux 6.9-rc1
[linux-2.6-microblaze.git] / mm / memory_hotplug.c
index 1213d0c..a444e2d 100644 (file)
@@ -13,7 +13,6 @@
 #include <linux/pagemap.h>
 #include <linux/compiler.h>
 #include <linux/export.h>
-#include <linux/pagevec.h>
 #include <linux/writeback.h>
 #include <linux/slab.h>
 #include <linux/sysctl.h>
 #include "internal.h"
 #include "shuffle.h"
 
-#ifdef CONFIG_MHP_MEMMAP_ON_MEMORY
-static int memmap_on_memory_set(const char *val, const struct kernel_param *kp)
+enum {
+       MEMMAP_ON_MEMORY_DISABLE = 0,
+       MEMMAP_ON_MEMORY_ENABLE,
+       MEMMAP_ON_MEMORY_FORCE,
+};
+
+static int memmap_mode __read_mostly = MEMMAP_ON_MEMORY_DISABLE;
+
+static inline unsigned long memory_block_memmap_size(void)
 {
-       if (hugetlb_optimize_vmemmap_enabled())
-               return 0;
-       return param_set_bool(val, kp);
+       return PHYS_PFN(memory_block_size_bytes()) * sizeof(struct page);
 }
 
-static const struct kernel_param_ops memmap_on_memory_ops = {
-       .flags  = KERNEL_PARAM_OPS_FL_NOARG,
-       .set    = memmap_on_memory_set,
-       .get    = param_get_bool,
-};
+static inline unsigned long memory_block_memmap_on_memory_pages(void)
+{
+       unsigned long nr_pages = PFN_UP(memory_block_memmap_size());
+
+       /*
+        * In "forced" memmap_on_memory mode, we add extra pages to align the
+        * vmemmap size to cover full pageblocks. That way, we can add memory
+        * even if the vmemmap size is not properly aligned, however, we might waste
+        * memory.
+        */
+       if (memmap_mode == MEMMAP_ON_MEMORY_FORCE)
+               return pageblock_align(nr_pages);
+       return nr_pages;
+}
 
+#ifdef CONFIG_MHP_MEMMAP_ON_MEMORY
 /*
  * memory_hotplug.memmap_on_memory parameter
  */
-static bool memmap_on_memory __ro_after_init;
-module_param_cb(memmap_on_memory, &memmap_on_memory_ops, &memmap_on_memory, 0444);
-MODULE_PARM_DESC(memmap_on_memory, "Enable memmap on memory for memory hotplug");
+static int set_memmap_mode(const char *val, const struct kernel_param *kp)
+{
+       int ret, mode;
+       bool enabled;
+
+       if (sysfs_streq(val, "force") ||  sysfs_streq(val, "FORCE")) {
+               mode = MEMMAP_ON_MEMORY_FORCE;
+       } else {
+               ret = kstrtobool(val, &enabled);
+               if (ret < 0)
+                       return ret;
+               if (enabled)
+                       mode = MEMMAP_ON_MEMORY_ENABLE;
+               else
+                       mode = MEMMAP_ON_MEMORY_DISABLE;
+       }
+       *((int *)kp->arg) = mode;
+       if (mode == MEMMAP_ON_MEMORY_FORCE) {
+               unsigned long memmap_pages = memory_block_memmap_on_memory_pages();
 
-bool mhp_memmap_on_memory(void)
+               pr_info_once("Memory hotplug will waste %ld pages in each memory block\n",
+                            memmap_pages - PFN_UP(memory_block_memmap_size()));
+       }
+       return 0;
+}
+
+static int get_memmap_mode(char *buffer, const struct kernel_param *kp)
 {
-       return memmap_on_memory;
+       int mode = *((int *)kp->arg);
+
+       if (mode == MEMMAP_ON_MEMORY_FORCE)
+               return sprintf(buffer, "force\n");
+       return sprintf(buffer, "%c\n", mode ? 'Y' : 'N');
+}
+
+static const struct kernel_param_ops memmap_mode_ops = {
+       .set = set_memmap_mode,
+       .get = get_memmap_mode,
+};
+module_param_cb(memmap_on_memory, &memmap_mode_ops, &memmap_mode, 0444);
+MODULE_PARM_DESC(memmap_on_memory, "Enable memmap on memory for memory hotplug\n"
+                "With value \"force\" it could result in memory wastage due "
+                "to memmap size limitations (Y/N/force)");
+
+static inline bool mhp_memmap_on_memory(void)
+{
+       return memmap_mode != MEMMAP_ON_MEMORY_DISABLE;
+}
+#else
+static inline bool mhp_memmap_on_memory(void)
+{
+       return false;
 }
 #endif
 
@@ -237,8 +296,7 @@ static void release_memory_resource(struct resource *res)
        kfree(res);
 }
 
-static int check_pfn_span(unsigned long pfn, unsigned long nr_pages,
-               const char *reason)
+static int check_pfn_span(unsigned long pfn, unsigned long nr_pages)
 {
        /*
         * Disallow all operations smaller than a sub-section and only
@@ -255,12 +313,8 @@ static int check_pfn_span(unsigned long pfn, unsigned long nr_pages,
                min_align = PAGES_PER_SUBSECTION;
        else
                min_align = PAGES_PER_SECTION;
-       if (!IS_ALIGNED(pfn, min_align)
-                       || !IS_ALIGNED(nr_pages, min_align)) {
-               WARN(1, "Misaligned __%s_pages start: %#lx end: #%lx\n",
-                               reason, pfn, pfn + nr_pages - 1);
+       if (!IS_ALIGNED(pfn | nr_pages, min_align))
                return -EINVAL;
-       }
        return 0;
 }
 
@@ -337,9 +391,10 @@ int __ref __add_pages(int nid, unsigned long pfn, unsigned long nr_pages,
                altmap->alloc = 0;
        }
 
-       err = check_pfn_span(pfn, nr_pages, "add");
-       if (err)
-               return err;
+       if (check_pfn_span(pfn, nr_pages)) {
+               WARN(1, "Misaligned %s start: %#lx end: %#lx\n", __func__, pfn, pfn + nr_pages - 1);
+               return -EINVAL;
+       }
 
        for (; pfn < end_pfn; pfn += cur_nr_pages) {
                /* Select all remaining pages up to the next section boundary */
@@ -504,18 +559,6 @@ void __ref remove_pfn_range_from_zone(struct zone *zone,
        set_zone_contiguous(zone);
 }
 
-static void __remove_section(unsigned long pfn, unsigned long nr_pages,
-                            unsigned long map_offset,
-                            struct vmem_altmap *altmap)
-{
-       struct mem_section *ms = __pfn_to_section(pfn);
-
-       if (WARN_ON_ONCE(!valid_section(ms)))
-               return;
-
-       sparse_remove_section(ms, pfn, nr_pages, map_offset, altmap);
-}
-
 /**
  * __remove_pages() - remove sections of pages
  * @pfn: starting pageframe (must be aligned to start of a section)
@@ -532,20 +575,18 @@ void __remove_pages(unsigned long pfn, unsigned long nr_pages,
 {
        const unsigned long end_pfn = pfn + nr_pages;
        unsigned long cur_nr_pages;
-       unsigned long map_offset = 0;
 
-       map_offset = vmem_altmap_offset(altmap);
-
-       if (check_pfn_span(pfn, nr_pages, "remove"))
+       if (check_pfn_span(pfn, nr_pages)) {
+               WARN(1, "Misaligned %s start: %#lx end: %#lx\n", __func__, pfn, pfn + nr_pages - 1);
                return;
+       }
 
        for (; pfn < end_pfn; pfn += cur_nr_pages) {
                cond_resched();
                /* Select all remaining pages up to the next section boundary */
                cur_nr_pages = min(end_pfn - pfn,
                                   SECTION_ALIGN_UP(pfn + 1) - pfn);
-               __remove_section(pfn, cur_nr_pages, map_offset, altmap);
-               map_offset = 0;
+               sparse_remove_section(pfn, cur_nr_pages, altmap);
        }
 }
 
@@ -606,7 +647,7 @@ static void online_pages_range(unsigned long start_pfn, unsigned long nr_pages)
        unsigned long pfn;
 
        /*
-        * Online the pages in MAX_ORDER - 1 aligned chunks. The callback might
+        * Online the pages in MAX_PAGE_ORDER aligned chunks. The callback might
         * decide to not expose all pages to the buddy (e.g., expose them
         * later). We account all pages as being online and belonging to this
         * zone ("present").
@@ -615,7 +656,19 @@ static void online_pages_range(unsigned long start_pfn, unsigned long nr_pages)
         * this and the first chunk to online will be pageblock_nr_pages.
         */
        for (pfn = start_pfn; pfn < end_pfn;) {
-               int order = min(MAX_ORDER - 1UL, __ffs(pfn));
+               int order;
+
+               /*
+                * Free to online pages in the largest chunks alignment allows.
+                *
+                * __ffs() behaviour is undefined for 0. start == 0 is
+                * MAX_PAGE_ORDER-aligned, Set order to MAX_PAGE_ORDER for
+                * the case.
+                */
+               if (pfn)
+                       order = min_t(int, MAX_PAGE_ORDER, __ffs(pfn));
+               else
+                       order = MAX_PAGE_ORDER;
 
                (*online_page_callback)(pfn_to_page(pfn), order);
                pfn += (1UL << order);
@@ -672,12 +725,18 @@ static void __meminit resize_pgdat_range(struct pglist_data *pgdat, unsigned lon
 
 }
 
+#ifdef CONFIG_ZONE_DEVICE
 static void section_taint_zone_device(unsigned long pfn)
 {
        struct mem_section *ms = __pfn_to_section(pfn);
 
        ms->section_mem_map |= SECTION_TAINT_ZONE_DEVICE;
 }
+#else
+static inline void section_taint_zone_device(unsigned long pfn)
+{
+}
+#endif
 
 /*
  * Associate the pfn range with the given zone, initializing the memmaps
@@ -936,7 +995,7 @@ static struct zone *auto_movable_zone_for_pfn(int nid,
                        if (!page)
                                continue;
                        /* If anything is !MOVABLE online the rest !MOVABLE. */
-                       if (page_zonenum(page) != ZONE_MOVABLE)
+                       if (!is_zone_movable_page(page))
                                goto kernel_zone;
                        online_pages += PAGES_PER_SECTION;
                }
@@ -1028,17 +1087,29 @@ void adjust_present_page_count(struct page *page, struct memory_group *group,
 }
 
 int mhp_init_memmap_on_memory(unsigned long pfn, unsigned long nr_pages,
-                             struct zone *zone)
+                             struct zone *zone, bool mhp_off_inaccessible)
 {
        unsigned long end_pfn = pfn + nr_pages;
-       int ret;
+       int ret, i;
 
        ret = kasan_add_zero_shadow(__va(PFN_PHYS(pfn)), PFN_PHYS(nr_pages));
        if (ret)
                return ret;
 
+       /*
+        * Memory block is accessible at this stage and hence poison the struct
+        * pages now.  If the memory block is accessible during memory hotplug
+        * addition phase, then page poisining is already performed in
+        * sparse_add_section().
+        */
+       if (mhp_off_inaccessible)
+               page_init_poison(pfn_to_page(pfn), sizeof(struct page) * nr_pages);
+
        move_pfn_range_to_zone(zone, pfn, nr_pages, NULL, MIGRATE_UNMOVABLE);
 
+       for (i = 0; i < nr_pages; i++)
+               SetPageVmemmapSelfHosted(pfn_to_page(pfn + i));
+
        /*
         * It might be that the vmemmap_pages fully span sections. If that is
         * the case, mark those sections online here as otherwise they will be
@@ -1070,6 +1141,9 @@ void mhp_deinit_memmap_on_memory(unsigned long pfn, unsigned long nr_pages)
        kasan_remove_zero_shadow(__va(PFN_PHYS(pfn)), PFN_PHYS(nr_pages));
 }
 
+/*
+ * Must be called with mem_hotplug_lock in write mode.
+ */
 int __ref online_pages(unsigned long pfn, unsigned long nr_pages,
                       struct zone *zone, struct memory_group *group)
 {
@@ -1086,12 +1160,10 @@ int __ref online_pages(unsigned long pfn, unsigned long nr_pages,
         * of the physical memory space for vmemmaps. That space is pageblock
         * aligned.
         */
-       if (WARN_ON_ONCE(!nr_pages ||
-                        !IS_ALIGNED(pfn, pageblock_nr_pages) ||
+       if (WARN_ON_ONCE(!nr_pages || !pageblock_aligned(pfn) ||
                         !IS_ALIGNED(pfn + nr_pages, PAGES_PER_SECTION)))
                return -EINVAL;
 
-       mem_hotplug_begin();
 
        /* associate pfn range with the zone */
        move_pfn_range_to_zone(zone, pfn, nr_pages, NULL, MIGRATE_ISOLATE);
@@ -1150,7 +1222,6 @@ int __ref online_pages(unsigned long pfn, unsigned long nr_pages,
        writeback_set_ratelimit();
 
        memory_notify(MEM_ONLINE, &arg);
-       mem_hotplug_done();
        return 0;
 
 failed_addition:
@@ -1159,20 +1230,9 @@ failed_addition:
                 (((unsigned long long) pfn + nr_pages) << PAGE_SHIFT) - 1);
        memory_notify(MEM_CANCEL_ONLINE, &arg);
        remove_pfn_range_from_zone(zone, pfn, nr_pages);
-       mem_hotplug_done();
        return ret;
 }
 
-static void reset_node_present_pages(pg_data_t *pgdat)
-{
-       struct zone *z;
-
-       for (z = pgdat->node_zones; z < pgdat->node_zones + MAX_NR_ZONES; z++)
-               z->present_pages = 0;
-
-       pgdat->node_present_pages = 0;
-}
-
 /* we are OK calling __meminit stuff here - we have CONFIG_MEMORY_HOTPLUG */
 static pg_data_t __ref *hotadd_init_pgdat(int nid)
 {
@@ -1195,15 +1255,6 @@ static pg_data_t __ref *hotadd_init_pgdat(int nid)
         */
        build_all_zonelists(pgdat);
 
-       /*
-        * When memory is hot-added, all the memory is in offline state. So
-        * clear all zones' present_pages because they will be updated in
-        * online_pages() and offline_pages().
-        * TODO: should be in free_area_init_core_hotplug?
-        */
-       reset_node_managed_pages(pgdat);
-       reset_node_present_pages(pgdat);
-
        return pgdat;
 }
 
@@ -1274,27 +1325,32 @@ static int online_memory_block(struct memory_block *mem, void *arg)
        return device_online(&mem->dev);
 }
 
-bool mhp_supports_memmap_on_memory(unsigned long size)
+#ifndef arch_supports_memmap_on_memory
+static inline bool arch_supports_memmap_on_memory(unsigned long vmemmap_size)
+{
+       /*
+        * As default, we want the vmemmap to span a complete PMD such that we
+        * can map the vmemmap using a single PMD if supported by the
+        * architecture.
+        */
+       return IS_ALIGNED(vmemmap_size, PMD_SIZE);
+}
+#endif
+
+bool mhp_supports_memmap_on_memory(void)
 {
-       unsigned long nr_vmemmap_pages = size / PAGE_SIZE;
-       unsigned long vmemmap_size = nr_vmemmap_pages * sizeof(struct page);
-       unsigned long remaining_size = size - vmemmap_size;
+       unsigned long vmemmap_size = memory_block_memmap_size();
+       unsigned long memmap_pages = memory_block_memmap_on_memory_pages();
 
        /*
         * Besides having arch support and the feature enabled at runtime, we
         * need a few more assumptions to hold true:
         *
-        * a) We span a single memory block: memory onlining/offlinin;g happens
-        *    in memory block granularity. We don't want the vmemmap of online
-        *    memory blocks to reside on offline memory blocks. In the future,
-        *    we might want to support variable-sized memory blocks to make the
-        *    feature more versatile.
-        *
-        * b) The vmemmap pages span complete PMDs: We don't want vmemmap code
+        * a) The vmemmap pages span complete PMDs: We don't want vmemmap code
         *    to populate memory from the altmap for unrelated parts (i.e.,
         *    other memory blocks)
         *
-        * c) The vmemmap pages (and thereby the pages that will be exposed to
+        * b) The vmemmap pages (and thereby the pages that will be exposed to
         *    the buddy) have to cover full pageblocks: memory onlining/offlining
         *    code requires applicable ranges to be page-aligned, for example, to
         *    set the migratetypes properly.
@@ -1306,10 +1362,110 @@ bool mhp_supports_memmap_on_memory(unsigned long size)
         *       altmap as an alternative source of memory, and we do not exactly
         *       populate a single PMD.
         */
-       return mhp_memmap_on_memory() &&
-              size == memory_block_size_bytes() &&
-              IS_ALIGNED(vmemmap_size, PMD_SIZE) &&
-              IS_ALIGNED(remaining_size, (pageblock_nr_pages << PAGE_SHIFT));
+       if (!mhp_memmap_on_memory())
+               return false;
+
+       /*
+        * Make sure the vmemmap allocation is fully contained
+        * so that we always allocate vmemmap memory from altmap area.
+        */
+       if (!IS_ALIGNED(vmemmap_size, PAGE_SIZE))
+               return false;
+
+       /*
+        * start pfn should be pageblock_nr_pages aligned for correctly
+        * setting migrate types
+        */
+       if (!pageblock_aligned(memmap_pages))
+               return false;
+
+       if (memmap_pages == PHYS_PFN(memory_block_size_bytes()))
+               /* No effective hotplugged memory doesn't make sense. */
+               return false;
+
+       return arch_supports_memmap_on_memory(vmemmap_size);
+}
+EXPORT_SYMBOL_GPL(mhp_supports_memmap_on_memory);
+
+static void __ref remove_memory_blocks_and_altmaps(u64 start, u64 size)
+{
+       unsigned long memblock_size = memory_block_size_bytes();
+       u64 cur_start;
+
+       /*
+        * For memmap_on_memory, the altmaps were added on a per-memblock
+        * basis; we have to process each individual memory block.
+        */
+       for (cur_start = start; cur_start < start + size;
+            cur_start += memblock_size) {
+               struct vmem_altmap *altmap = NULL;
+               struct memory_block *mem;
+
+               mem = find_memory_block(pfn_to_section_nr(PFN_DOWN(cur_start)));
+               if (WARN_ON_ONCE(!mem))
+                       continue;
+
+               altmap = mem->altmap;
+               mem->altmap = NULL;
+
+               remove_memory_block_devices(cur_start, memblock_size);
+
+               arch_remove_memory(cur_start, memblock_size, altmap);
+
+               /* Verify that all vmemmap pages have actually been freed. */
+               WARN(altmap->alloc, "Altmap not fully unmapped");
+               kfree(altmap);
+       }
+}
+
+static int create_altmaps_and_memory_blocks(int nid, struct memory_group *group,
+                                           u64 start, u64 size, mhp_t mhp_flags)
+{
+       unsigned long memblock_size = memory_block_size_bytes();
+       u64 cur_start;
+       int ret;
+
+       for (cur_start = start; cur_start < start + size;
+            cur_start += memblock_size) {
+               struct mhp_params params = { .pgprot =
+                                                    pgprot_mhp(PAGE_KERNEL) };
+               struct vmem_altmap mhp_altmap = {
+                       .base_pfn = PHYS_PFN(cur_start),
+                       .end_pfn = PHYS_PFN(cur_start + memblock_size - 1),
+               };
+
+               mhp_altmap.free = memory_block_memmap_on_memory_pages();
+               if (mhp_flags & MHP_OFFLINE_INACCESSIBLE)
+                       mhp_altmap.inaccessible = true;
+               params.altmap = kmemdup(&mhp_altmap, sizeof(struct vmem_altmap),
+                                       GFP_KERNEL);
+               if (!params.altmap) {
+                       ret = -ENOMEM;
+                       goto out;
+               }
+
+               /* call arch's memory hotadd */
+               ret = arch_add_memory(nid, cur_start, memblock_size, &params);
+               if (ret < 0) {
+                       kfree(params.altmap);
+                       goto out;
+               }
+
+               /* create memory block devices after memory was added */
+               ret = create_memory_block_devices(cur_start, memblock_size,
+                                                 params.altmap, group);
+               if (ret) {
+                       arch_remove_memory(cur_start, memblock_size, NULL);
+                       kfree(params.altmap);
+                       goto out;
+               }
+       }
+
+       return 0;
+out:
+       if (ret && cur_start != start)
+               remove_memory_blocks_and_altmaps(start, cur_start - start);
+       return ret;
 }
 
 /*
@@ -1322,7 +1478,6 @@ int __ref add_memory_resource(int nid, struct resource *res, mhp_t mhp_flags)
 {
        struct mhp_params params = { .pgprot = pgprot_mhp(PAGE_KERNEL) };
        enum memblock_flags memblock_flags = MEMBLOCK_NONE;
-       struct vmem_altmap mhp_altmap = {};
        struct memory_group *group = NULL;
        u64 start, size;
        bool new_node = false;
@@ -1365,27 +1520,22 @@ int __ref add_memory_resource(int nid, struct resource *res, mhp_t mhp_flags)
        /*
         * Self hosted memmap array
         */
-       if (mhp_flags & MHP_MEMMAP_ON_MEMORY) {
-               if (!mhp_supports_memmap_on_memory(size)) {
-                       ret = -EINVAL;
+       if ((mhp_flags & MHP_MEMMAP_ON_MEMORY) &&
+           mhp_supports_memmap_on_memory()) {
+               ret = create_altmaps_and_memory_blocks(nid, group, start, size, mhp_flags);
+               if (ret)
+                       goto error;
+       } else {
+               ret = arch_add_memory(nid, start, size, &params);
+               if (ret < 0)
                        goto error;
-               }
-               mhp_altmap.free = PHYS_PFN(size);
-               mhp_altmap.base_pfn = PHYS_PFN(start);
-               params.altmap = &mhp_altmap;
-       }
-
-       /* call arch's memory hotadd */
-       ret = arch_add_memory(nid, start, size, &params);
-       if (ret < 0)
-               goto error;
 
-       /* create memory block devices after memory was added */
-       ret = create_memory_block_devices(start, size, mhp_altmap.alloc,
-                                         group);
-       if (ret) {
-               arch_remove_memory(start, size, NULL);
-               goto error;
+               /* create memory block devices after memory was added */
+               ret = create_memory_block_devices(start, size, NULL, group);
+               if (ret) {
+                       arch_remove_memory(start, size, params.altmap);
+                       goto error;
+               }
        }
 
        if (new_node) {
@@ -1613,7 +1763,7 @@ static int scan_movable_pages(unsigned long start, unsigned long end,
                 */
                if (HPageMigratable(head))
                        goto found;
-               skip = compound_nr(head) - (page - head);
+               skip = compound_nr(head) - (pfn - page_to_pfn(head));
                pfn += skip - 1;
        }
        return -ENOENT;
@@ -1622,18 +1772,17 @@ found:
        return 0;
 }
 
-static int
-do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
+static void do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
 {
        unsigned long pfn;
        struct page *page, *head;
-       int ret = 0;
        LIST_HEAD(source);
        static DEFINE_RATELIMIT_STATE(migrate_rs, DEFAULT_RATELIMIT_INTERVAL,
                                      DEFAULT_RATELIMIT_BURST);
 
        for (pfn = start_pfn; pfn < end_pfn; pfn++) {
                struct folio *folio;
+               bool isolated;
 
                if (!pfn_valid(pfn))
                        continue;
@@ -1643,7 +1792,7 @@ do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
 
                if (PageHuge(page)) {
                        pfn = page_to_pfn(head) + compound_nr(head) - 1;
-                       isolate_huge_page(head, &source);
+                       isolate_hugetlb(folio, &source);
                        continue;
                } else if (PageTransHuge(page))
                        pfn = page_to_pfn(head) + thp_nr_pages(page) - 1;
@@ -1670,10 +1819,10 @@ do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
                 * LRU and non-lru movable pages.
                 */
                if (PageLRU(page))
-                       ret = isolate_lru_page(page);
+                       isolated = isolate_lru_page(page);
                else
-                       ret = isolate_movable_page(page, ISOLATE_UNEVICTABLE);
-               if (!ret) { /* Success */
+                       isolated = isolate_movable_page(page, ISOLATE_UNEVICTABLE);
+               if (isolated) {
                        list_add_tail(&page->lru, &source);
                        if (!__PageMovable(page))
                                inc_node_page_state(page, NR_ISOLATED_ANON +
@@ -1693,6 +1842,7 @@ do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
                        .nmask = &nmask,
                        .gfp_mask = GFP_USER | __GFP_MOVABLE | __GFP_RETRY_MAYFAIL,
                };
+               int ret;
 
                /*
                 * We have checked that migration range is on a single zone so
@@ -1721,8 +1871,6 @@ do_migrate_range(unsigned long start_pfn, unsigned long end_pfn)
                        putback_movable_pages(&source);
                }
        }
-
-       return ret;
 }
 
 static int __init cmdline_parse_movable_node(char *p)
@@ -1789,6 +1937,9 @@ static int count_system_ram_pages_cb(unsigned long start_pfn,
        return 0;
 }
 
+/*
+ * Must be called with mem_hotplug_lock in write mode.
+ */
 int __ref offline_pages(unsigned long start_pfn, unsigned long nr_pages,
                        struct zone *zone, struct memory_group *group)
 {
@@ -1807,13 +1958,10 @@ int __ref offline_pages(unsigned long start_pfn, unsigned long nr_pages,
         * of the physical memory space for vmemmaps. That space is pageblock
         * aligned.
         */
-       if (WARN_ON_ONCE(!nr_pages ||
-                        !IS_ALIGNED(start_pfn, pageblock_nr_pages) ||
+       if (WARN_ON_ONCE(!nr_pages || !pageblock_aligned(start_pfn) ||
                         !IS_ALIGNED(start_pfn + nr_pages, PAGES_PER_SECTION)))
                return -EINVAL;
 
-       mem_hotplug_begin();
-
        /*
         * Don't allow to offline memory blocks that contain holes.
         * Consequently, memory blocks with holes can never get onlined
@@ -1873,6 +2021,11 @@ int __ref offline_pages(unsigned long start_pfn, unsigned long nr_pages,
        do {
                pfn = start_pfn;
                do {
+                       /*
+                        * Historically we always checked for any signal and
+                        * can't limit it to fatal signals without eventually
+                        * breaking user space.
+                        */
                        if (signal_pending(current)) {
                                ret = -EINTR;
                                reason = "signal backoff";
@@ -1934,22 +2087,25 @@ int __ref offline_pages(unsigned long start_pfn, unsigned long nr_pages,
        /* reinitialise watermarks and update pcp limits */
        init_per_zone_wmark_min();
 
+       /*
+        * Make sure to mark the node as memory-less before rebuilding the zone
+        * list. Otherwise this node would still appear in the fallback lists.
+        */
+       node_states_clear_node(node, &arg);
        if (!populated_zone(zone)) {
                zone_pcp_reset(zone);
                build_all_zonelists(NULL);
        }
 
-       node_states_clear_node(node, &arg);
        if (arg.status_change_nid >= 0) {
-               kswapd_stop(node);
                kcompactd_stop(node);
+               kswapd_stop(node);
        }
 
        writeback_set_ratelimit();
 
        memory_notify(MEM_OFFLINE, &arg);
        remove_pfn_range_from_zone(zone, start_pfn, nr_pages);
-       mem_hotplug_done();
        return 0;
 
 failed_removal_isolated:
@@ -1964,17 +2120,15 @@ failed_removal:
                 (unsigned long long) start_pfn << PAGE_SHIFT,
                 ((unsigned long long) end_pfn << PAGE_SHIFT) - 1,
                 reason);
-       mem_hotplug_done();
        return ret;
 }
 
 static int check_memblock_offlined_cb(struct memory_block *mem, void *arg)
 {
-       int ret = !is_memblock_offlined(mem);
        int *nid = arg;
 
        *nid = mem->nid;
-       if (unlikely(ret)) {
+       if (unlikely(mem->state != MEM_OFFLINE)) {
                phys_addr_t beginpa, endpa;
 
                beginpa = PFN_PHYS(section_nr_to_pfn(mem->start_section_nr));
@@ -1987,12 +2141,14 @@ static int check_memblock_offlined_cb(struct memory_block *mem, void *arg)
        return 0;
 }
 
-static int get_nr_vmemmap_pages_cb(struct memory_block *mem, void *arg)
+static int count_memory_range_altmaps_cb(struct memory_block *mem, void *arg)
 {
-       /*
-        * If not set, continue with the next block.
-        */
-       return mem->nr_vmemmap_pages;
+       u64 *num_altmaps = (u64 *)arg;
+
+       if (mem->altmap)
+               *num_altmaps += 1;
+
+       return 0;
 }
 
 static int check_cpu_on_node(int nid)
@@ -2065,12 +2221,29 @@ void try_offline_node(int nid)
 }
 EXPORT_SYMBOL(try_offline_node);
 
+static int memory_blocks_have_altmaps(u64 start, u64 size)
+{
+       u64 num_memblocks = size / memory_block_size_bytes();
+       u64 num_altmaps = 0;
+
+       if (!mhp_memmap_on_memory())
+               return 0;
+
+       walk_memory_blocks(start, size, &num_altmaps,
+                          count_memory_range_altmaps_cb);
+
+       if (num_altmaps == 0)
+               return 0;
+
+       if (WARN_ON_ONCE(num_memblocks != num_altmaps))
+               return -EINVAL;
+
+       return 1;
+}
+
 static int __ref try_remove_memory(u64 start, u64 size)
 {
-       struct vmem_altmap mhp_altmap = {};
-       struct vmem_altmap *altmap = NULL;
-       unsigned long nr_vmemmap_pages;
-       int rc = 0, nid = NUMA_NO_NODE;
+       int rc, nid = NUMA_NO_NODE;
 
        BUG_ON(check_hotplug_memory_range(start, size));
 
@@ -2087,43 +2260,27 @@ static int __ref try_remove_memory(u64 start, u64 size)
        if (rc)
                return rc;
 
-       /*
-        * We only support removing memory added with MHP_MEMMAP_ON_MEMORY in
-        * the same granularity it was added - a single memory block.
-        */
-       if (mhp_memmap_on_memory()) {
-               nr_vmemmap_pages = walk_memory_blocks(start, size, NULL,
-                                                     get_nr_vmemmap_pages_cb);
-               if (nr_vmemmap_pages) {
-                       if (size != memory_block_size_bytes()) {
-                               pr_warn("Refuse to remove %#llx - %#llx,"
-                                       "wrong granularity\n",
-                                       start, start + size);
-                               return -EINVAL;
-                       }
-
-                       /*
-                        * Let remove_pmd_table->free_hugepage_table do the
-                        * right thing if we used vmem_altmap when hot-adding
-                        * the range.
-                        */
-                       mhp_altmap.alloc = nr_vmemmap_pages;
-                       altmap = &mhp_altmap;
-               }
-       }
-
        /* remove memmap entry */
        firmware_map_remove(start, start + size, "System RAM");
 
-       /*
-        * Memory block device removal under the device_hotplug_lock is
-        * a barrier against racing online attempts.
-        */
-       remove_memory_block_devices(start, size);
-
        mem_hotplug_begin();
 
-       arch_remove_memory(start, size, altmap);
+       rc = memory_blocks_have_altmaps(start, size);
+       if (rc < 0) {
+               mem_hotplug_done();
+               return rc;
+       } else if (!rc) {
+               /*
+                * Memory block device removal under the device_hotplug_lock is
+                * a barrier against racing online attempts.
+                * No altmaps present, do the removal directly
+                */
+               remove_memory_block_devices(start, size);
+               arch_remove_memory(start, size, NULL);
+       } else {
+               /* all memblocks in the range have altmaps */
+               remove_memory_blocks_and_altmaps(start, size);
+       }
 
        if (IS_ENABLED(CONFIG_ARCH_KEEP_MEMBLOCK)) {
                memblock_phys_free(start, size);