4 files changed, 15 insertions, 11 deletions
diff --git a/arch_init.c b/arch_init.c
index 606b848474..f6bc3fe69c 100644
--- a/arch_init.c
+++ b/arch_init.c
@@ -296,6 +296,11 @@ static void migration_end(void)
     memory_global_dirty_log_stop();
 }
 
+static void ram_migration_cancel(void *opaque)
+{
+    migration_end();
+}
+
 #define MAX_WAIT 50 /* ms, half buffered_file limit */
 
 static int ram_save_live(QEMUFile *f, int stage, void *opaque)
@@ -306,11 +311,6 @@ static int ram_save_live(QEMUFile *f, int stage, void *opaque)
     int ret;
     int i;
 
-    if (stage < 0) {
-        migration_end();
-        return 0;
-    }
-
     memory_global_sync_dirty_bitmap(get_system_memory());
 
     if (stage == 1) {
@@ -537,6 +537,7 @@ done:
 SaveVMHandlers savevm_ram_handlers = {
     .save_live_state = ram_save_live,
     .load_state = ram_load,
+    .cancel = ram_migration_cancel,
 };
 
 #ifdef HAS_AUDIO
diff --git a/block-migration.c b/block-migration.c
index 00151a07d9..cd8a8dd6cb 100644
--- a/block-migration.c
+++ b/block-migration.c
@@ -536,6 +536,11 @@ static void blk_mig_cleanup(void)
     }
 }
 
+static void block_migration_cancel(void *opaque)
+{
+    blk_mig_cleanup();
+}
+
 static int block_save_live(QEMUFile *f, int stage, void *opaque)
 {
     int ret;
@@ -543,10 +548,6 @@ static int block_save_live(QEMUFile *f, int stage, void *opaque)
     DPRINTF("Enter save live stage %d submitted %d transferred %d\n",
             stage, block_mig_state.submitted, block_mig_state.transferred);
 
-    if (stage < 0) {
-        blk_mig_cleanup();
-        return 0;
-    }
 
     if (block_mig_state.blk_enable != 1) {
         /* no need to migrate storage */
@@ -713,6 +714,7 @@ SaveVMHandlers savevm_block_handlers = {
     .set_params = block_set_params,
     .save_live_state = block_save_live,
     .load_state = block_load,
+    .cancel = block_migration_cancel,
 };
 
 void blk_mig_init(void)
diff --git a/savevm.c b/savevm.c
index a451be22a6..888c5a2faf 100644
--- a/savevm.c
+++ b/savevm.c
@@ -1703,8 +1703,8 @@ void qemu_savevm_state_cancel(QEMUFile *f)
     SaveStateEntry *se;
 
     QTAILQ_FOREACH(se, &savevm_handlers, entry) {
-        if (se->ops && se->ops->save_live_state) {
-            se->ops->save_live_state(f, -1, se->opaque);
+        if (se->ops && se->ops->cancel) {
+            se->ops->cancel(se->opaque);
         }
     }
 }
diff --git a/vmstate.h b/vmstate.h
index 0e248344fd..1dd42f516a 100644
--- a/vmstate.h
+++ b/vmstate.h
@@ -33,6 +33,7 @@ typedef struct SaveVMHandlers {
     void (*set_params)(const MigrationParams *params, void * opaque);
     SaveStateHandler *save_state;
     int (*save_live_state)(QEMUFile *f, int stage, void *opaque);
+    void (*cancel)(void *opaque);
     LoadStateHandler *load_state;
 } SaveVMHandlers;