diff --git a/drivers/scsi/sd_zbc.c b/drivers/scsi/sd_zbc.c
index bc3cb81a9c7d..ea9e1e0ed5b8 100644
--- a/drivers/scsi/sd_zbc.c
+++ b/drivers/scsi/sd_zbc.c
@@ -423,7 +423,16 @@ static int sd_zbc_check_capacity(struct scsi_disk *sdkp,
 
 #define SD_ZBC_BUF_SIZE 131072
 
-static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
+/**
+ * sd_zbc_check_zone_size - Check the device zone sizes
+ * @sdkp: Target disk
+ *
+ * Check that all zones of the device are equal. The last zone can however
+ * be smaller. The zone size must also be a power of two number of LBAs.
+ *
+ * Returns the zone size in bytes upon success or an error code upon failure.
+ */
+static s64 sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 {
 	u64 zone_blocks = 0;
 	sector_t block = 0;
@@ -434,8 +443,6 @@ static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 	int ret;
 	u8 same;
 
-	sdkp->zone_blocks = 0;
-
 	/* Get a buffer */
 	buf = kmalloc(SD_ZBC_BUF_SIZE, GFP_KERNEL);
 	if (!buf)
@@ -470,16 +477,17 @@ static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 
 		/* Parse zone descriptors */
 		while (rec < buf + buf_len) {
-			zone_blocks = get_unaligned_be64(&rec[8]);
-			if (sdkp->zone_blocks == 0) {
-				sdkp->zone_blocks = zone_blocks;
-			} else if (zone_blocks != sdkp->zone_blocks &&
-				   (block + zone_blocks < sdkp->capacity
-				    || zone_blocks > sdkp->zone_blocks)) {
+			u64 this_zone_blocks = get_unaligned_be64(&rec[8]);
+
+			if (zone_blocks == 0) {
+				zone_blocks = this_zone_blocks;
+			} else if (this_zone_blocks != zone_blocks &&
+				   (block + this_zone_blocks < sdkp->capacity
+				    || this_zone_blocks > zone_blocks)) {
 				zone_blocks = 0;
 				goto out;
 			}
-			block += zone_blocks;
+			block += this_zone_blocks;
 			rec += 64;
 		}
 
@@ -492,8 +500,6 @@ static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 
 	} while (block < sdkp->capacity);
 
-	zone_blocks = sdkp->zone_blocks;
-
 out:
 	if (!zone_blocks) {
 		if (sdkp->first_scan)
@@ -513,8 +519,7 @@ static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 				  "Zone size too large\n");
 		ret = -ENODEV;
 	} else {
-		sdkp->zone_blocks = zone_blocks;
-		sdkp->zone_shift = ilog2(zone_blocks);
+		ret = zone_blocks;
 	}
 
 out_free:
@@ -523,23 +528,44 @@ static int sd_zbc_check_zone_size(struct scsi_disk *sdkp)
 	return ret;
 }
 
-static int sd_zbc_setup(struct scsi_disk *sdkp)
+static int sd_zbc_setup(struct scsi_disk *sdkp, u32 zone_blocks)
 {
+	struct request_queue *q = sdkp->disk->queue;
+	u32 zone_shift = ilog2(zone_blocks);
+	u32 nr_zones;
 
 	/* chunk_sectors indicates the zone size */
-	blk_queue_chunk_sectors(sdkp->disk->queue,
-			logical_to_sectors(sdkp->device, sdkp->zone_blocks));
-	sdkp->zone_shift = ilog2(sdkp->zone_blocks);
-	sdkp->nr_zones = sdkp->capacity >> sdkp->zone_shift;
-	if (sdkp->capacity & (sdkp->zone_blocks - 1))
-		sdkp->nr_zones++;
-
-	if (!sdkp->zones_wlock) {
-		sdkp->zones_wlock = kcalloc(BITS_TO_LONGS(sdkp->nr_zones),
-					    sizeof(unsigned long),
-					    GFP_KERNEL);
-		if (!sdkp->zones_wlock)
-			return -ENOMEM;
+	blk_queue_chunk_sectors(q,
+			logical_to_sectors(sdkp->device, zone_blocks));
+	nr_zones = round_up(sdkp->capacity, zone_blocks) >> zone_shift;
+
+	/*
+	 * Initialize the disk zone write lock bitmap if the number
+	 * of zones changed.
+	 */
+	if (nr_zones != sdkp->nr_zones) {
+		unsigned long *zones_wlock = NULL;
+
+		if (nr_zones) {
+			zones_wlock = kcalloc(BITS_TO_LONGS(nr_zones),
+					      sizeof(unsigned long),
+					      GFP_KERNEL);
+			if (!zones_wlock)
+				return -ENOMEM;
+		}
+
+		blk_mq_freeze_queue(q);
+		sdkp->zone_blocks = zone_blocks;
+		sdkp->zone_shift = zone_shift;
+		sdkp->nr_zones = nr_zones;
+		swap(sdkp->zones_wlock, zones_wlock);
+		blk_mq_unfreeze_queue(q);
+
+		kfree(zones_wlock);
+
+		/* READ16/WRITE16 is mandatory for ZBC disks */
+		sdkp->device->use_16_for_rw = 1;
+		sdkp->device->use_10_for_rw = 0;
 	}
 
 	return 0;
@@ -548,6 +574,7 @@ static int sd_zbc_setup(struct scsi_disk *sdkp)
 int sd_zbc_read_zones(struct scsi_disk *sdkp,
 		      unsigned char *buf)
 {
+	int64_t zone_blocks;
 	int ret;
 
 	if (!sd_is_zoned(sdkp))
@@ -585,19 +612,19 @@ int sd_zbc_read_zones(struct scsi_disk *sdkp,
 	 * Check zone size: only devices with a constant zone size (except
 	 * an eventual last runt zone) that is a power of 2 are supported.
 	 */
-	ret = sd_zbc_check_zone_size(sdkp);
-	if (ret)
+	zone_blocks = sd_zbc_check_zone_size(sdkp);
+	ret = -EFBIG;
+	if (zone_blocks != (u32)zone_blocks)
+		goto err;
+	ret = zone_blocks;
+	if (ret < 0)
 		goto err;
 
 	/* The drive satisfies the kernel restrictions: set it up */
-	ret = sd_zbc_setup(sdkp);
+	ret = sd_zbc_setup(sdkp, zone_blocks);
 	if (ret)
 		goto err;
 
-	/* READ16/WRITE16 is mandatory for ZBC disks */
-	sdkp->device->use_16_for_rw = 1;
-	sdkp->device->use_10_for_rw = 0;
-
 	return 0;
 
 err:
@@ -610,6 +637,7 @@ void sd_zbc_remove(struct scsi_disk *sdkp)
 {
 	kfree(sdkp->zones_wlock);
 	sdkp->zones_wlock = NULL;
+	sdkp->nr_zones = 0;
 }
 
 void sd_zbc_print_zones(struct scsi_disk *sdkp)
