1 files changed, 5 insertions, 14 deletions

diff --git a/hw/virtio/virtio-balloon.c b/hw/virtio/virtio-balloon.c
index 1a22e6d993..5af429a58a 100644
--- a/hw/virtio/virtio-balloon.c
+++ b/hw/virtio/virtio-balloon.c
@@ -396,11 +396,6 @@ static void virtio_balloon_to_target(void *opaque, ram_addr_t target)
     trace_virtio_balloon_to_target(target, dev->num_pages);
 }
 
-static void virtio_balloon_save(QEMUFile *f, void *opaque)
-{
-    virtio_save(VIRTIO_DEVICE(opaque), f);
-}
-
 static void virtio_balloon_save_device(VirtIODevice *vdev, QEMUFile *f)
 {
     VirtIOBalloon *s = VIRTIO_BALLOON(vdev);
@@ -409,12 +404,9 @@ static void virtio_balloon_save_device(VirtIODevice *vdev, QEMUFile *f)
     qemu_put_be32(f, s->actual);
 }
 
-static int virtio_balloon_load(QEMUFile *f, void *opaque, int version_id)
+static int virtio_balloon_load(QEMUFile *f, void *opaque, size_t size)
 {
-    if (version_id != 1)
-        return -EINVAL;
-
-    return virtio_load(VIRTIO_DEVICE(opaque), f, version_id);
+    return virtio_load(VIRTIO_DEVICE(opaque), f, 1);
 }
 
 static int virtio_balloon_load_device(VirtIODevice *vdev, QEMUFile *f,
@@ -454,9 +446,6 @@ static void virtio_balloon_device_realize(DeviceState *dev, Error **errp)
     s->svq = virtio_add_queue(vdev, 128, virtio_balloon_receive_stats);
 
     reset_stats(s);
-
-    register_savevm(dev, "virtio-balloon", -1, 1,
-                    virtio_balloon_save, virtio_balloon_load, s);
 }
 
 static void virtio_balloon_device_unrealize(DeviceState *dev, Error **errp)
@@ -466,7 +455,6 @@ static void virtio_balloon_device_unrealize(DeviceState *dev, Error **errp)
 
     balloon_stats_destroy_timer(s);
     qemu_remove_balloon_handler(s);
-    unregister_savevm(dev, "virtio-balloon", s);
     virtio_cleanup(vdev);
 }
 
@@ -493,6 +481,8 @@ static void virtio_balloon_instance_init(Object *obj)
                         NULL, s, NULL);
 }
 
+VMSTATE_VIRTIO_DEVICE(balloon, 1, virtio_balloon_load, virtio_vmstate_save);
+
 static Property virtio_balloon_properties[] = {
     DEFINE_PROP_BIT("deflate-on-oom", VirtIOBalloon, host_features,
                     VIRTIO_BALLOON_F_DEFLATE_ON_OOM, false),
@@ -505,6 +495,7 @@ static void virtio_balloon_class_init(ObjectClass *klass, void *data)
     VirtioDeviceClass *vdc = VIRTIO_DEVICE_CLASS(klass);
 
     dc->props = virtio_balloon_properties;
+    dc->vmsd = &vmstate_virtio_balloon;
     set_bit(DEVICE_CATEGORY_MISC, dc->categories);
     vdc->realize = virtio_balloon_device_realize;
     vdc->unrealize = virtio_balloon_device_unrealize;