#include "print-tree.h"
 #include "volumes.h"
 
-struct stripe {
-       struct btrfs_device *dev;
-       u64 physical;
-};
-
-struct multi_bio {
-       atomic_t stripes;
-       bio_end_io_t *end_io;
-       void *private;
-       int error;
-};
-
 struct map_lookup {
        u64 type;
        int io_align;
        int stripe_len;
        int sector_size;
        int num_stripes;
-       struct stripe stripes[];
+       struct btrfs_bio_stripe stripes[];
 };
 
 #define map_lookup_size(n) (sizeof(struct map_lookup) + \
-                           (sizeof(struct stripe) * (n)))
+                           (sizeof(struct btrfs_bio_stripe) * (n)))
 
 static DEFINE_MUTEX(uuid_mutex);
 static LIST_HEAD(fs_uuids);
 }
 
 int btrfs_map_block(struct btrfs_mapping_tree *map_tree, int rw,
-                   int dev_nr, u64 logical, u64 *phys, u64 *length,
-                   struct btrfs_device **dev, int *total_devs)
+                   u64 logical, u64 *length,
+                   struct btrfs_multi_bio **multi_ret)
 {
        struct extent_map *em;
        struct map_lookup *map;
        u64 offset;
        u64 stripe_offset;
        u64 stripe_nr;
+       int stripes_allocated = 8;
        int stripe_index;
+       int i;
+       struct btrfs_multi_bio *multi = NULL;
 
+       if (multi_ret && !(rw & (1 << BIO_RW))) {
+               stripes_allocated = 1;
+       }
+again:
+       if (multi_ret) {
+               multi = kzalloc(btrfs_multi_bio_size(stripes_allocated),
+                               GFP_NOFS);
+               if (!multi)
+                       return -ENOMEM;
+       }
 
        spin_lock(&em_tree->lock);
        em = lookup_extent_mapping(em_tree, logical, *length);
        map = (struct map_lookup *)em->bdev;
        offset = logical - em->start;
 
+       /* if our multi bio struct is too small, back off and try again */
+       if (multi_ret && (rw & (1 << BIO_RW)) &&
+           stripes_allocated < map->num_stripes &&
+           ((map->type & BTRFS_BLOCK_GROUP_RAID1) ||
+            (map->type & BTRFS_BLOCK_GROUP_DUP))) {
+               stripes_allocated = map->num_stripes;
+               spin_unlock(&em_tree->lock);
+               free_extent_map(em);
+               kfree(multi);
+               goto again;
+       }
        stripe_nr = offset;
        /*
         * stripe_nr counts the total number of stripes we have to stride
        /* stripe_offset is the offset of this block in its stripe*/
        stripe_offset = offset - stripe_offset;
 
+       if (map->type & (BTRFS_BLOCK_GROUP_RAID0 | BTRFS_BLOCK_GROUP_RAID1 |
+                        BTRFS_BLOCK_GROUP_DUP)) {
+               /* we limit the length of each bio to what fits in a stripe */
+               *length = min_t(u64, em->len - offset,
+                             map->stripe_len - stripe_offset);
+       } else {
+               *length = em->len - offset;
+       }
+       if (!multi_ret)
+               goto out;
+
+       multi->num_stripes = 1;
+       stripe_index = 0;
        if (map->type & BTRFS_BLOCK_GROUP_RAID1) {
-               stripe_index = dev_nr;
                if (rw & (1 << BIO_RW))
-                       *total_devs = map->num_stripes;
+                       multi->num_stripes = map->num_stripes;
                else {
                        int i;
                        u64 least = (u64)-1;
                                }
                                spin_unlock(&cur->io_lock);
                        }
-                       *total_devs = 1;
                }
        } else if (map->type & BTRFS_BLOCK_GROUP_DUP) {
-               if (rw == WRITE) {
-                       *total_devs = map->num_stripes;
-                       stripe_index = dev_nr;
-               } else {
-                       stripe_index = 0;
-                       *total_devs = 1;
-               }
+               if (rw & (1 << BIO_RW))
+                       multi->num_stripes = map->num_stripes;
        } else {
                /*
                 * after this do_div call, stripe_nr is the number of stripes
                stripe_index = do_div(stripe_nr, map->num_stripes);
        }
        BUG_ON(stripe_index >= map->num_stripes);
-       *phys = map->stripes[stripe_index].physical + stripe_offset +
-               stripe_nr * map->stripe_len;
-
-       if (map->type & (BTRFS_BLOCK_GROUP_RAID0 | BTRFS_BLOCK_GROUP_RAID1 |
-                        BTRFS_BLOCK_GROUP_DUP)) {
-               /* we limit the length of each bio to what fits in a stripe */
-               *length = min_t(u64, em->len - offset,
-                             map->stripe_len - stripe_offset);
-       } else {
-               *length = em->len - offset;
+       BUG_ON(stripe_index != 0 && multi->num_stripes > 1);
+
+       for (i = 0; i < multi->num_stripes; i++) {
+               multi->stripes[i].physical =
+                       map->stripes[stripe_index].physical + stripe_offset +
+                       stripe_nr * map->stripe_len;
+               multi->stripes[i].dev = map->stripes[stripe_index].dev;
+               stripe_index++;
        }
-       *dev = map->stripes[stripe_index].dev;
+       *multi_ret = multi;
+out:
        free_extent_map(em);
        spin_unlock(&em_tree->lock);
        return 0;
                                   unsigned int bytes_done, int err)
 #endif
 {
-       struct multi_bio *multi = bio->bi_private;
+       struct btrfs_multi_bio *multi = bio->bi_private;
 
 #if LINUX_VERSION_CODE <= KERNEL_VERSION(2,6,23)
        if (bio->bi_size)
        if (err)
                multi->error = err;
 
-       if (atomic_dec_and_test(&multi->stripes)) {
+       if (atomic_dec_and_test(&multi->stripes_pending)) {
                bio->bi_private = multi->private;
                bio->bi_end_io = multi->end_io;
 
        struct btrfs_device *dev;
        struct bio *first_bio = bio;
        u64 logical = bio->bi_sector << 9;
-       u64 physical;
        u64 length = 0;
        u64 map_length;
        struct bio_vec *bvec;
-       struct multi_bio *multi = NULL;
+       struct btrfs_multi_bio *multi = NULL;
        int i;
        int ret;
        int dev_nr = 0;
 
        map_tree = &root->fs_info->mapping_tree;
        map_length = length;
+
+       ret = btrfs_map_block(map_tree, rw, logical, &map_length, &multi);
+       BUG_ON(ret);
+
+       total_devs = multi->num_stripes;
+       if (map_length < length) {
+               printk("mapping failed logical %Lu bio len %Lu "
+                      "len %Lu\n", logical, length, map_length);
+               BUG();
+       }
+       multi->end_io = first_bio->bi_end_io;
+       multi->private = first_bio->bi_private;
+       atomic_set(&multi->stripes_pending, multi->num_stripes);
+
        while(dev_nr < total_devs) {
-               ret = btrfs_map_block(map_tree, rw, dev_nr, logical,
-                                     &physical, &map_length, &dev,
-                                     &total_devs);
-               if (map_length < length) {
-                       printk("mapping failed logical %Lu bio len %Lu physical %Lu "
-                              "len %Lu\n", logical, length, physical, map_length);
-                       BUG();
-               }
-               BUG_ON(map_length < length);
                if (total_devs > 1) {
-                       if (!multi) {
-                               multi = kmalloc(sizeof(*multi), GFP_NOFS);
-                               atomic_set(&multi->stripes, 1);
-                               multi->end_io = bio->bi_end_io;
-                               multi->private = first_bio->bi_private;
-                               multi->error = 0;
-                       } else {
-                               atomic_inc(&multi->stripes);
-                       }
                        if (dev_nr < total_devs - 1) {
                                bio = bio_clone(first_bio, GFP_NOFS);
                                BUG_ON(!bio);
                        bio->bi_private = multi;
                        bio->bi_end_io = end_bio_multi_stripe;
                }
-               bio->bi_sector = physical >> 9;
+               bio->bi_sector = multi->stripes[dev_nr].physical >> 9;
+               dev = multi->stripes[dev_nr].dev;
                bio->bi_bdev = dev->bdev;
                spin_lock(&dev->io_lock);
                dev->total_ios++;
                submit_bio(rw, bio);
                dev_nr++;
        }
+       if (total_devs == 1)
+               kfree(multi);
        return 0;
 }
 
 
 #ifndef __BTRFS_VOLUMES_
 #define __BTRFS_VOLUMES_
 
+#include <linux/bio.h>
+
 struct btrfs_device {
        struct list_head dev_list;
        struct btrfs_root *dev_root;
        struct list_head list;
 };
 
+struct btrfs_bio_stripe {
+       struct btrfs_device *dev;
+       u64 physical;
+};
+
+struct btrfs_multi_bio {
+       atomic_t stripes_pending;
+       bio_end_io_t *end_io;
+       void *private;
+       int error;
+       int num_stripes;
+       struct btrfs_bio_stripe stripes[];
+};
+
+#define btrfs_multi_bio_size(n) (sizeof(struct btrfs_multi_bio) + \
+                           (sizeof(struct btrfs_bio_stripe) * (n)))
+
 int btrfs_alloc_dev_extent(struct btrfs_trans_handle *trans,
                           struct btrfs_device *device,
                           u64 owner, u64 num_bytes, u64 *start);
-int btrfs_map_block(struct btrfs_mapping_tree *map_tree, int rw, int stripe_nr,
-                   u64 logical, u64 *phys, u64 *length,
-                   struct btrfs_device **dev, int *total_stripes);
+int btrfs_map_block(struct btrfs_mapping_tree *map_tree, int rw,
+                   u64 logical, u64 *length,
+                   struct btrfs_multi_bio **multi_ret);
 int btrfs_read_sys_array(struct btrfs_root *root);
 int btrfs_read_chunk_tree(struct btrfs_root *root);
 int btrfs_alloc_chunk(struct btrfs_trans_handle *trans,