diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 90a45ef7203bd..bee3809e85e93 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -571,6 +571,9 @@ struct kvm_memory_slot {
 	gfn_t base_gfn;
 	unsigned long npages;
 	unsigned long *dirty_bitmap;
+	/* Protects use_secondary, which tells which bitmap to use */
+	spinlock_t bm_switch_lock;
+	bool use_secondary;
 	struct kvm_arch_memory_slot arch;
 	unsigned long userspace_addr;
 	u32 flags;
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index a49df8988cd6a..5ea543d8cfec2 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -1560,6 +1560,9 @@ static int kvm_prepare_memory_region(struct kvm *kvm,
 
 	r = kvm_arch_prepare_memory_region(kvm, old, new, change);
 
+	if (new)
+		spin_lock_init(&new->bm_switch_lock);
+
 	/* Free the bitmap on failure if it was allocated above. */
 	if (r && new && new->dirty_bitmap && (!old || !old->dirty_bitmap))
 		kvm_destroy_dirty_bitmap(new);
@@ -2053,7 +2056,6 @@ static int kvm_get_dirty_log_protect(struct kvm *kvm, struct kvm_dirty_log *log)
 	int i, as_id, id;
 	unsigned long n;
 	unsigned long *dirty_bitmap;
-	unsigned long *dirty_bitmap_buffer;
 	bool flush;
 
 	/* Dirty ring tracking is exclusive to dirty log tracking */
@@ -2070,12 +2072,11 @@ static int kvm_get_dirty_log_protect(struct kvm *kvm, struct kvm_dirty_log *log)
 	if (!memslot || !memslot->dirty_bitmap)
 		return -ENOENT;
 
-	dirty_bitmap = memslot->dirty_bitmap;
-
 	kvm_arch_sync_dirty_log(kvm, memslot);
 
 	n = kvm_dirty_bitmap_bytes(memslot);
 	flush = false;
+
 	if (kvm->manual_dirty_log_protect) {
 		/*
 		 * Unlike kvm_get_dirty_log, we always return false in *flush,
@@ -2085,35 +2086,49 @@ static int kvm_get_dirty_log_protect(struct kvm *kvm, struct kvm_dirty_log *log)
 		 * transition to kvm_get_dirty_log_protect and kvm_get_dirty_log
 		 * can be eliminated.
 		 */
-		dirty_bitmap_buffer = dirty_bitmap;
-	} else {
-		dirty_bitmap_buffer = kvm_second_dirty_bitmap(memslot);
-		memset(dirty_bitmap_buffer, 0, n);
+		if (copy_to_user(log->dirty_bitmap, memslot->dirty_bitmap, n))
+			return -EFAULT;
 
-		KVM_MMU_LOCK(kvm);
-		for (i = 0; i < n / sizeof(long); i++) {
-			unsigned long mask;
-			gfn_t offset;
+		return 0;
+	}
 
-			if (!dirty_bitmap[i])
-				continue;
+	/*
+	 * Switches between primary and secondary dirty_bitmap.
+	 * After unlocking, all new dirty bits are written to an empty bitmap.
+	 */
+	spin_lock(&memslot->bm_switch_lock);
+	if (memslot->use_secondary)
+		dirty_bitmap = kvm_second_dirty_bitmap(memslot);
+	else
+		dirty_bitmap = memslot->dirty_bitmap;
 
-			flush = true;
-			mask = xchg(&dirty_bitmap[i], 0);
-			dirty_bitmap_buffer[i] = mask;
+	memslot->use_secondary = !memslot->use_secondary;
+	spin_unlock(&memslot->bm_switch_lock);
 
-			offset = i * BITS_PER_LONG;
-			kvm_arch_mmu_enable_log_dirty_pt_masked(kvm, memslot,
-								offset, mask);
-		}
-		KVM_MMU_UNLOCK(kvm);
+	if (copy_to_user(log->dirty_bitmap, dirty_bitmap, n))
+		return -EFAULT;
+
+	KVM_MMU_LOCK(kvm);
+	for (i = 0; i < n / sizeof(long); i++) {
+		unsigned long mask;
+		gfn_t offset;
+
+		if (!dirty_bitmap[i])
+			continue;
+
+		flush = true;
+		mask = dirty_bitmap[i];
+		dirty_bitmap[i] = 0;
+
+		offset = i * BITS_PER_LONG;
+		kvm_arch_mmu_enable_log_dirty_pt_masked(kvm, memslot,
+							offset, mask);
 	}
+	KVM_MMU_UNLOCK(kvm);
 
 	if (flush)
 		kvm_arch_flush_remote_tlbs_memslot(kvm, memslot);
 
-	if (copy_to_user(log->dirty_bitmap, dirty_bitmap_buffer, n))
-		return -EFAULT;
 	return 0;
 }
 
@@ -3203,11 +3218,19 @@ void mark_page_dirty_in_slot(struct kvm *kvm,
 		unsigned long rel_gfn = gfn - memslot->base_gfn;
 		u32 slot = (memslot->as_id << 16) | memslot->id;
 
-		if (kvm->dirty_ring_size)
+		if (kvm->dirty_ring_size) {
 			kvm_dirty_ring_push(&vcpu->dirty_ring,
 					    slot, rel_gfn);
-		else
-			set_bit_le(rel_gfn, memslot->dirty_bitmap);
+		} else {
+			spin_lock(&memslot->bm_switch_lock);
+			if (memslot->use_secondary)
+				set_bit_le(rel_gfn,
+					   kvm_second_dirty_bitmap(memslot));
+			else
+				set_bit_le(rel_gfn, memslot->dirty_bitmap);
+
+			spin_unlock(&memslot->bm_switch_lock);
+		}
 	}
 }
 EXPORT_SYMBOL_GPL(mark_page_dirty_in_slot);
