#define CREATE_TRACE_POINTS
 #include "null_blk_trace.h"
 
-/* zone_size in MBs to sectors. */
-#define ZONE_SIZE_SHIFT                11
+#define MB_TO_SECTS(mb) (((sector_t)mb * SZ_1M) >> SECTOR_SHIFT)
 
 static inline unsigned int null_zone_no(struct nullb_device *dev, sector_t sect)
 {
 
 int null_init_zoned_dev(struct nullb_device *dev, struct request_queue *q)
 {
-       sector_t dev_size = (sector_t)dev->size * 1024 * 1024;
+       sector_t dev_capacity_sects, zone_capacity_sects;
        sector_t sector = 0;
        unsigned int i;
 
                return -EINVAL;
        }
 
-       dev->zone_size_sects = dev->zone_size << ZONE_SIZE_SHIFT;
-       dev->nr_zones = dev_size >>
-                               (SECTOR_SHIFT + ilog2(dev->zone_size_sects));
+       zone_capacity_sects = MB_TO_SECTS(dev->zone_capacity);
+       dev_capacity_sects = MB_TO_SECTS(dev->size);
+       dev->zone_size_sects = MB_TO_SECTS(dev->zone_size);
+       dev->nr_zones = dev_capacity_sects >> ilog2(dev->zone_size_sects);
+       if (dev_capacity_sects & (dev->zone_size_sects - 1))
+               dev->nr_zones++;
+
        dev->zones = kvmalloc_array(dev->nr_zones, sizeof(struct blk_zone),
                        GFP_KERNEL | __GFP_ZERO);
        if (!dev->zones)
                struct blk_zone *zone = &dev->zones[i];
 
                zone->start = zone->wp = sector;
-               zone->len = dev->zone_size_sects;
-               zone->capacity = dev->zone_capacity << ZONE_SIZE_SHIFT;
+               if (zone->start + dev->zone_size_sects > dev_capacity_sects)
+                       zone->len = dev_capacity_sects - zone->start;
+               else
+                       zone->len = dev->zone_size_sects;
+               zone->capacity =
+                       min_t(sector_t, zone->len, zone_capacity_sects);
                zone->type = BLK_ZONE_TYPE_SEQWRITE_REQ;
                zone->cond = BLK_ZONE_COND_EMPTY;