3 files changed, 26 insertions, 88 deletions

diff --git a/drivers/virtio/config.c b/drivers/virtio/config.c
index 983d482fba4..f70bcd2ff98 100644
--- a/drivers/virtio/config.c
+++ b/drivers/virtio/config.c
@@ -9,5 +9,4 @@
 #include <linux/virtio.h>
 #include <linux/virtio_config.h>
 #include <linux/bug.h>
-#include <asm/system.h>
 
diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 95aeedf198f..05f0a80818a 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -367,42 +367,40 @@ 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 virtballoon_thaw(struct virtio_device *vdev)
-{
-	return init_vqs(vdev->priv);
-}
-
-static int virtballoon_restore(struct virtio_device *vdev)
+static int restore_common(struct virtio_device *vdev)
 {
 	struct virtio_balloon *vb = vdev->priv;
-	struct page *page, *page2;
+	int ret;
 
-	/* We're starting from a clean slate */
-	vb->num_pages = 0;
+	ret = init_vqs(vdev->priv);
+	if (ret)
+		return ret;
 
-	/*
-	 * If a request wasn't complete at the time of freezing, this
-	 * could have been set.
-	 */
-	vb->need_stats_update = 0;
+	fill_balloon(vb, towards_target(vb));
+	update_balloon_size(vb);
+	return 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);
+static int virtballoon_restore(struct virtio_device *vdev)
+{
+	return restore_common(vdev);
 }
 #endif
 
@@ -423,7 +421,6 @@ static struct virtio_driver virtio_balloon_driver = {
 #ifdef CONFIG_PM
 	.freeze	=	virtballoon_freeze,
 	.restore =	virtballoon_restore,
-	.thaw =		virtballoon_thaw,
 #endif
 };
 
diff --git a/drivers/virtio/virtio_pci.c b/drivers/virtio/virtio_pci.c
index 635e1efb379..2e03d416b9a 100644
--- a/drivers/virtio/virtio_pci.c
+++ b/drivers/virtio/virtio_pci.c
@@ -720,24 +720,6 @@ static void __devexit virtio_pci_remove(struct pci_dev *pci_dev)
 }
 
 #ifdef CONFIG_PM
-static int virtio_pci_suspend(struct device *dev)
-{
-	struct pci_dev *pci_dev = to_pci_dev(dev);
-
-	pci_save_state(pci_dev);
-	pci_set_power_state(pci_dev, PCI_D3hot);
-	return 0;
-}
-
-static int virtio_pci_resume(struct device *dev)
-{
-	struct pci_dev *pci_dev = to_pci_dev(dev);
-
-	pci_restore_state(pci_dev);
-	pci_set_power_state(pci_dev, PCI_D0);
-	return 0;
-}
-
 static int virtio_pci_freeze(struct device *dev)
 {
 	struct pci_dev *pci_dev = to_pci_dev(dev);
@@ -758,59 +740,24 @@ static int virtio_pci_freeze(struct device *dev)
 	return ret;
 }
 
-static int restore_common(struct device *dev)
-{
-	struct pci_dev *pci_dev = to_pci_dev(dev);
-	struct virtio_pci_device *vp_dev = pci_get_drvdata(pci_dev);
-	int ret;
-
-	ret = pci_enable_device(pci_dev);
-	if (ret)
-		return ret;
-	pci_set_master(pci_dev);
-	vp_finalize_features(&vp_dev->vdev);
-
-	return ret;
-}
-
-static int virtio_pci_thaw(struct device *dev)
+static int virtio_pci_restore(struct device *dev)
 {
 	struct pci_dev *pci_dev = to_pci_dev(dev);
 	struct virtio_pci_device *vp_dev = pci_get_drvdata(pci_dev);
 	struct virtio_driver *drv;
 	int ret;
 
-	ret = restore_common(dev);
-	if (ret)
-		return ret;
-
 	drv = container_of(vp_dev->vdev.dev.driver,
 			   struct virtio_driver, driver);
 
-	if (drv && drv->thaw)
-		ret = drv->thaw(&vp_dev->vdev);
-	else if (drv && drv->restore)
-		ret = drv->restore(&vp_dev->vdev);
-
-	/* Finally, tell the device we're all set */
-	if (!ret)
-		vp_set_status(&vp_dev->vdev, vp_dev->saved_status);
-
-	return ret;
-}
-
-static int virtio_pci_restore(struct device *dev)
-{
-	struct pci_dev *pci_dev = to_pci_dev(dev);
-	struct virtio_pci_device *vp_dev = pci_get_drvdata(pci_dev);
-	struct virtio_driver *drv;
-	int ret;
+	ret = pci_enable_device(pci_dev);
+	if (ret)
+		return ret;
 
-	drv = container_of(vp_dev->vdev.dev.driver,
-			   struct virtio_driver, driver);
+	pci_set_master(pci_dev);
+	vp_finalize_features(&vp_dev->vdev);
 
-	ret = restore_common(dev);
-	if (!ret && drv && drv->restore)
+	if (drv && drv->restore)
 		ret = drv->restore(&vp_dev->vdev);
 
 	/* Finally, tell the device we're all set */
@@ -821,12 +768,7 @@ static int virtio_pci_restore(struct device *dev)
 }
 
 static const struct dev_pm_ops virtio_pci_pm_ops = {
-	.suspend	= virtio_pci_suspend,
-	.resume		= virtio_pci_resume,
-	.freeze		= virtio_pci_freeze,
-	.thaw		= virtio_pci_thaw,
-	.restore	= virtio_pci_restore,
-	.poweroff	= virtio_pci_suspend,
+	SET_SYSTEM_SLEEP_PM_OPS(virtio_pci_freeze, virtio_pci_restore)
 };
 #endif