[PATCH] KVM: Add helpers to allocate memslot metadata

[Fred Griffoul <griffoul@gmail.com>]

From: Fred Griffoul <fgriffo@amazon.co.uk>

__vcalloc() makes every allocation at least one page, so a memslot
that spans a single page uses 8 KiB of vmalloc for 8 bytes of
lpage_info, plus 4 KiB for a 16-byte dirty bitmap once dirty logging
is enabled. The waste grows with the number of memslots and VMs on a
host. kvcalloc() avoids it for small arrays, but refuses allocations
above INT_MAX bytes, which the rmap and write-tracking arrays of a
large memslot exceed.

Add kvm_memslot_metadata_alloc() and kvm_memslot_metadata_free(),
which use kvzalloc() up to INT_MAX bytes and vmalloc above, and route
the dirty bitmap, rmaps, lpage_info and write-tracking arrays through
them. That way, how KVM allocates memslot metadata is decided in one
place instead of at each array.

Suggested-by: Sean Christopherson <seanjc@google.com>
Link: https://lore.kernel.org/all/arr1VgGZmgUmY4hc@google.com/
Signed-off-by: Fred Griffoul <fgriffo@amazon.co.uk>
---
Mushahid's patch [1] switched the lpage_info arrays and dirty bitmaps to
kvcalloc(). Sean suggested dedicated helpers for memslot metadata
instead [2], which this patch adds.

[1] https://lore.kernel.org/all/20260815142218.85067-1-hmushi@amazon.co.uk/
[2] https://lore.kernel.org/all/arr1VgGZmgUmY4hc@google.com/
---
 arch/x86/kvm/mmu/page_track.c |  6 +++---
 arch/x86/kvm/x86.c            | 10 +++++-----
 include/linux/kvm_host.h      |  2 ++
 virt/kvm/kvm_main.c           | 23 +++++++++++++++++++++--
 4 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/arch/x86/kvm/mmu/page_track.c b/arch/x86/kvm/mmu/page_track.c
index f12558dfcd81..f33547964b98 100644
--- a/arch/x86/kvm/mmu/page_track.c
+++ b/arch/x86/kvm/mmu/page_track.c
@@ -43,7 +43,7 @@ bool kvm_page_track_write_tracking_enabled(struct kvm *kvm)
 
 void kvm_page_track_free_memslot(struct kvm_memory_slot *slot)
 {
-	vfree(slot->arch.gfn_write_track);
+	kvm_memslot_metadata_free(slot->arch.gfn_write_track);
 	slot->arch.gfn_write_track = NULL;
 }
 
@@ -53,8 +53,8 @@ static int __kvm_page_track_write_tracking_alloc(struct kvm_memory_slot *slot,
 	const size_t size = sizeof(*slot->arch.gfn_write_track);
 
 	if (!slot->arch.gfn_write_track)
-		slot->arch.gfn_write_track = __vcalloc(npages, size,
-						       GFP_KERNEL_ACCOUNT);
+		slot->arch.gfn_write_track =
+			kvm_memslot_metadata_alloc(npages, size);
 
 	return slot->arch.gfn_write_track ? 0 : -ENOMEM;
 }
diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
index 26e7beaff8a0..dc716ad83c8f 100644
--- a/arch/x86/kvm/x86.c
+++ b/arch/x86/kvm/x86.c
@@ -10146,7 +10146,7 @@ static void memslot_rmap_free(struct kvm_memory_slot *slot)
 	int i;
 
 	for (i = 0; i < KVM_NR_PAGE_SIZES; ++i) {
-		vfree(slot->arch.rmap[i]);
+		kvm_memslot_metadata_free(slot->arch.rmap[i]);
 		slot->arch.rmap[i] = NULL;
 	}
 }
@@ -10158,7 +10158,7 @@ void kvm_arch_free_memslot(struct kvm *kvm, struct kvm_memory_slot *slot)
 	memslot_rmap_free(slot);
 
 	for (i = 1; i < KVM_NR_PAGE_SIZES; ++i) {
-		vfree(slot->arch.lpage_info[i - 1]);
+		kvm_memslot_metadata_free(slot->arch.lpage_info[i - 1]);
 		slot->arch.lpage_info[i - 1] = NULL;
 	}
 
