Staging: hv: blkvsc: Fix bugs in the module unload path

Fix bugs in the module unload path for the blkvsc driver.
Signed-off-by: K. Y. Srinivasan <kys@microsoft.com>
Signed-off-by: Haiyang Zhang <haiyangz@microsoft.com>
Signed-off-by: Abhishek Kane <v-abkane@microsoft.com>
Signed-off-by: Greg Kroah-Hartman <gregkh@suse.de>

drivers/staging/hv/blkvsc_drv.c

@@ -518,22 +518,18 @@ static int blkvsc_remove(struct hv_device *dev)
 
 	blkvsc_do_operation(blkdev, DO_FLUSH);
 
-	blk_cleanup_queue(blkdev->gd->queue);
+	if (blkdev->users == 0) {
+		del_gendisk(blkdev->gd);
+		put_disk(blkdev->gd);
+		blk_cleanup_queue(blkdev->gd->queue);
 
-	/*
-	 * Call to the vsc driver to let it know that the device is being
-	 * removed
-	 */
-	storvsc_dev_remove(dev);
-
-	del_gendisk(blkdev->gd);
+		storvsc_dev_remove(blkdev->device_ctx);
 
-	kmem_cache_destroy(blkdev->request_pool);
-
-	kfree(blkdev);
+		kmem_cache_destroy(blkdev->request_pool);
+		kfree(blkdev);
+	}
 
 	return 0;
-
 }
 
 static void blkvsc_shutdown(struct hv_device *dev)
@@ -568,13 +564,23 @@ static int blkvsc_release(struct gendisk *disk, fmode_t mode)
 	struct block_device_context *blkdev = disk->private_data;
 	unsigned long flags;
 
-	if (blkdev->users == 1) {
+	spin_lock_irqsave(&blkdev->lock, flags);
+
+	if ((--blkdev->users == 0) && (blkdev->shutting_down)) {
+		blk_stop_queue(blkdev->gd->queue);
+		spin_unlock_irqrestore(&blkdev->lock, flags);
+
 		blkvsc_do_operation(blkdev, DO_FLUSH);
-	}
+		del_gendisk(blkdev->gd);
+		put_disk(blkdev->gd);
+		blk_cleanup_queue(blkdev->gd->queue);
 
-	spin_lock_irqsave(&blkdev->lock, flags);
-	blkdev->users--;
-	spin_unlock_irqrestore(&blkdev->lock, flags);
+		storvsc_dev_remove(blkdev->device_ctx);
+
+		kmem_cache_destroy(blkdev->request_pool);
+		kfree(blkdev);
+	} else
+		spin_unlock_irqrestore(&blkdev->lock, flags);
 
 	return 0;
 }