]> git.sur5r.net Git - u-boot/blobdiff - drivers/mtd/ubi/fastmap.c
ubi,ubifs: sync with linux v4.2
[u-boot] / drivers / mtd / ubi / fastmap.c
index a2166e431c22fc463280958db829cf1b9baf0f63..48ba98a6b55b69b31c660a1e5e1b2c0084418df9 100644 (file)
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2012 Linutronix GmbH
+ * Copyright (c) 2014 sigma star gmbh
  * Author: Richard Weinberger <richard@nod.at>
  *
  * SPDX-License-Identifier:    GPL-2.0+
 #include <linux/math64.h>
 #include "ubi.h"
 
+/**
+ * init_seen - allocate memory for used for debugging.
+ * @ubi: UBI device description object
+ */
+static inline int *init_seen(struct ubi_device *ubi)
+{
+       int *ret;
+
+       if (!ubi_dbg_chk_fastmap(ubi))
+               return NULL;
+
+       ret = kcalloc(ubi->peb_count, sizeof(int), GFP_KERNEL);
+       if (!ret)
+               return ERR_PTR(-ENOMEM);
+
+       return ret;
+}
+
+/**
+ * free_seen - free the seen logic integer array.
+ * @seen: integer array of @ubi->peb_count size
+ */
+static inline void free_seen(int *seen)
+{
+       kfree(seen);
+}
+
+/**
+ * set_seen - mark a PEB as seen.
+ * @ubi: UBI device description object
+ * @pnum: The PEB to be makred as seen
+ * @seen: integer array of @ubi->peb_count size
+ */
+static inline void set_seen(struct ubi_device *ubi, int pnum, int *seen)
+{
+       if (!ubi_dbg_chk_fastmap(ubi) || !seen)
+               return;
+
+       seen[pnum] = 1;
+}
+
+/**
+ * self_check_seen - check whether all PEB have been seen by fastmap.
+ * @ubi: UBI device description object
+ * @seen: integer array of @ubi->peb_count size
+ */
+static int self_check_seen(struct ubi_device *ubi, int *seen)
+{
+       int pnum, ret = 0;
+
+       if (!ubi_dbg_chk_fastmap(ubi) || !seen)
+               return 0;
+
+       for (pnum = 0; pnum < ubi->peb_count; pnum++) {
+               if (!seen[pnum] && ubi->lookuptbl[pnum]) {
+                       ubi_err(ubi, "self-check failed for PEB %d, fastmap didn't see it", pnum);
+                       ret = -EINVAL;
+               }
+       }
+
+       return ret;
+}
+
 /**
  * ubi_calc_fm_size - calculates the fastmap size in bytes for an UBI device.
  * @ubi: UBI device description object
@@ -26,12 +90,13 @@ size_t ubi_calc_fm_size(struct ubi_device *ubi)
 {
        size_t size;
 
-       size = sizeof(struct ubi_fm_hdr) + \
-               sizeof(struct ubi_fm_scan_pool) + \
-               sizeof(struct ubi_fm_scan_pool) + \
-               (ubi->peb_count * sizeof(struct ubi_fm_ec)) + \
-               (sizeof(struct ubi_fm_eba) + \
-               (ubi->peb_count * sizeof(__be32))) + \
+       size = sizeof(struct ubi_fm_sb) +
+               sizeof(struct ubi_fm_hdr) +
+               sizeof(struct ubi_fm_scan_pool) +
+               sizeof(struct ubi_fm_scan_pool) +
+               (ubi->peb_count * sizeof(struct ubi_fm_ec)) +
+               (sizeof(struct ubi_fm_eba) +
+               (ubi->peb_count * sizeof(__be32))) +
                sizeof(struct ubi_fm_volhdr) * UBI_MAX_VOLUMES;
        return roundup(size, ubi->leb_size);
 }
@@ -129,22 +194,25 @@ static struct ubi_ainf_volume *add_vol(struct ubi_attach_info *ai, int vol_id,
 
                if (vol_id > av->vol_id)
                        p = &(*p)->rb_left;
-               else if (vol_id > av->vol_id)
+               else if (vol_id < av->vol_id)
                        p = &(*p)->rb_right;
+               else
+                       return ERR_PTR(-EINVAL);
        }
 
        av = kmalloc(sizeof(struct ubi_ainf_volume), GFP_KERNEL);
        if (!av)
                goto out;
 
-       av->highest_lnum = av->leb_count = 0;
+       av->highest_lnum = av->leb_count = av->used_ebs = 0;
        av->vol_id = vol_id;
-       av->used_ebs = used_ebs;
        av->data_pad = data_pad;
        av->last_data_size = last_eb_bytes;
        av->compat = 0;
        av->vol_type = vol_type;
        av->root = RB_ROOT;
+       if (av->vol_type == UBI_STATIC_VOLUME)
+               av->used_ebs = used_ebs;
 
        dbg_bld("found volume (ID %i)", vol_id);
 
@@ -250,7 +318,7 @@ static int update_vol(struct ubi_device *ubi, struct ubi_attach_info *ai,
                        list_add_tail(&victim->u.list, &ai->erase);
 
                        if (av->highest_lnum == be32_to_cpu(new_vh->lnum))
-                               av->last_data_size = \
+                               av->last_data_size =
                                        be32_to_cpu(new_vh->data_size);
 
                        dbg_bld("vol %i: AEB %i's PEB %i is the newer",
@@ -331,7 +399,8 @@ static int process_pool_aeb(struct ubi_device *ubi, struct ubi_attach_info *ai,
        if (found)
                av = tmp_av;
        else {
-               ubi_err("orphaned volume in fastmap pool!");
+               ubi_err(ubi, "orphaned volume in fastmap pool!");
+               kmem_cache_free(ai->aeb_slab_cache, new_aeb);
                return UBI_BAD_FASTMAP;
        }
 
@@ -362,6 +431,7 @@ static void unmap_peb(struct ubi_attach_info *ai, int pnum)
                        aeb = rb_entry(node2, struct ubi_ainf_peb, u.rb);
                        if (aeb->pnum == pnum) {
                                rb_erase(&aeb->u.rb, &av->root);
+                               av->leb_count--;
                                kmem_cache_free(ai->aeb_slab_cache, aeb);
                                return;
                        }
@@ -376,7 +446,6 @@ static void unmap_peb(struct ubi_attach_info *ai, int pnum)
  * @pebs: an array of all PEB numbers in the to be scanned pool
  * @pool_size: size of the pool (number of entries in @pebs)
  * @max_sqnum: pointer to the maximal sequence number
- * @eba_orphans: list of PEBs which need to be scanned
  * @free: list of PEBs which are most likely free (and go into @ai->free)
  *
  * Returns 0 on success, if the pool is unusable UBI_BAD_FASTMAP is returned.
@@ -385,17 +454,17 @@ static void unmap_peb(struct ubi_attach_info *ai, int pnum)
 #ifndef __UBOOT__
 static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                     int *pebs, int pool_size, unsigned long long *max_sqnum,
-                    struct list_head *eba_orphans, struct list_head *freef)
+                    struct list_head *free)
 #else
 static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
-                    __be32 *pebs, int pool_size, unsigned long long *max_sqnum,
-                    struct list_head *eba_orphans, struct list_head *freef)
+                    int *pebs, int pool_size, unsigned long long *max_sqnum,
+                    struct list_head *free)
 #endif
 {
        struct ubi_vid_hdr *vh;
        struct ubi_ec_hdr *ech;
-       struct ubi_ainf_peb *new_aeb, *tmp_aeb;
-       int i, pnum, err, found_orphan, ret = 0;
+       struct ubi_ainf_peb *new_aeb;
+       int i, pnum, err, ret = 0;
 
        ech = kzalloc(ubi->ec_hdr_alsize, GFP_KERNEL);
        if (!ech)
@@ -420,18 +489,18 @@ static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                pnum = be32_to_cpu(pebs[i]);
 
                if (ubi_io_is_bad(ubi, pnum)) {
-                       ubi_err("bad PEB in fastmap pool!");
+                       ubi_err(ubi, "bad PEB in fastmap pool!");
                        ret = UBI_BAD_FASTMAP;
                        goto out;
                }
 
                err = ubi_io_read_ec_hdr(ubi, pnum, ech, 0);
                if (err && err != UBI_IO_BITFLIPS) {
-                       ubi_err("unable to read EC header! PEB:%i err:%i",
+                       ubi_err(ubi, "unable to read EC header! PEB:%i err:%i",
                                pnum, err);
                        ret = err > 0 ? UBI_BAD_FASTMAP : err;
                        goto out;
-               } else if (ret == UBI_IO_BITFLIPS)
+               } else if (err == UBI_IO_BITFLIPS)
                        scrub = 1;
 
                /*
@@ -441,7 +510,7 @@ static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                image_seq = be32_to_cpu(ech->image_seq);
 
                if (image_seq && (image_seq != ubi->image_seq)) {
-                       ubi_err("bad image seq: 0x%x, expected: 0x%x",
+                       ubi_err(ubi, "bad image seq: 0x%x, expected: 0x%x",
                                be32_to_cpu(ech->image_seq), ubi->image_seq);
                        ret = UBI_BAD_FASTMAP;
                        goto out;
@@ -453,9 +522,9 @@ static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                        unmap_peb(ai, pnum);
                        dbg_bld("Adding PEB to free: %i", pnum);
                        if (err == UBI_IO_FF_BITFLIPS)
-                               add_aeb(ai, freef, pnum, ec, 1);
+                               add_aeb(ai, free, pnum, ec, 1);
                        else
-                               add_aeb(ai, freef, pnum, ec, 0);
+                               add_aeb(ai, free, pnum, ec, 0);
                        continue;
                } else if (err == 0 || err == UBI_IO_BITFLIPS) {
                        dbg_bld("Found non empty PEB:%i in pool", pnum);
@@ -463,18 +532,6 @@ static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                        if (err == UBI_IO_BITFLIPS)
                                scrub = 1;
 
-                       found_orphan = 0;
-                       list_for_each_entry(tmp_aeb, eba_orphans, u.list) {
-                               if (tmp_aeb->pnum == pnum) {
-                                       found_orphan = 1;
-                                       break;
-                               }
-                       }
-                       if (found_orphan) {
-                               list_del(&tmp_aeb->u.list);
-                               kmem_cache_free(ai->aeb_slab_cache, tmp_aeb);
-                       }
-
                        new_aeb = kmem_cache_alloc(ai->aeb_slab_cache,
                                                   GFP_KERNEL);
                        if (!new_aeb) {
@@ -499,7 +556,7 @@ static int scan_pool(struct ubi_device *ubi, struct ubi_attach_info *ai,
                        }
                } else {
                        /* We are paranoid and fall back to scanning mode */
