#include <linux/miscdevice.h>
 #include <linux/log2.h>
 #include <linux/kthread.h>
+#include <linux/reboot.h>
 #include "ubi.h"
 
 /* Maximum length of the 'mtd=' parameter */
        return 0;
 }
 
+/**
+ * ubi_reboot_notifier - halt UBI transactions immediately prior to a reboot.
+ * @n: reboot notifier object
+ * @state: SYS_RESTART, SYS_HALT, or SYS_POWER_OFF
+ * @cmd: pointer to command string for RESTART2
+ *
+ * This function stops the UBI background thread so that the flash device
+ * remains quiescent when Linux restarts the system. Any queued work will be
+ * discarded, but this function will block until do_work() finishes if an
+ * operation is already in progress.
+ *
+ * This function solves a real-life problem observed on NOR flashes when an
+ * PEB erase operation starts, then the system is rebooted before the erase is
+ * finishes, and the boot loader gets confused and dies. So we prefer to finish
+ * the ongoing operation before rebooting.
+ */
+static int ubi_reboot_notifier(struct notifier_block *n, unsigned long state,
+                              void *cmd)
+{
+       struct ubi_device *ubi;
+
+       ubi = container_of(n, struct ubi_device, reboot_notifier);
+       if (ubi->bgt_thread)
+               kthread_stop(ubi->bgt_thread);
+       ubi_sync(ubi->ubi_num);
+       return NOTIFY_DONE;
+}
+
 /**
  * ubi_attach_mtd_dev - attach an MTD device.
  * @mtd: MTD device description object
        wake_up_process(ubi->bgt_thread);
        spin_unlock(&ubi->wl_lock);
 
+       /* Flash device priority is 0 - UBI needs to shut down first */
+       ubi->reboot_notifier.priority = 1;
+       ubi->reboot_notifier.notifier_call = ubi_reboot_notifier;
+       register_reboot_notifier(&ubi->reboot_notifier);
+
        ubi_devices[ubi_num] = ubi;
        ubi_notify_all(ubi, UBI_VOLUME_ADDED, NULL);
        return ubi_num;
         * Before freeing anything, we have to stop the background thread to
         * prevent it from doing anything on this device while we are freeing.
         */
+       unregister_reboot_notifier(&ubi->reboot_notifier);
        if (ubi->bgt_thread)
                kthread_stop(ubi->bgt_thread);
 
 
 #include <linux/device.h>
 #include <linux/string.h>
 #include <linux/vmalloc.h>
+#include <linux/notifier.h>
 #include <linux/mtd/mtd.h>
 #include <linux/mtd/ubi.h>
-#include <linux/notifier.h>
 
 #include "ubi-media.h"
 #include "scan.h"
  * @bgt_thread: background thread description object
  * @thread_enabled: if the background thread is enabled
  * @bgt_name: background thread name
+ * @reboot_notifier: notifier to terminate background thread before rebooting
  *
  * @flash_size: underlying MTD device size (in bytes)
  * @peb_count: count of physical eraseblocks on the MTD device
        struct task_struct *bgt_thread;
        int thread_enabled;
        char bgt_name[sizeof(UBI_BGT_NAME_PATTERN)+2];
+       struct notifier_block reboot_notifier;
 
        /* I/O sub-system's stuff */
        long long flash_size;