virtio: balloon: separate out common code between remove and freeze f…

…unctions

The remove and freeze functions have a lot of shared code; put it into a
common function that gets called by both.

Signed-off-by: Amit Shah <amit.shah@redhat.com>
Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>

drivers/virtio/virtio_balloon.c

@@ -381,21 +381,25 @@ static int virtballoon_probe(struct virtio_device *vdev)
 	return err;
 }
 
-static void __devexit virtballoon_remove(struct virtio_device *vdev)
+static void remove_common(struct virtio_balloon *vb)
 {
-	struct virtio_balloon *vb = vdev->priv;
-
-	kthread_stop(vb->thread);
-
 	/* There might be pages left in the balloon: free them. */
 	while (vb->num_pages)
 		leak_balloon(vb, vb->num_pages);
 	update_balloon_size(vb);
 
 	/* Now we reset the device so we can clean up the queues. */
-	vdev->config->reset(vdev);
+	vb->vdev->config->reset(vb->vdev);
 
-	vdev->config->del_vqs(vdev);
+	vb->vdev->config->del_vqs(vb->vdev);
+}
+
+static void __devexit virtballoon_remove(struct virtio_device *vdev)
+{
+	struct virtio_balloon *vb = vdev->priv;
+
+	kthread_stop(vb->thread);
+	remove_common(vb);
 	kfree(vb);
 }
 
@@ -409,13 +413,7 @@ static int virtballoon_freeze(struct virtio_device *vdev)
 	 * function is called.
 	 */
 
-	while (vb->num_pages)
-		leak_balloon(vb, vb->num_pages);
-	update_balloon_size(vb);
-
-	/* Ensure we don't get any more requests from the host */
-	vdev->config->reset(vdev);
-	vdev->config->del_vqs(vdev);
+	remove_common(vb);
 	return 0;
 }