return disk_to_dev(nbd->disk);
 }
 
+static bool nbd_is_connected(struct nbd_device *nbd)
+{
+       return !!nbd->task_recv;
+}
+
 static const char *nbdcmd_to_ascii(int cmd)
 {
        switch (cmd) {
        return "invalid";
 }
 
+static int nbd_size_clear(struct nbd_device *nbd, struct block_device *bdev)
+{
+       bdev->bd_inode->i_size = 0;
+       set_capacity(nbd->disk, 0);
+       kobject_uevent(&nbd_to_dev(nbd)->kobj, KOBJ_CHANGE);
+
+       return 0;
+}
+
+static void nbd_size_update(struct nbd_device *nbd, struct block_device *bdev)
+{
+       if (!nbd_is_connected(nbd))
+               return;
+
+       bdev->bd_inode->i_size = nbd->bytesize;
+       set_capacity(nbd->disk, nbd->bytesize >> 9);
+       kobject_uevent(&nbd_to_dev(nbd)->kobj, KOBJ_CHANGE);
+}
+
+static int nbd_size_set(struct nbd_device *nbd, struct block_device *bdev,
+                       int blocksize, int nr_blocks)
+{
+       int ret;
+
+       ret = set_blocksize(bdev, blocksize);
+       if (ret)
+               return ret;
+
+       nbd->blksize = blocksize;
+       nbd->bytesize = (loff_t)blocksize * (loff_t)nr_blocks;
+
+       nbd_size_update(nbd, bdev);
+
+       return 0;
+}
+
 static void nbd_end_request(struct nbd_device *nbd, struct request *req)
 {
        int error = req->errors ? -EIO : 0;
        .show = pid_show,
 };
 
-static int nbd_thread_recv(struct nbd_device *nbd)
+static int nbd_thread_recv(struct nbd_device *nbd, struct block_device *bdev)
 {
        struct request *req;
        int ret;
                return ret;
        }
 
+       nbd_size_update(nbd, bdev);
+
        while (1) {
                req = nbd_read_stat(nbd);
                if (IS_ERR(req)) {
                nbd_end_request(nbd, req);
        }
 
+       nbd_size_clear(nbd, bdev);
+
        device_remove_file(disk_to_dev(nbd->disk), &pid_attr);
 
        nbd->task_recv = NULL;
                return err;
        }
 
-       case NBD_SET_BLKSIZE:
-               nbd->blksize = arg;
-               nbd->bytesize &= ~(nbd->blksize-1);
-               bdev->bd_inode->i_size = nbd->bytesize;
-               set_blocksize(bdev, nbd->blksize);
-               set_capacity(nbd->disk, nbd->bytesize >> 9);
-               return 0;
+       case NBD_SET_BLKSIZE: {
+               loff_t bsize = nbd->bytesize;
+               do_div(bsize, arg);
+
+               return nbd_size_set(nbd, bdev, arg, bsize);
+       }
 
        case NBD_SET_SIZE:
-               nbd->bytesize = arg & ~(nbd->blksize-1);
-               bdev->bd_inode->i_size = nbd->bytesize;
-               set_blocksize(bdev, nbd->blksize);
-               set_capacity(nbd->disk, nbd->bytesize >> 9);
-               return 0;
+               return nbd_size_set(nbd, bdev, nbd->blksize,
+                                   arg / nbd->blksize);
+
+       case NBD_SET_SIZE_BLOCKS:
+               return nbd_size_set(nbd, bdev, nbd->blksize, arg);
 
        case NBD_SET_TIMEOUT:
                nbd->xmit_timeout = arg * HZ;
                nbd->flags = arg;
                return 0;
 
-       case NBD_SET_SIZE_BLOCKS:
-               nbd->bytesize = ((u64) arg) * nbd->blksize;
-               bdev->bd_inode->i_size = nbd->bytesize;
-               set_blocksize(bdev, nbd->blksize);
-               set_capacity(nbd->disk, nbd->bytesize >> 9);
-               return 0;
-
        case NBD_DO_IT: {
                struct task_struct *thread;
                int error;
                }
 
                nbd_dev_dbg_init(nbd);
-               error = nbd_thread_recv(nbd);
+               error = nbd_thread_recv(nbd, bdev);
                nbd_dev_dbg_close(nbd);
                kthread_stop(thread);