-                       ubi_err("fastmap pool PEBs contains damaged PEBs!");
+                       ubi_err(ubi, "fastmap pool PEBs contains damaged PEBs!");
                        ret = err > 0 ? UBI_BAD_FASTMAP : err;
                        goto out;
                }
@@ -549,13 +606,12 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                              struct ubi_attach_info *ai,
                              struct ubi_fastmap_layout *fm)
 {
-       struct list_head used, eba_orphans, freef;
+       struct list_head used, free;
        struct ubi_ainf_volume *av;
        struct ubi_ainf_peb *aeb, *tmp_aeb, *_tmp_aeb;
-       struct ubi_ec_hdr *ech;
        struct ubi_fm_sb *fmsb;
        struct ubi_fm_hdr *fmhdr;
-       struct ubi_fm_scan_pool *fmpl1, *fmpl2;
+       struct ubi_fm_scan_pool *fmpl, *fmpl_wl;
        struct ubi_fm_ec *fmec;
        struct ubi_fm_volhdr *fmvhdr;
        struct ubi_fm_eba *fm_eba;
@@ -565,23 +621,9 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
        void *fm_raw = ubi->fm_buf;
 
        INIT_LIST_HEAD(&used);
-       INIT_LIST_HEAD(&freef);
-       INIT_LIST_HEAD(&eba_orphans);
-       INIT_LIST_HEAD(&ai->corr);
-       INIT_LIST_HEAD(&ai->free);
-       INIT_LIST_HEAD(&ai->erase);
-       INIT_LIST_HEAD(&ai->alien);
-       ai->volumes = RB_ROOT;
+       INIT_LIST_HEAD(&free);
        ai->min_ec = UBI_MAX_ERASECOUNTER;
 
-       ai->aeb_slab_cache = kmem_cache_create("ubi_ainf_peb_slab",
-                                              sizeof(struct ubi_ainf_peb),
-                                              0, 0, NULL);
-       if (!ai->aeb_slab_cache) {
-               ret = -ENOMEM;
-               goto fail;
-       }
-
        fmsb = (struct ubi_fm_sb *)(fm_raw);
        ai->max_sqnum = fmsb->sqnum;
        fm_pos += sizeof(struct ubi_fm_sb);
@@ -594,56 +636,57 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                goto fail_bad;
 
        if (be32_to_cpu(fmhdr->magic) != UBI_FM_HDR_MAGIC) {
-               ubi_err("bad fastmap header magic: 0x%x, expected: 0x%x",
+               ubi_err(ubi, "bad fastmap header magic: 0x%x, expected: 0x%x",
                        be32_to_cpu(fmhdr->magic), UBI_FM_HDR_MAGIC);
                goto fail_bad;
        }
 
-       fmpl1 = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
-       fm_pos += sizeof(*fmpl1);
+       fmpl = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
+       fm_pos += sizeof(*fmpl);
        if (fm_pos >= fm_size)
                goto fail_bad;
-       if (be32_to_cpu(fmpl1->magic) != UBI_FM_POOL_MAGIC) {
-               ubi_err("bad fastmap pool magic: 0x%x, expected: 0x%x",
-                       be32_to_cpu(fmpl1->magic), UBI_FM_POOL_MAGIC);
+       if (be32_to_cpu(fmpl->magic) != UBI_FM_POOL_MAGIC) {
+               ubi_err(ubi, "bad fastmap pool magic: 0x%x, expected: 0x%x",
+                       be32_to_cpu(fmpl->magic), UBI_FM_POOL_MAGIC);
                goto fail_bad;
        }
 
-       fmpl2 = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
-       fm_pos += sizeof(*fmpl2);
+       fmpl_wl = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
+       fm_pos += sizeof(*fmpl_wl);
        if (fm_pos >= fm_size)
                goto fail_bad;
-       if (be32_to_cpu(fmpl2->magic) != UBI_FM_POOL_MAGIC) {
-               ubi_err("bad fastmap pool magic: 0x%x, expected: 0x%x",
-                       be32_to_cpu(fmpl2->magic), UBI_FM_POOL_MAGIC);
+       if (be32_to_cpu(fmpl_wl->magic) != UBI_FM_POOL_MAGIC) {
+               ubi_err(ubi, "bad fastmap WL pool magic: 0x%x, expected: 0x%x",
+                       be32_to_cpu(fmpl_wl->magic), UBI_FM_POOL_MAGIC);
                goto fail_bad;
        }
 
-       pool_size = be16_to_cpu(fmpl1->size);
-       wl_pool_size = be16_to_cpu(fmpl2->size);
-       fm->max_pool_size = be16_to_cpu(fmpl1->max_size);
-       fm->max_wl_pool_size = be16_to_cpu(fmpl2->max_size);
+       pool_size = be16_to_cpu(fmpl->size);
+       wl_pool_size = be16_to_cpu(fmpl_wl->size);
+       fm->max_pool_size = be16_to_cpu(fmpl->max_size);
+       fm->max_wl_pool_size = be16_to_cpu(fmpl_wl->max_size);
 
        if (pool_size > UBI_FM_MAX_POOL_SIZE || pool_size < 0) {
-               ubi_err("bad pool size: %i", pool_size);
+               ubi_err(ubi, "bad pool size: %i", pool_size);
                goto fail_bad;
        }
 
        if (wl_pool_size > UBI_FM_MAX_POOL_SIZE || wl_pool_size < 0) {
-               ubi_err("bad WL pool size: %i", wl_pool_size);
+               ubi_err(ubi, "bad WL pool size: %i", wl_pool_size);
                goto fail_bad;
        }
 
 
        if (fm->max_pool_size > UBI_FM_MAX_POOL_SIZE ||
            fm->max_pool_size < 0) {
-               ubi_err("bad maximal pool size: %i", fm->max_pool_size);
+               ubi_err(ubi, "bad maximal pool size: %i", fm->max_pool_size);
                goto fail_bad;
        }
 
        if (fm->max_wl_pool_size > UBI_FM_MAX_POOL_SIZE ||
            fm->max_wl_pool_size < 0) {
-               ubi_err("bad maximal WL pool size: %i", fm->max_wl_pool_size);
+               ubi_err(ubi, "bad maximal WL pool size: %i",
+                       fm->max_wl_pool_size);
                goto fail_bad;
        }
 
@@ -702,8 +745,7 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                        goto fail_bad;
 
                if (be32_to_cpu(fmvhdr->magic) != UBI_FM_VHDR_MAGIC) {
-                       ubi_err("bad fastmap vol header magic: 0x%x, " \
-                               "expected: 0x%x",
+                       ubi_err(ubi, "bad fastmap vol header magic: 0x%x, expected: 0x%x",
                                be32_to_cpu(fmvhdr->magic), UBI_FM_VHDR_MAGIC);
                        goto fail_bad;
                }
@@ -716,6 +758,11 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
 
                if (!av)
                        goto fail_bad;
+               if (PTR_ERR(av) == -EINVAL) {
+                       ubi_err(ubi, "volume (ID %i) already exists",
+                               fmvhdr->vol_id);
+                       goto fail_bad;
+               }
 
                ai->vols_found++;
                if (ai->highest_vol_id < be32_to_cpu(fmvhdr->vol_id))
@@ -728,8 +775,7 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                        goto fail_bad;
 
                if (be32_to_cpu(fm_eba->magic) != UBI_FM_EBA_MAGIC) {
-                       ubi_err("bad fastmap EBA header magic: 0x%x, " \
-                               "expected: 0x%x",
+                       ubi_err(ubi, "bad fastmap EBA header magic: 0x%x, expected: 0x%x",
                                be32_to_cpu(fm_eba->magic), UBI_FM_EBA_MAGIC);
                        goto fail_bad;
                }
@@ -748,28 +794,9 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                                }
                        }
 
-                       /* This can happen if a PEB is already in an EBA known
-                        * by this fastmap but the PEB itself is not in the used
-                        * list.
-                        * In this case the PEB can be within the fastmap pool
-                        * or while writing the fastmap it was in the protection
-                        * queue.
-                        */
                        if (!aeb) {
-                               aeb = kmem_cache_alloc(ai->aeb_slab_cache,
-                                                      GFP_KERNEL);
-                               if (!aeb) {
-                                       ret = -ENOMEM;
-
-                                       goto fail;
-                               }
-
-                               aeb->lnum = j;
-                               aeb->pnum = be32_to_cpu(fm_eba->pnum[j]);
-                               aeb->ec = -1;
-                               aeb->scrub = aeb->copy_flag = aeb->sqnum = 0;
-                               list_add_tail(&aeb->u.list, &eba_orphans);
-                               continue;
+                               ubi_err(ubi, "PEB %i is in EBA but not in used list", pnum);
+                               goto fail_bad;
                        }
 
                        aeb->lnum = j;
@@ -782,61 +809,26 @@ static int ubi_attach_fastmap(struct ubi_device *ubi,
                        dbg_bld("inserting PEB:%i (LEB %i) to vol %i",
                                aeb->pnum, aeb->lnum, av->vol_id);
                }
-
-               ech = kzalloc(ubi->ec_hdr_alsize, GFP_KERNEL);
-               if (!ech) {
-                       ret = -ENOMEM;
-                       goto fail;
-               }
-
-               list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &eba_orphans,
-                                        u.list) {
-                       int err;
-
-                       if (ubi_io_is_bad(ubi, tmp_aeb->pnum)) {
-                               ubi_err("bad PEB in fastmap EBA orphan list");
-                               ret = UBI_BAD_FASTMAP;
-                               kfree(ech);
-                               goto fail;
-                       }
-
-                       err = ubi_io_read_ec_hdr(ubi, tmp_aeb->pnum, ech, 0);
-                       if (err && err != UBI_IO_BITFLIPS) {
-                               ubi_err("unable to read EC header! PEB:%i " \
-                                       "err:%i", tmp_aeb->pnum, err);
-                               ret = err > 0 ? UBI_BAD_FASTMAP : err;
-                               kfree(ech);
-
-                               goto fail;
-                       } else if (err == UBI_IO_BITFLIPS)
-                               tmp_aeb->scrub = 1;
-
-                       tmp_aeb->ec = be64_to_cpu(ech->ec);
-                       assign_aeb_to_av(ai, tmp_aeb, av);
-               }
-
-               kfree(ech);
        }
 
-       ret = scan_pool(ubi, ai, fmpl1->pebs, pool_size, &max_sqnum,
-                       &eba_orphans, &freef);
+       ret = scan_pool(ubi, ai, fmpl->pebs, pool_size, &max_sqnum, &free);
        if (ret)
                goto fail;
 
-       ret = scan_pool(ubi, ai, fmpl2->pebs, wl_pool_size, &max_sqnum,
-                       &eba_orphans, &freef);
+       ret = scan_pool(ubi, ai, fmpl_wl->pebs, wl_pool_size, &max_sqnum, &free);
        if (ret)
                goto fail;
 
        if (max_sqnum > ai->max_sqnum)
                ai->max_sqnum = max_sqnum;
 
-       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &freef, u.list)
+       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &free, u.list)
                list_move_tail(&tmp_aeb->u.list, &ai->free);
 
