[PATCH 2/5] mm, devm_memremap_pages: handle errors allocating final devres action

[Dan Williams]

The last step before devm_memremap_pages() returns success is to
allocate a release action to tear the entire setup down. However, the
result from devm_add_action() is not checked.

Checking the error also means that we need to handle the fact that the
percpu_ref may not be killed by the time devm_memremap_pages_release()
runs. Add a new state flag for this case.

Cc: <sta...@vger.kernel.org>
Fixes: e8d513483300 ("memremap: change devm_memremap_pages interface...")
Cc: Christoph Hellwig <h...@lst.de>
Cc: "Jérôme Glisse" <jgli...@redhat.com>
Cc: Logan Gunthorpe <log...@deltatee.com>
Signed-off-by: Dan Williams <dan.j.willi...@intel.com>
---
 include/linux/memremap.h |    1 +
 kernel/memremap.c        |    8 ++++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/include/linux/memremap.h b/include/linux/memremap.h
index 7b4899c06f49..44a7ee517513 100644
--- a/include/linux/memremap.h
+++ b/include/linux/memremap.h
@@ -115,6 +115,7 @@ struct dev_pagemap {
        dev_page_free_t page_free;
        struct vmem_altmap altmap;
        bool altmap_valid;
+       bool registered;
        struct resource res;
        struct percpu_ref *ref;
        struct device *dev;
diff --git a/kernel/memremap.c b/kernel/memremap.c
index c614645227a7..30d96be5a965 100644
--- a/kernel/memremap.c
+++ b/kernel/memremap.c
@@ -296,7 +296,7 @@ static void devm_memremap_pages_release(void *data)
        for_each_device_pfn(pfn, pgmap)
                put_page(pfn_to_page(pfn));
 
-       if (percpu_ref_tryget_live(pgmap->ref)) {
+       if (pgmap->registered && percpu_ref_tryget_live(pgmap->ref)) {
                dev_WARN(dev, "%s: page mapping is still live!\n", __func__);
                percpu_ref_put(pgmap->ref);
        }
@@ -418,7 +418,11 @@ void *devm_memremap_pages(struct device *dev, struct 
dev_pagemap *pgmap)
                percpu_ref_get(pgmap->ref);
        }
 
-       devm_add_action(dev, devm_memremap_pages_release, pgmap);
+       error = devm_add_action_or_reset(dev, devm_memremap_pages_release,
+                       pgmap);
+       if (error)
+               return ERR_PTR(error);
+       pgmap->registered = true;
 
        return __va(res->start);