diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 65eb26a05..8843dc35c 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -2534,6 +2534,52 @@ static bool kvm_pre_set_memory_attributes(struct kvm *kvm,
 	return kvm_arch_pre_set_memory_attributes(kvm, range);
 }
 
+/*
+ * Reserve or store @entry for every gfn in [@start, @end) that is backed by
+ * a memslot.
+ *
+ * Only gfns that are backed by a memslot are tracked; attributes for any
+ * other gfn are meaningless, as such gfns cannot be mapped into the guest.
+ * Walking the entire, user-controlled range instead would let userspace
+ * force KVM to hold slots_lock for an unbounded amount of time while
+ * allocating an unbounded amount of memory.  That can starve other
+ * slots_lock users, e.g. kvm_gmem_release() can get stuck waiting for
+ * slots_lock and trigger a hung task splat.
+ */
+static int kvm_set_mem_attributes_walk(struct kvm *kvm, gfn_t start, gfn_t end,
+				       void *entry, bool store)
+{
+	int as_id;
+
+	for (as_id = 0; as_id < kvm_arch_nr_memslot_as_ids(kvm); as_id++) {
+		struct kvm_memslots *slots = __kvm_memslots(kvm, as_id);
+		struct kvm_memslot_iter iter;
+
+		kvm_for_each_memslot_in_gfn_range(&iter, slots, start, end) {
+			gfn_t gfn = max(start, iter.slot->base_gfn);
+			gfn_t gfn_end = min(end, iter.slot->base_gfn + iter.slot->npages);
+
+			for (; gfn < gfn_end; gfn++) {
+				int r;
+
+				if (store)
+					r = xa_err(xa_store(&kvm->mem_attr_array,
+							    gfn, entry,
+							    GFP_KERNEL_ACCOUNT));
+				else
+					r = xa_reserve(&kvm->mem_attr_array,
+						       gfn, GFP_KERNEL_ACCOUNT);
+				if (r)
+					return r;
+
+				cond_resched();
+			}
+		}
+	}
+
+	return 0;
+}
+
 /* Set @attributes for the gfn range [@start, @end). */
 static int kvm_vm_set_mem_attributes(struct kvm *kvm, gfn_t start, gfn_t end,
 				     unsigned long attributes)
@@ -2555,7 +2601,6 @@ static int kvm_vm_set_mem_attributes(struct kvm *kvm, gfn_t start, gfn_t end,
 		.on_lock = kvm_mmu_invalidate_end,
 		.may_block = true,
 	};
-	unsigned long i;
 	void *entry;
 	int r = 0;
 
@@ -2571,24 +2616,17 @@ static int kvm_vm_set_mem_attributes(struct kvm *kvm, gfn_t start, gfn_t end,
 
 	/*
 	 * Reserve memory ahead of time to avoid having to deal with failures
-	 * partway through setting the new attributes.
+	 * partway through setting the new attributes.  Note, only gfns that
+	 * are backed by a memslot are reserved, see the walker.
 	 */
-	for (i = start; i < end; i++) {
-		r = xa_reserve(&kvm->mem_attr_array, i, GFP_KERNEL_ACCOUNT);
-		if (r)
-			goto out_unlock;
-
-		cond_resched();
-	}
+	r = kvm_set_mem_attributes_walk(kvm, start, end, entry, false);
+	if (r)
+		goto out_unlock;
 
 	kvm_handle_gfn_range(kvm, &pre_set_range);
 
-	for (i = start; i < end; i++) {
-		r = xa_err(xa_store(&kvm->mem_attr_array, i, entry,
-				    GFP_KERNEL_ACCOUNT));
-		KVM_BUG_ON(r, kvm);
-		cond_resched();
-	}
+	r = kvm_set_mem_attributes_walk(kvm, start, end, entry, true);
+	KVM_BUG_ON(r, kvm);
 
 	kvm_handle_gfn_range(kvm, &post_set_range);