diff --git a/include/migration/migration.h b/include/migration/migration.h
index 8334621..a4117d6 100644
--- a/include/migration/migration.h
+++ b/include/migration/migration.h
@@ -114,6 +114,10 @@ void migrate_fd_connect(MigrationState *s);
 
 int migrate_fd_close(MigrationState *s);
 
+void migration_end(void);
+
+void blk_mig_cleanup(void);
+
 void add_migration_state_change_notifier(Notifier *notify);
 void remove_migration_state_change_notifier(Notifier *notify);
 bool migration_in_setup(MigrationState *);
diff --git a/migration/block.c b/migration/block.c
index f7bb1e0..267953f 100644
--- a/migration/block.c
+++ b/migration/block.c
@@ -591,7 +591,7 @@ static int64_t get_remaining_dirty(void)
 
 /* Called with iothread lock taken.  */
 
-static void blk_mig_cleanup(void)
+void blk_mig_cleanup(void)
 {
     BlkMigDevState *bmds;
     BlkMigBlock *blk;
@@ -751,6 +751,7 @@ static int block_save_complete(QEMUFile *f, void *opaque)
     qemu_put_be64(f, BLK_MIG_FLAG_EOS);
 
     blk_mig_cleanup();
+    //blk_mig_cleanup();
     return 0;
 }
 
diff --git a/migration/migration.c b/migration/migration.c
index b092f38..caad0e3 100644
--- a/migration/migration.c
+++ b/migration/migration.c
@@ -613,12 +613,9 @@ static void migrate_fd_cleanup(void *opaque)
 
     assert(s->state != MIGRATION_STATUS_ACTIVE);
 
-    if (s->state != MIGRATION_STATUS_COMPLETED) {
-        qemu_savevm_state_cancel();
-        if (s->state == MIGRATION_STATUS_CANCELLING) {
-            migrate_set_state(s, MIGRATION_STATUS_CANCELLING,
-                              MIGRATION_STATUS_CANCELLED);
-        }
+    if (s->state == MIGRATION_STATUS_CANCELLING) {
+        migrate_set_state(s, MIGRATION_STATUS_CANCELLING,
+                          MIGRATION_STATUS_CANCELLED);
     }
 
     notifier_list_notify(&migration_state_notifiers, s);
@@ -1028,6 +1025,7 @@ static void *migration_thread(void *opaque)
     int64_t initial_bytes = 0;
     int64_t max_size = 0;
     int64_t start_time = initial_time;
+    int64_t end_time;
     bool old_vm_running = false;
 
     rcu_register_thread();
@@ -1089,10 +1087,12 @@ static void *migration_thread(void *opaque)
 
     /* If we enabled cpu throttling for auto-converge, turn it off. */
     cpu_throttle_stop();
+    end_time = qemu_clock_get_ms(QEMU_CLOCK_REALTIME);
 
     qemu_mutex_lock_iothread();
+
+    migration_end();
     if (s->state == MIGRATION_STATUS_COMPLETED) {
-        int64_t end_time = qemu_clock_get_ms(QEMU_CLOCK_REALTIME);
         uint64_t transferred_bytes = qemu_ftell(s->file);
         s->total_time = end_time - s->total_time;
         s->downtime = end_time - start_time;
diff --git a/migration/ram.c b/migration/ram.c
index a25bcc7..23cc30c 100644
--- a/migration/ram.c
+++ b/migration/ram.c
@@ -1100,7 +1100,7 @@ static void migration_bitmap_free(struct BitmapRcu *bmap)
     g_free(bmap);
 }
 
-static void migration_end(void)
+void migration_end(void)
 {
     /* caller have hold iothread lock or is in a bh, so there is
      * no writing race against this migration_bitmap
@@ -1344,7 +1344,7 @@ static int ram_save_complete(QEMUFile *f, void *opaque)
 
     rcu_read_unlock();
 
-    migration_end();
+    //migration_end();
     qemu_put_be64(f, RAM_SAVE_FLAG_EOS);
 
     return 0;