2 files changed, 48 insertions, 14 deletions

diff --git a/block/badblocks.c b/block/badblocks.c
index 37e5c0a2ef69..7be53cb1cc3c 100644
--- a/block/badblocks.c
+++ b/block/badblocks.c
@@ -17,6 +17,7 @@
 
 #include <linux/badblocks.h>
 #include <linux/seqlock.h>
+#include <linux/device.h>
 #include <linux/kernel.h>
 #include <linux/module.h>
 #include <linux/stddef.h>
@@ -522,24 +523,20 @@ ssize_t badblocks_store(struct badblocks *bb, const char *page, size_t len,
 }
 EXPORT_SYMBOL_GPL(badblocks_store);
 
-/**
- * badblocks_init() - initialize the badblocks structure
- * @bb:         the badblocks structure that holds all badblock information
- * @enable:     weather to enable badblocks accounting
- *
- * Return:
- *  0: success
- *  -ve errno: on error
- */
-int badblocks_init(struct badblocks *bb, int enable)
+static int __badblocks_init(struct device *dev, struct badblocks *bb,
+                int enable)
 {
+        bb->dev = dev;
         bb->count = 0;
         if (enable)
                 bb->shift = 0;
         else
                 bb->shift = -1;
-        bb->page = kzalloc(PAGE_SIZE, GFP_KERNEL);
-        if (bb->page == (u64 *)0) {
+        if (dev)
+                bb->page = devm_kzalloc(dev, PAGE_SIZE, GFP_KERNEL);
+        else
+                bb->page = kzalloc(PAGE_SIZE, GFP_KERNEL);
+        if (!bb->page) {
                 bb->shift = -1;
                 return -ENOMEM;
         }
@@ -547,8 +544,30 @@ int badblocks_init(struct badblocks *bb, int enable)
 
         return 0;
 }
+
+/**
+ * badblocks_init() - initialize the badblocks structure
+ * @bb:         the badblocks structure that holds all badblock information
+ * @enable:     weather to enable badblocks accounting
+ *
+ * Return:
+ *  0: success
+ *  -ve errno: on error
+ */
+int badblocks_init(struct badblocks *bb, int enable)
+{
+        return __badblocks_init(NULL, bb, enable);
+}
 EXPORT_SYMBOL_GPL(badblocks_init);
 
+int devm_init_badblocks(struct device *dev, struct badblocks *bb)
+{
+        if (!bb)
+                return -EINVAL;
+        return __badblocks_init(dev, bb, 1);
+}
+EXPORT_SYMBOL_GPL(devm_init_badblocks);
+
 /**
  * badblocks_exit() - free the badblocks structure
  * @bb:         the badblocks structure that holds all badblock information
@@ -557,7 +576,10 @@ void badblocks_exit(struct badblocks *bb)
 {
         if (!bb)
                 return;
-        kfree(bb->page);
+        if (bb->dev)
+                devm_kfree(bb->dev, bb->page);
+        else
+                kfree(bb->page);
         bb->page = NULL;
 }
 EXPORT_SYMBOL_GPL(badblocks_exit);
diff --git a/include/linux/badblocks.h b/include/linux/badblocks.h
index 2d98c026c57f..c3bdf8c59480 100644
--- a/include/linux/badblocks.h
+++ b/include/linux/badblocks.h
@@ -2,6 +2,7 @@
 #define _LINUX_BADBLOCKS_H
 
 #include <linux/seqlock.h>
+#include <linux/device.h>
 #include <linux/kernel.h>
 #include <linux/stddef.h>
 #include <linux/types.h>
@@ -23,6 +24,7 @@
 #define MAX_BADBLOCKS   (PAGE_SIZE/8)
 
 struct badblocks {
+        struct device *dev;     /* set by devm_init_badblocks */
         int count;              /* count of bad blocks */
         int unacked_exist;      /* there probably are unacknowledged
                                  * bad blocks.  This is only cleared
@@ -49,5 +51,15 @@ ssize_t badblocks_store(struct badblocks *bb, const char *page, size_t len,
                         int unack);
 int badblocks_init(struct badblocks *bb, int enable);
 void badblocks_exit(struct badblocks *bb);
-
+struct device;
+int devm_init_badblocks(struct device *dev, struct badblocks *bb);
+static inline void devm_exit_badblocks(struct device *dev, struct badblocks *bb)
+{
+        if (bb->dev != dev) {
+                dev_WARN_ONCE(dev, 1, "%s: badblocks instance not associated\n",
+                                __func__);
+                return;
+        }
+        badblocks_exit(bb);
+}
 #endif