@@ -10177,7 +10177,7 @@ int memslot_rmap_alloc(struct kvm_memory_slot *slot, unsigned long npages)
 		if (slot->arch.rmap[i])
 			continue;
 
-		slot->arch.rmap[i] = __vcalloc(lpages, sz, GFP_KERNEL_ACCOUNT);
+		slot->arch.rmap[i] = kvm_memslot_metadata_alloc(lpages, sz);
 		if (!slot->arch.rmap[i]) {
 			memslot_rmap_free(slot);
 			return -ENOMEM;
@@ -10214,7 +10214,7 @@ static int kvm_alloc_memslot_metadata(struct kvm *kvm,
 
 		lpages = __kvm_mmu_slot_lpages(slot, npages, level);
 
-		linfo = __vcalloc(lpages, sizeof(*linfo), GFP_KERNEL_ACCOUNT);
+		linfo = kvm_memslot_metadata_alloc(lpages, sizeof(*linfo));
 		if (!linfo)
 			goto out_free;
 
@@ -10250,7 +10250,7 @@ static int kvm_alloc_memslot_metadata(struct kvm *kvm,
 	memslot_rmap_free(slot);
 
 	for (i = 1; i < KVM_NR_PAGE_SIZES; ++i) {
-		vfree(slot->arch.lpage_info[i - 1]);
+		kvm_memslot_metadata_free(slot->arch.lpage_info[i - 1]);
 		slot->arch.lpage_info[i - 1] = NULL;
 	}
 	return -ENOMEM;
diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index e7fc94e2910a..11a6f501a78b 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -1253,6 +1253,8 @@ enum kvm_mr_change {
 
 int kvm_set_internal_memslot(struct kvm *kvm,
 			     const struct kvm_userspace_memory_region2 *mem);
+void *kvm_memslot_metadata_alloc(size_t n, size_t size);
+void kvm_memslot_metadata_free(void *p);
 void kvm_arch_free_memslot(struct kvm *kvm, struct kvm_memory_slot *slot);
 void kvm_arch_memslots_updated(struct kvm *kvm, u64 gen);
 int kvm_arch_prepare_memory_region(struct kvm *kvm,
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index cf4cb586f140..ce872ffc779f 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -941,12 +941,31 @@ static void kvm_destroy_pm_notifier(struct kvm *kvm)
 }
 #endif /* CONFIG_HAVE_KVM_PM_NOTIFIER */
 
+void *kvm_memslot_metadata_alloc(size_t n, size_t size)
+{
+	size_t bytes;
+
+	if (unlikely(check_mul_overflow(n, size, &bytes)))
+		return NULL;
+
+	/* kvzalloc() refuses sizes above INT_MAX, which large memslots need. */
+	if (bytes <= INT_MAX)
+		return kvzalloc(bytes, GFP_KERNEL_ACCOUNT);
+
+	return __vmalloc(bytes, GFP_KERNEL_ACCOUNT | __GFP_ZERO);
+}
+
+void kvm_memslot_metadata_free(void *p)
+{
+	kvfree(p);
+}
+
 static void kvm_destroy_dirty_bitmap(struct kvm_memory_slot *memslot)
 {
 	if (!memslot->dirty_bitmap)
 		return;
 
-	vfree(memslot->dirty_bitmap);
+	kvm_memslot_metadata_free(memslot->dirty_bitmap);
 	memslot->dirty_bitmap = NULL;
 }
 
@@ -1449,7 +1468,7 @@ static int kvm_alloc_dirty_bitmap(struct kvm_memory_slot *memslot)
 {
 	unsigned long dirty_bytes = kvm_dirty_bitmap_bytes(memslot);
 
-	memslot->dirty_bitmap = __vcalloc(2, dirty_bytes, GFP_KERNEL_ACCOUNT);
+	memslot->dirty_bitmap = kvm_memslot_metadata_alloc(2, dirty_bytes);
 	if (!memslot->dirty_bitmap)
 		return -ENOMEM;
 

base-commit: d21eef2084a9af87ac2a783bc7c2b175ba114e4e
-- 
2.47.3