2 files changed, 26 insertions, 15 deletions

diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 95aeedf198f..958e5129c60 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -367,29 +367,45 @@ static void __devexit virtballoon_remove(struct virtio_device *vdev)
 #ifdef CONFIG_PM
 static int virtballoon_freeze(struct virtio_device *vdev)
 {
+	struct virtio_balloon *vb = vdev->priv;
+
 	/*
 	 * The kthread is already frozen by the PM core before this
 	 * 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);
 	return 0;
 }
 
+static int restore_common(struct virtio_device *vdev)
+{
+	struct virtio_balloon *vb = vdev->priv;
+	int ret;
+
+	ret = init_vqs(vdev->priv);
+	if (ret)
+		return ret;
+
+	fill_balloon(vb, towards_target(vb));
+	update_balloon_size(vb);
+	return 0;
+}
+
 static int virtballoon_thaw(struct virtio_device *vdev)
 {
-	return init_vqs(vdev->priv);
+	return restore_common(vdev);
 }
 
 static int virtballoon_restore(struct virtio_device *vdev)
 {
 	struct virtio_balloon *vb = vdev->priv;
-	struct page *page, *page2;
-
-	/* We're starting from a clean slate */
-	vb->num_pages = 0;
 
 	/*
 	 * If a request wasn't complete at the time of freezing, this
@@ -397,12 +413,7 @@ static int virtballoon_restore(struct virtio_device *vdev)
 	 */
 	vb->need_stats_update = 0;
 
-	/* We don't have these pages in the balloon anymore! */
-	list_for_each_entry_safe(page, page2, &vb->pages, lru) {
-		list_del(&page->lru);
-		totalram_pages++;
-	}
-	return init_vqs(vdev->priv);
+	return restore_common(vdev);
 }
 #endif
 
diff --git a/drivers/virtio/virtio_ring.c b/drivers/virtio/virtio_ring.c
index 79e1b292c03..5aa43c3392a 100644
--- a/drivers/virtio/virtio_ring.c
+++ b/drivers/virtio/virtio_ring.c
@@ -35,7 +35,7 @@
 #define virtio_rmb(vq) \
 	do { if ((vq)->weak_barriers) smp_rmb(); else rmb(); } while(0)
 #define virtio_wmb(vq) \
-	do { if ((vq)->weak_barriers) smp_rmb(); else rmb(); } while(0)
+	do { if ((vq)->weak_barriers) smp_wmb(); else wmb(); } while(0)
 #else
 /* We must force memory ordering even if guest is UP since host could be
  * running on another CPU, but SMP barriers are defined to barrier() in that
@@ -308,9 +308,9 @@ bool virtqueue_kick_prepare(struct virtqueue *_vq)
 	bool needs_kick;
 
 	START_USE(vq);
-	/* Descriptors and available array need to be set before we expose the
-	 * new available array entries. */
-	virtio_wmb(vq);
+	/* We need to expose available array entries before checking avail
+	 * event. */
+	virtio_mb(vq);
 
 	old = vq->vring.avail->idx - vq->num_added;
 	new = vq->vring.avail->idx;