-       ubi_assert(list_empty(&used));
-       ubi_assert(list_empty(&eba_orphans));
-       ubi_assert(list_empty(&freef));
+       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &used, u.list)
+               list_move_tail(&tmp_aeb->u.list, &ai->erase);
+
+       ubi_assert(list_empty(&free));
 
        /*
         * If fastmap is leaking PEBs (must not happen), raise a
@@ -865,11 +857,7 @@ fail:
                list_del(&tmp_aeb->u.list);
                kmem_cache_free(ai->aeb_slab_cache, tmp_aeb);
        }
-       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &eba_orphans, u.list) {
-               list_del(&tmp_aeb->u.list);
-               kmem_cache_free(ai->aeb_slab_cache, tmp_aeb);
-       }
-       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &freef, u.list) {
+       list_for_each_entry_safe(tmp_aeb, _tmp_aeb, &free, u.list) {
                list_del(&tmp_aeb->u.list);
                kmem_cache_free(ai->aeb_slab_cache, tmp_aeb);
        }
@@ -899,7 +887,7 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
        __be32 crc, tmp_crc;
        unsigned long long sqnum = 0;
 
-       mutex_lock(&ubi->fm_mutex);
+       down_write(&ubi->fm_protect);
        memset(ubi->fm_buf, 0, ubi->fm_size);
 
        fmsb = kmalloc(sizeof(*fmsb), GFP_KERNEL);
@@ -922,14 +910,14 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
                fm->to_be_tortured[0] = 1;
 
        if (be32_to_cpu(fmsb->magic) != UBI_FM_SB_MAGIC) {
-               ubi_err("bad super block magic: 0x%x, expected: 0x%x",
+               ubi_err(ubi, "bad super block magic: 0x%x, expected: 0x%x",
                        be32_to_cpu(fmsb->magic), UBI_FM_SB_MAGIC);
                ret = UBI_BAD_FASTMAP;
                goto free_fm_sb;
        }
 
        if (fmsb->version != UBI_FM_FMT_VERSION) {
-               ubi_err("bad fastmap version: %i, expected: %i",
+               ubi_err(ubi, "bad fastmap version: %i, expected: %i",
                        fmsb->version, UBI_FM_FMT_VERSION);
                ret = UBI_BAD_FASTMAP;
                goto free_fm_sb;
@@ -937,15 +925,16 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
 
        used_blocks = be32_to_cpu(fmsb->used_blocks);
        if (used_blocks > UBI_FM_MAX_BLOCKS || used_blocks < 1) {
-               ubi_err("number of fastmap blocks is invalid: %i", used_blocks);
+               ubi_err(ubi, "number of fastmap blocks is invalid: %i",
+                       used_blocks);
                ret = UBI_BAD_FASTMAP;
                goto free_fm_sb;
        }
 
        fm_size = ubi->leb_size * used_blocks;
        if (fm_size != ubi->fm_size) {
-               ubi_err("bad fastmap size: %zi, expected: %zi", fm_size,
-                       ubi->fm_size);
+               ubi_err(ubi, "bad fastmap size: %zi, expected: %zi",
+                       fm_size, ubi->fm_size);
                ret = UBI_BAD_FASTMAP;
                goto free_fm_sb;
        }
@@ -974,7 +963,7 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
 
                ret = ubi_io_read_ec_hdr(ubi, pnum, ech, 0);
                if (ret && ret != UBI_IO_BITFLIPS) {
-                       ubi_err("unable to read fastmap block# %i EC (PEB: %i)",
+                       ubi_err(ubi, "unable to read fastmap block# %i EC (PEB: %i)",
                                i, pnum);
                        if (ret > 0)
                                ret = UBI_BAD_FASTMAP;
@@ -991,7 +980,7 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
                 * we shouldn't fail if image_seq == 0.
                 */
                if (image_seq && (image_seq != ubi->image_seq)) {
-                       ubi_err("wrong image seq:%d instead of %d",
+                       ubi_err(ubi, "wrong image seq:%d instead of %d",
                                be32_to_cpu(ech->image_seq), ubi->image_seq);
                        ret = UBI_BAD_FASTMAP;
                        goto free_hdr;
@@ -999,15 +988,14 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
 
                ret = ubi_io_read_vid_hdr(ubi, pnum, vh, 0);
                if (ret && ret != UBI_IO_BITFLIPS) {
-                       ubi_err("unable to read fastmap block# %i (PEB: %i)",
+                       ubi_err(ubi, "unable to read fastmap block# %i (PEB: %i)",
                                i, pnum);
                        goto free_hdr;
                }
 
                if (i == 0) {
                        if (be32_to_cpu(vh->vol_id) != UBI_FM_SB_VOLUME_ID) {
-                               ubi_err("bad fastmap anchor vol_id: 0x%x," \
-                                       " expected: 0x%x",
+                               ubi_err(ubi, "bad fastmap anchor vol_id: 0x%x, expected: 0x%x",
                                        be32_to_cpu(vh->vol_id),
                                        UBI_FM_SB_VOLUME_ID);
                                ret = UBI_BAD_FASTMAP;
@@ -1015,8 +1003,7 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
                        }
                } else {
                        if (be32_to_cpu(vh->vol_id) != UBI_FM_DATA_VOLUME_ID) {
-                               ubi_err("bad fastmap data vol_id: 0x%x," \
-                                       " expected: 0x%x",
+                               ubi_err(ubi, "bad fastmap data vol_id: 0x%x, expected: 0x%x",
                                        be32_to_cpu(vh->vol_id),
                                        UBI_FM_DATA_VOLUME_ID);
                                ret = UBI_BAD_FASTMAP;
@@ -1030,7 +1017,7 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
                ret = ubi_io_read(ubi, ubi->fm_buf + (ubi->leb_size * i), pnum,
                                  ubi->leb_start, ubi->leb_size);
                if (ret && ret != UBI_IO_BITFLIPS) {
-                       ubi_err("unable to read fastmap block# %i (PEB: %i, " \
+                       ubi_err(ubi, "unable to read fastmap block# %i (PEB: %i, "
                                "err: %i)", i, pnum, ret);
                        goto free_hdr;
                }
@@ -1044,8 +1031,9 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
        fmsb2->data_crc = 0;
        crc = crc32(UBI_CRC32_INIT, ubi->fm_buf, fm_size);
        if (crc != tmp_crc) {
-               ubi_err("fastmap data CRC is invalid");
-               ubi_err("CRC should be: 0x%x, calc: 0x%x", tmp_crc, crc);
+               ubi_err(ubi, "fastmap data CRC is invalid");
+               ubi_err(ubi, "CRC should be: 0x%x, calc: 0x%x",
+                       tmp_crc, crc);
                ret = UBI_BAD_FASTMAP;
                goto free_hdr;
        }
@@ -1081,17 +1069,18 @@ int ubi_scan_fastmap(struct ubi_device *ubi, struct ubi_attach_info *ai,
        ubi->fm = fm;
        ubi->fm_pool.max_size = ubi->fm->max_pool_size;
        ubi->fm_wl_pool.max_size = ubi->fm->max_wl_pool_size;
-       ubi_msg("attached by fastmap");
-       ubi_msg("fastmap pool size: %d", ubi->fm_pool.max_size);
-       ubi_msg("fastmap WL pool size: %d", ubi->fm_wl_pool.max_size);
+       ubi_msg(ubi, "attached by fastmap");
+       ubi_msg(ubi, "fastmap pool size: %d", ubi->fm_pool.max_size);
+       ubi_msg(ubi, "fastmap WL pool size: %d",
+               ubi->fm_wl_pool.max_size);
        ubi->fm_disabled = 0;
 
        ubi_free_vid_hdr(ubi, vh);
        kfree(ech);
 out:
-       mutex_unlock(&ubi->fm_mutex);
+       up_write(&ubi->fm_protect);
        if (ret == UBI_BAD_FASTMAP)
-               ubi_err("Attach by fastmap failed, doing a full scan!");
+               ubi_err(ubi, "Attach by fastmap failed, doing a full scan!");
        return ret;
 
 free_hdr:
@@ -1117,17 +1106,18 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        void *fm_raw;
        struct ubi_fm_sb *fmsb;
        struct ubi_fm_hdr *fmh;
-       struct ubi_fm_scan_pool *fmpl1, *fmpl2;
+       struct ubi_fm_scan_pool *fmpl, *fmpl_wl;
        struct ubi_fm_ec *fec;
        struct ubi_fm_volhdr *fvh;
        struct ubi_fm_eba *feba;
-       struct rb_node *node;
        struct ubi_wl_entry *wl_e;
        struct ubi_volume *vol;
        struct ubi_vid_hdr *avhdr, *dvhdr;
        struct ubi_work *ubi_wrk;
+       struct rb_node *tmp_rb;
        int ret, i, j, free_peb_count, used_peb_count, vol_count;
        int scrub_peb_count, erase_peb_count;
+       int *seen_pebs = NULL;
 
        fm_raw = ubi->fm_buf;
        memset(ubi->fm_buf, 0, ubi->fm_size);
@@ -1144,6 +1134,12 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
                goto out_kfree;
        }
 
+       seen_pebs = init_seen(ubi);
+       if (IS_ERR(seen_pebs)) {
+               ret = PTR_ERR(seen_pebs);
+               goto out_kfree;
+       }
+
        spin_lock(&ubi->volumes_lock);
        spin_lock(&ubi->wl_lock);
 
@@ -1168,29 +1164,33 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        erase_peb_count = 0;
        vol_count = 0;
 
-       fmpl1 = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
-       fm_pos += sizeof(*fmpl1);
-       fmpl1->magic = cpu_to_be32(UBI_FM_POOL_MAGIC);
-       fmpl1->size = cpu_to_be16(ubi->fm_pool.size);
-       fmpl1->max_size = cpu_to_be16(ubi->fm_pool.max_size);
+       fmpl = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
+       fm_pos += sizeof(*fmpl);
+       fmpl->magic = cpu_to_be32(UBI_FM_POOL_MAGIC);
+       fmpl->size = cpu_to_be16(ubi->fm_pool.size);
+       fmpl->max_size = cpu_to_be16(ubi->fm_pool.max_size);
 
-       for (i = 0; i < ubi->fm_pool.size; i++)
-               fmpl1->pebs[i] = cpu_to_be32(ubi->fm_pool.pebs[i]);
+       for (i = 0; i < ubi->fm_pool.size; i++) {
+               fmpl->pebs[i] = cpu_to_be32(ubi->fm_pool.pebs[i]);
+               set_seen(ubi, ubi->fm_pool.pebs[i], seen_pebs);
+       }
 
-       fmpl2 = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
-       fm_pos += sizeof(*fmpl2);
-       fmpl2->magic = cpu_to_be32(UBI_FM_POOL_MAGIC);
-       fmpl2->size = cpu_to_be16(ubi->fm_wl_pool.size);
-       fmpl2->max_size = cpu_to_be16(ubi->fm_wl_pool.max_size);
+       fmpl_wl = (struct ubi_fm_scan_pool *)(fm_raw + fm_pos);
+       fm_pos += sizeof(*fmpl_wl);
+       fmpl_wl->magic = cpu_to_be32(UBI_FM_POOL_MAGIC);
+       fmpl_wl->size = cpu_to_be16(ubi->fm_wl_pool.size);
+       fmpl_wl->max_size = cpu_to_be16(ubi->fm_wl_pool.max_size);
 
-       for (i = 0; i < ubi->fm_wl_pool.size; i++)
-               fmpl2->pebs[i] = cpu_to_be32(ubi->fm_wl_pool.pebs[i]);
+       for (i = 0; i < ubi->fm_wl_pool.size; i++) {
+               fmpl_wl->pebs[i] = cpu_to_be32(ubi->fm_wl_pool.pebs[i]);
+               set_seen(ubi, ubi->fm_wl_pool.pebs[i], seen_pebs);
+       }
 
-       for (node = rb_first(&ubi->free); node; node = rb_next(node)) {
-               wl_e = rb_entry(node, struct ubi_wl_entry, u.rb);
+       ubi_for_each_free_peb(ubi, wl_e, tmp_rb) {
                fec = (struct ubi_fm_ec *)(fm_raw + fm_pos);
 
                fec->pnum = cpu_to_be32(wl_e->pnum);
+               set_seen(ubi, wl_e->pnum, seen_pebs);
                fec->ec = cpu_to_be32(wl_e->ec);
 
                free_peb_count++;
@@ -1199,11 +1199,23 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        }
        fmh->free_peb_count = cpu_to_be32(free_peb_count);
 
-       for (node = rb_first(&ubi->used); node; node = rb_next(node)) {
-               wl_e = rb_entry(node, struct ubi_wl_entry, u.rb);
+       ubi_for_each_used_peb(ubi, wl_e, tmp_rb) {
+               fec = (struct ubi_fm_ec *)(fm_raw + fm_pos);
+
+               fec->pnum = cpu_to_be32(wl_e->pnum);
+               set_seen(ubi, wl_e->pnum, seen_pebs);
+               fec->ec = cpu_to_be32(wl_e->ec);
+
+               used_peb_count++;
+               fm_pos += sizeof(*fec);
+               ubi_assert(fm_pos <= ubi->fm_size);
+       }
+
+       ubi_for_each_protected_peb(ubi, i, wl_e) {
                fec = (struct ubi_fm_ec *)(fm_raw + fm_pos);
 
                fec->pnum = cpu_to_be32(wl_e->pnum);
+               set_seen(ubi, wl_e->pnum, seen_pebs);
                fec->ec = cpu_to_be32(wl_e->ec);
 
                used_peb_count++;
@@ -1212,11 +1224,11 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        }
        fmh->used_peb_count = cpu_to_be32(used_peb_count);
 
-       for (node = rb_first(&ubi->scrub); node; node = rb_next(node)) {
-               wl_e = rb_entry(node, struct ubi_wl_entry, u.rb);
+       ubi_for_each_scrub_peb(ubi, wl_e, tmp_rb) {
                fec = (struct ubi_fm_ec *)(fm_raw + fm_pos);
 
                fec->pnum = cpu_to_be32(wl_e->pnum);
+               set_seen(ubi, wl_e->pnum, seen_pebs);
                fec->ec = cpu_to_be32(wl_e->ec);
 
                scrub_peb_count++;
@@ -1234,6 +1246,7 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
                        fec = (struct ubi_fm_ec *)(fm_raw + fm_pos);
 
                        fec->pnum = cpu_to_be32(wl_e->pnum);
+                       set_seen(ubi, wl_e->pnum, seen_pebs);
                        fec->ec = cpu_to_be32(wl_e->ec);
 
                        erase_peb_count++;
@@ -1287,12 +1300,13 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        dbg_bld("writing fastmap SB to PEB %i", new_fm->e[0]->pnum);
        ret = ubi_io_write_vid_hdr(ubi, new_fm->e[0]->pnum, avhdr);
        if (ret) {
-               ubi_err("unable to write vid_hdr to fastmap SB!");
+               ubi_err(ubi, "unable to write vid_hdr to fastmap SB!");
                goto out_kfree;
        }
 
        for (i = 0; i < new_fm->used_blocks; i++) {
                fmsb->block_loc[i] = cpu_to_be32(new_fm->e[i]->pnum);
+               set_seen(ubi, new_fm->e[i]->pnum, seen_pebs);
                fmsb->block_ec[i] = cpu_to_be32(new_fm->e[i]->ec);
        }
 
@@ -1307,7 +1321,7 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
                        new_fm->e[i]->pnum, be64_to_cpu(dvhdr->sqnum));
                ret = ubi_io_write_vid_hdr(ubi, new_fm->e[i]->pnum, dvhdr);
                if (ret) {
-                       ubi_err("unable to write vid_hdr to PEB %i!",
+                       ubi_err(ubi, "unable to write vid_hdr to PEB %i!",
                                new_fm->e[i]->pnum);
                        goto out_kfree;
                }
@@ -1317,7 +1331,7 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
                ret = ubi_io_write(ubi, fm_raw + (i * ubi->leb_size),
                        new_fm->e[i]->pnum, ubi->leb_start, ubi->leb_size);
                if (ret) {
-                       ubi_err("unable to write fastmap to PEB %i!",
+                       ubi_err(ubi, "unable to write fastmap to PEB %i!",
                                new_fm->e[i]->pnum);
                        goto out_kfree;
                }
@@ -1326,11 +1340,13 @@ static int ubi_write_fastmap(struct ubi_device *ubi,
        ubi_assert(new_fm);
        ubi->fm = new_fm;
 
+       ret = self_check_seen(ubi, seen_pebs);
        dbg_bld("fastmap written!");
 
 out_kfree:
        ubi_free_vid_hdr(ubi, avhdr);
        ubi_free_vid_hdr(ubi, dvhdr);
+       free_seen(seen_pebs);
 out:
        return ret;
 }
@@ -1385,31 +1401,87 @@ out:
 /**
  * invalidate_fastmap - destroys a fastmap.
  * @ubi: UBI device object
- * @fm: the fastmap to be destroyed
  *
+ * This function ensures that upon next UBI attach a full scan
+ * is issued. We need this if UBI is about to write a new fastmap
+ * but is unable to do so. In this case we have two options:
+ * a) Make sure that the current fastmap will not be usued upon
+ * attach time and contine or b) fall back to RO mode to have the
+ * current fastmap in a valid state.
  * Returns 0 on success, < 0 indicates an internal error.
  */
-static int invalidate_fastmap(struct ubi_device *ubi,
-                             struct ubi_fastmap_layout *fm)
+static int invalidate_fastmap(struct ubi_device *ubi)
 {
        int ret;
-       struct ubi_vid_hdr *vh;
+       struct ubi_fastmap_layout *fm;
+       struct ubi_wl_entry *e;
+       struct ubi_vid_hdr *vh = NULL;
 
-       ret = erase_block(ubi, fm->e[0]->pnum);
-       if (ret < 0)
-               return ret;
+       if (!ubi->fm)
+               return 0;
+
+       ubi->fm = NULL;
+
+       ret = -ENOMEM;
+       fm = kzalloc(sizeof(*fm), GFP_KERNEL);
+       if (!fm)
+               goto out;
 
        vh = new_fm_vhdr(ubi, UBI_FM_SB_VOLUME_ID);
        if (!vh)
-               return -ENOMEM;
+               goto out_free_fm;
 
-       /* deleting the current fastmap SB is not enough, an old SB may exist,
-        * so create a (corrupted) SB such that fastmap will find it and fall
-        * back to scanning mode in any case */
+       ret = -ENOSPC;
+       e = ubi_wl_get_fm_peb(ubi, 1);
+       if (!e)
+               goto out_free_fm;
+
+       /*
+        * Create fake fastmap such that UBI will fall back
+        * to scanning mode.
+        */
        vh->sqnum = cpu_to_be64(ubi_next_sqnum(ubi));
-       ret = ubi_io_write_vid_hdr(ubi, fm->e[0]->pnum, vh);
+       ret = ubi_io_write_vid_hdr(ubi, e->pnum, vh);
+       if (ret < 0) {
+               ubi_wl_put_fm_peb(ubi, e, 0, 0);
+               goto out_free_fm;
+       }
+
+       fm->used_blocks = 1;
+       fm->e[0] = e;
+
+       ubi->fm = fm;
 
+out:
+       ubi_free_vid_hdr(ubi, vh);
        return ret;
+
+out_free_fm:
+       kfree(fm);
+       goto out;
+}
+
+/**
+ * return_fm_pebs - returns all PEBs used by a fastmap back to the
+ * WL sub-system.
+ * @ubi: UBI device object
+ * @fm: fastmap layout object
+ */
+static void return_fm_pebs(struct ubi_device *ubi,
+                          struct ubi_fastmap_layout *fm)
+{
+       int i;
+
+       if (!fm)
+               return;
+
+       for (i = 0; i < fm->used_blocks; i++) {
+               if (fm->e[i]) {
+                       ubi_wl_put_fm_peb(ubi, fm->e[i], i,
+                                         fm->to_be_tortured[i]);
+                       fm->e[i] = NULL;
+               }
+       }
 }
 
 /**
@@ -1421,50 +1493,37 @@ static int invalidate_fastmap(struct ubi_device *ubi,
  */
 int ubi_update_fastmap(struct ubi_device *ubi)
 {
-       int ret, i;
+       int ret, i, j;
        struct ubi_fastmap_layout *new_fm, *old_fm;
        struct ubi_wl_entry *tmp_e;
 
-       mutex_lock(&ubi->fm_mutex);
+       down_write(&ubi->fm_protect);
 
        ubi_refill_pools(ubi);
 
        if (ubi->ro_mode || ubi->fm_disabled) {
-               mutex_unlock(&ubi->fm_mutex);
+               up_write(&ubi->fm_protect);
                return 0;
        }
 
        ret = ubi_ensure_anchor_pebs(ubi);
        if (ret) {
-               mutex_unlock(&ubi->fm_mutex);
+               up_write(&ubi->fm_protect);
                return ret;
        }
 
        new_fm = kzalloc(sizeof(*new_fm), GFP_KERNEL);
        if (!new_fm) {
-               mutex_unlock(&ubi->fm_mutex);
+               up_write(&ubi->fm_protect);
                return -ENOMEM;
        }
 
        new_fm->used_blocks = ubi->fm_size / ubi->leb_size;
-
-       for (i = 0; i < new_fm->used_blocks; i++) {
-               new_fm->e[i] = kmem_cache_alloc(ubi_wl_entry_slab, GFP_KERNEL);
-               if (!new_fm->e[i]) {
-                       while (i--)
-                               kfree(new_fm->e[i]);
-
-                       kfree(new_fm);
-                       mutex_unlock(&ubi->fm_mutex);
-                       return -ENOMEM;
-               }
-       }
-
        old_fm = ubi->fm;
        ubi->fm = NULL;
 
        if (new_fm->used_blocks > UBI_FM_MAX_BLOCKS) {
-               ubi_err("fastmap too large");
+               ubi_err(ubi, "fastmap too large");
                ret = -ENOSPC;
                goto err;
        }
@@ -1474,37 +1533,49 @@ int ubi_update_fastmap(struct ubi_device *ubi)
                tmp_e = ubi_wl_get_fm_peb(ubi, 0);
                spin_unlock(&ubi->wl_lock);
 
-               if (!tmp_e && !old_fm) {
-                       int j;
-                       ubi_err("could not get any free erase block");
-
-                       for (j = 1; j < i; j++)
-                               ubi_wl_put_fm_peb(ubi, new_fm->e[j], j, 0);
-
-                       ret = -ENOSPC;
-                       goto err;
-               } else if (!tmp_e && old_fm) {
-                       ret = erase_block(ubi, old_fm->e[i]->pnum);
-                       if (ret < 0) {
-                               int j;
-
-                               for (j = 1; j < i; j++)
-                                       ubi_wl_put_fm_peb(ubi, new_fm->e[j],
-                                                         j, 0);
+               if (!tmp_e) {
+                       if (old_fm && old_fm->e[i]) {
+                               ret = erase_block(ubi, old_fm->e[i]->pnum);
+                               if (ret < 0) {
+                                       ubi_err(ubi, "could not erase old fastmap PEB");
+
+                                       for (j = 1; j < i; j++) {
+                                               ubi_wl_put_fm_peb(ubi, new_fm->e[j],
+                                                                 j, 0);
+                                               new_fm->e[j] = NULL;
+                                       }
+                                       goto err;
+                               }
+                               new_fm->e[i] = old_fm->e[i];
+                               old_fm->e[i] = NULL;
+                       } else {
+                               ubi_err(ubi, "could not get any free erase block");
+
+                               for (j = 1; j < i; j++) {
+                                       ubi_wl_put_fm_peb(ubi, new_fm->e[j], j, 0);
+                                       new_fm->e[j] = NULL;
+                               }
 
-                               ubi_err("could not erase old fastmap PEB");
+                               ret = -ENOSPC;
                                goto err;
                        }
-
-                       new_fm->e[i]->pnum = old_fm->e[i]->pnum;
-                       new_fm->e[i]->ec = old_fm->e[i]->ec;
                } else {
-                       new_fm->e[i]->pnum = tmp_e->pnum;
-                       new_fm->e[i]->ec = tmp_e->ec;
+                       new_fm->e[i] = tmp_e;
 
-                       if (old_fm)
+                       if (old_fm && old_fm->e[i]) {
                                ubi_wl_put_fm_peb(ubi, old_fm->e[i], i,
                                                  old_fm->to_be_tortured[i]);
+                               old_fm->e[i] = NULL;
+                       }
+               }
+       }
+
+       /* Old fastmap is larger than the new one */
+       if (old_fm && new_fm->used_blocks < old_fm->used_blocks) {
+               for (i = new_fm->used_blocks; i < old_fm->used_blocks; i++) {
+                       ubi_wl_put_fm_peb(ubi, old_fm->e[i], i,
+                                         old_fm->to_be_tortured[i]);
+                       old_fm->e[i] = NULL;
                }
        }
 
@@ -1517,67 +1588,67 @@ int ubi_update_fastmap(struct ubi_device *ubi)
                if (!tmp_e) {
                        ret = erase_block(ubi, old_fm->e[0]->pnum);
                        if (ret < 0) {
-                               int i;
-                               ubi_err("could not erase old anchor PEB");
+                               ubi_err(ubi, "could not erase old anchor PEB");
 
-                               for (i = 1; i < new_fm->used_blocks; i++)
+                               for (i = 1; i < new_fm->used_blocks; i++) {
                                        ubi_wl_put_fm_peb(ubi, new_fm->e[i],
                                                          i, 0);
+                                       new_fm->e[i] = NULL;
+                               }
                                goto err;
                        }
-
-                       new_fm->e[0]->pnum = old_fm->e[0]->pnum;
+                       new_fm->e[0] = old_fm->e[0];
                        new_fm->e[0]->ec = ret;
+                       old_fm->e[0] = NULL;
                } else {
                        /* we've got a new anchor PEB, return the old one */
                        ubi_wl_put_fm_peb(ubi, old_fm->e[0], 0,
                                          old_fm->to_be_tortured[0]);
-
-                       new_fm->e[0]->pnum = tmp_e->pnum;
-                       new_fm->e[0]->ec = tmp_e->ec;
+                       new_fm->e[0] = tmp_e;
+                       old_fm->e[0] = NULL;
                }
        } else {
                if (!tmp_e) {
-                       int i;
-                       ubi_err("could not find any anchor PEB");
+                       ubi_err(ubi, "could not find any anchor PEB");
 
-                       for (i = 1; i < new_fm->used_blocks; i++)
+                       for (i = 1; i < new_fm->used_blocks; i++) {
                                ubi_wl_put_fm_peb(ubi, new_fm->e[i], i, 0);
+                               new_fm->e[i] = NULL;
+                       }
 
                        ret = -ENOSPC;
                        goto err;
                }
-
-               new_fm->e[0]->pnum = tmp_e->pnum;
-               new_fm->e[0]->ec = tmp_e->ec;
+               new_fm->e[0] = tmp_e;
        }
 
        down_write(&ubi->work_sem);
-       down_write(&ubi->fm_sem);
+       down_write(&ubi->fm_eba_sem);
        ret = ubi_write_fastmap(ubi, new_fm);
-       up_write(&ubi->fm_sem);
+       up_write(&ubi->fm_eba_sem);
        up_write(&ubi->work_sem);
 
        if (ret)
                goto err;
 
 out_unlock:
-       mutex_unlock(&ubi->fm_mutex);
+       up_write(&ubi->fm_protect);
        kfree(old_fm);
        return ret;
 
 err:
-       kfree(new_fm);
-
-       ubi_warn("Unable to write new fastmap, err=%i", ret);
+       ubi_warn(ubi, "Unable to write new fastmap, err=%i", ret);
 
-       ret = 0;
-       if (old_fm) {
-               ret = invalidate_fastmap(ubi, old_fm);
-               if (ret < 0)
-                       ubi_err("Unable to invalidiate current fastmap!");
-               else if (ret)
-                       ret = 0;
+       ret = invalidate_fastmap(ubi);
+       if (ret < 0) {
+               ubi_err(ubi, "Unable to invalidiate current fastmap!");
+               ubi_ro_mode(ubi);
+       } else {
+               return_fm_pebs(ubi, old_fm);
+               return_fm_pebs(ubi, new_fm);
+               ret = 0;
        }
+
+       kfree(new_fm);
        goto out_unlock;
 }