[PATCH 1/3] KVM: arm64: Don't walk the whole IPA space when destroying a pKVM stage-2

[Fuad Tabba <fuad.tabba@linux.dev>]

Since commit 4ddfab5436b6 ("KVM: arm64: Reschedule as needed when
destroying the stage-2 page-tables"), kvm_stage2_destroy() steps through
the IPA space in block-sized chunks, and with pKVM each chunk is a
lookup in the VM's mapping tree. However, pKVM gives every VM the host's
IPA limit, so with LPA2 a teardown makes 2^27 lookups with 16K pages
(2^22 with 4K), nearly all of them for empty space.

Walk the mapping tree once instead, and reschedule after passing each
mapping to the reclaim or unshare helper. The unshare helper can't
reschedule itself, as the map and unmap paths call it under the
mmu_lock. A failed unshare now leaves only its own mapping behind, not
the rest of its chunk.

Fixes: 4ddfab5436b69 ("KVM: arm64: Reschedule as needed when destroying the stage-2 page-tables")
Signed-off-by: Fuad Tabba <fuad.tabba@linux.dev>
---
 arch/arm64/kvm/mmu.c  |  7 ++++++-
 arch/arm64/kvm/pkvm.c | 17 +++++++++++++----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 1db10d8d6db73..42977cf8a7f5d 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -1022,7 +1022,12 @@ static void kvm_stage2_destroy(struct kvm_pgtable *pgt)
 {
 	unsigned int ia_bits = VTCR_EL2_IPA(pgt->mmu->vtcr);
 
-	stage2_destroy_range(pgt, 0, BIT(ia_bits));
+	/* pKVM walks its mapping tree, not the IPA space, and reschedules per mapping */
+	if (is_protected_kvm_enabled())
+		pkvm_pgtable_stage2_destroy_range(pgt, 0, BIT(ia_bits));
+	else
+		stage2_destroy_range(pgt, 0, BIT(ia_bits));
+
 	KVM_PGT_FN(kvm_pgtable_stage2_destroy_pgd)(pgt);
 }
 
diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index 016997288fd23..992ed3ca6bdd9 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -504,6 +504,7 @@ void pkvm_pgtable_stage2_destroy_range(struct kvm_pgtable *pgt,
 {
 	struct kvm *kvm = kvm_s2_mmu_to_kvm(pgt->mmu);
 	pkvm_handle_t handle = kvm->arch.pkvm.handle;
+	struct pkvm_mapping *mapping;
 
 	if (!handle)
 		return;
@@ -513,10 +514,18 @@ void pkvm_pgtable_stage2_destroy_range(struct kvm_pgtable *pgt,
 		kvm->arch.pkvm.is_dying = true;
 	}
 
-	if (kvm_vm_is_protected(kvm))
-		__pkvm_pgtable_stage2_reclaim(pgt, addr, addr + size);
-	else
-		__pkvm_pgtable_stage2_unshare(pgt, addr, addr + size);
+	/* Called once for the whole IPA space, so reschedule between mappings */
+	for_each_mapping_in_range_safe(pgt, addr, addr + size, mapping) {
+		u64 start = __pkvm_mapping_start(mapping);
+		u64 end = __pkvm_mapping_end(mapping) + 1;
+
+		if (kvm_vm_is_protected(kvm))
+			__pkvm_pgtable_stage2_reclaim(pgt, start, end);
+		else
+			__pkvm_pgtable_stage2_unshare(pgt, start, end);
+
+		cond_resched();
+	}
 }
 
 void pkvm_pgtable_stage2_destroy_pgd(struct kvm_pgtable *pgt)
-- 
2.47.3