[RFC PATCH v2 1/1] KVM: x86/mmu: Zap empty TDP leaf page tables

[Hao Zhang <hao_zhang_kdev@163.com>]

From: Hao Zhang <zhanghao1@kylinos.cn>

KVM's TDP MMU intentionally zaps only leaf SPTEs for some invalidations,
leaving the page table hierarchy intact so that subsequent faults can reuse
the existing shadow pages.

That behavior can be pathological if a later invalidation covers a range
whose retained 4K leaf tables contain only non-present SPTEs.  KVM still
descends into every retained leaf page table and scans all 512 entries even
though there is nothing left to zap.

Add an opt-in TDP iterator event that reports when the walk has finished a
child page table, and use that event in tdp_mmu_zap_leafs() to unlink empty
4K leaf page tables that are fully covered by the zap.

Do not unlink a leaf page table that had a leaf SPTE zapped in the current
pass.  Doing so would add allocation/free churn for rolling invalidation
workloads that refault the same page table.

Because unlinking a shadow page changes the old leaf-only RCU assumption,
flush any pending remote TLBs before dropping RCU protection if a shadow
page was removed.

Signed-off-by: Hao Zhang <zhanghao1@kylinos.cn>
---
 arch/x86/kvm/mmu/tdp_iter.c | 29 +++++++++++++--
 arch/x86/kvm/mmu/tdp_iter.h | 10 ++++++
 arch/x86/kvm/mmu/tdp_mmu.c  | 88 +++++++++++++++++++++++++++++++++++++--------
 3 files changed, 109 insertions(+), 18 deletions(-)

diff --git a/arch/x86/kvm/mmu/tdp_iter.c b/arch/x86/kvm/mmu/tdp_iter.c
index 9e17bfa80901..d9be4298410d 100644
--- a/arch/x86/kvm/mmu/tdp_iter.c
+++ b/arch/x86/kvm/mmu/tdp_iter.c
@@ -132,7 +132,8 @@ static bool try_step_side(struct tdp_iter *iter)
  * can continue from the next entry in the parent page table. Returns true on a
  * successful step up, false if already in the root page.
  */
-static bool try_step_up(struct tdp_iter *iter)
+static bool try_step_up(struct tdp_iter *iter,
+			struct tdp_iter_child_event *event)
 {
 	if (iter->level == iter->root_level)
 		return false;
@@ -141,6 +142,14 @@ static bool try_step_up(struct tdp_iter *iter)
 	iter->gfn = gfn_round_for_level(iter->gfn, iter->level);
 	tdp_iter_refresh_sptep(iter);
 
+	if (event && !event->valid) {
+		event->sptep = iter->sptep;
+		event->old_spte = iter->old_spte;
+		event->gfn = iter->gfn;
+		event->level = iter->level;
+		event->valid = true;
+	}
+
 	return true;
 }
 
@@ -160,8 +169,12 @@ static bool try_step_up(struct tdp_iter *iter)
  *    SPTE will have already been visited, and so the iterator must also step
  *    to the side again.
  */
-void tdp_iter_next(struct tdp_iter *iter)
+static void __tdp_iter_next(struct tdp_iter *iter,
+			    struct tdp_iter_child_event *event)
 {
+	if (event)
+		event->valid = false;
+
 	if (iter->yielded) {
 		tdp_iter_restart(iter);
 		return;
@@ -173,7 +186,17 @@ void tdp_iter_next(struct tdp_iter *iter)
 	do {
 		if (try_step_side(iter))
 			return;
-	} while (try_step_up(iter));
+	} while (try_step_up(iter, event));
 	iter->valid = false;
 }
 
+void tdp_iter_next(struct tdp_iter *iter)
+{
+	__tdp_iter_next(iter, NULL);
+}
+
+void tdp_iter_next_post_order(struct tdp_iter *iter,
+			      struct tdp_iter_child_event *event)
+{
+	__tdp_iter_next(iter, event);
+}
diff --git a/arch/x86/kvm/mmu/tdp_iter.h b/arch/x86/kvm/mmu/tdp_iter.h
index f898d8d0d93c..b7e2605fbf9e 100644
--- a/arch/x86/kvm/mmu/tdp_iter.h
+++ b/arch/x86/kvm/mmu/tdp_iter.h
@@ -123,6 +123,14 @@ struct tdp_iter {
 	bool yielded;
 };
 
+struct tdp_iter_child_event {
+	tdp_ptep_t sptep;
+	u64 old_spte;
+	gfn_t gfn;
+	int level;
+	bool valid;
+};
+
 /*
  * Iterates over every SPTE mapping the GFN range [start, end) in a
  * preorder traversal.
@@ -145,6 +153,8 @@ tdp_ptep_t spte_to_child_pt(u64 pte, int level);
 void tdp_iter_start(struct tdp_iter *iter, struct kvm_mmu_page *root,
 		    int min_level, gfn_t next_last_level_gfn, gfn_t gfn_bits);
 void tdp_iter_next(struct tdp_iter *iter);
+void tdp_iter_next_post_order(struct tdp_iter *iter,
+			      struct tdp_iter_child_event *event);
 void tdp_iter_restart(struct tdp_iter *iter);
 
 #endif /* __KVM_X86_MMU_TDP_ITER_H */
diff --git a/arch/x86/kvm/mmu/tdp_mmu.c b/arch/x86/kvm/mmu/tdp_mmu.c
index 44dad106fad1..82446c4737f7 100644
--- a/arch/x86/kvm/mmu/tdp_mmu.c
+++ b/arch/x86/kvm/mmu/tdp_mmu.c
@@ -923,6 +923,38 @@ bool kvm_tdp_mmu_zap_possible_nx_huge_page(struct kvm *kvm,
 	return true;
 }
 
+static bool tdp_mmu_zap_completed_empty_leaf_pt(struct kvm *kvm,
+						struct tdp_iter_child_event *event,
+						gfn_t start, gfn_t end,
+						tdp_ptep_t zapped_leaf_pt)
+{
+	struct kvm_mmu_page *child_sp;
+
+	if (event->level != PG_LEVEL_2M ||
+	    !is_shadow_present_pte(event->old_spte) ||
+	    is_last_spte(event->old_spte, event->level))
+		return false;
+
+	if (event->gfn < start ||
+	    event->gfn + KVM_PAGES_PER_HPAGE(event->level) > end)
+		return false;
+
+	/* Leave productive leaf tables in place to avoid refault churn. */
+	if (zapped_leaf_pt == event->sptep)
+		return false;
+
+	child_sp = spte_to_child_sp(event->old_spte);
+	if (is_mirror_sp(child_sp))
+		return false;
+
+	if (child_sp->role.level != PG_LEVEL_4K)
+		return false;
+
+	tdp_mmu_set_spte(kvm, event->sptep, event->old_spte,
+			 SHADOW_NONPRESENT_VALUE, event->gfn, event->level);
+	return true;
+}
+
 /*
  * If can_yield is true, will release the MMU lock and reschedule if the
  * scheduler needs the CPU or there is contention on the MMU lock. If this
@@ -933,7 +965,10 @@ bool kvm_tdp_mmu_zap_possible_nx_huge_page(struct kvm *kvm,
 static bool tdp_mmu_zap_leafs(struct kvm *kvm, struct kvm_mmu_page *root,
 			      gfn_t start, gfn_t end, bool can_yield, bool flush)
 {
+	struct tdp_iter_child_event event = {};
 	struct tdp_iter iter;
+	tdp_ptep_t zapped_leaf_pt = NULL;
+	bool removed_sp = false;
 
 	end = min(end, tdp_mmu_max_gfn_exclusive());
 
@@ -941,33 +976,56 @@ static bool tdp_mmu_zap_leafs(struct kvm *kvm, struct kvm_mmu_page *root,
 
 	rcu_read_lock();
 
-	for_each_tdp_pte_min_level(iter, kvm, root, PG_LEVEL_4K, start, end) {
+	tdp_iter_start(&iter, root, PG_LEVEL_4K, start,
+		       kvm_gfn_root_bits(kvm, root));
+	while (iter.valid && iter.gfn < end) {
 		if (can_yield &&
 		    tdp_mmu_iter_cond_resched(kvm, &iter, flush, false)) {
 			flush = false;
+			tdp_iter_next(&iter);
 			continue;
 		}
 
-		if (!is_shadow_present_pte(iter.old_spte) ||
-		    !is_last_spte(iter.old_spte, iter.level))
-			continue;
+		if (is_shadow_present_pte(iter.old_spte) &&
+		    is_last_spte(iter.old_spte, iter.level)) {
+			if (iter.level == PG_LEVEL_4K) {
+				struct kvm_mmu_page *sp;
 
-		tdp_mmu_iter_set_spte(kvm, &iter, SHADOW_NONPRESENT_VALUE);
+				sp = sptep_to_sp(rcu_dereference(iter.sptep));
+				zapped_leaf_pt = sp->ptep;
+			}
 
-		/*
-		 * Zappings SPTEs in invalid roots doesn't require a TLB flush,
-		 * see kvm_tdp_mmu_zap_invalidated_roots() for details.
-		 */
-		if (!root->role.invalid)
-			flush = true;
-	}
+			tdp_mmu_iter_set_spte(kvm, &iter,
+					      SHADOW_NONPRESENT_VALUE);
 
-	rcu_read_unlock();
+			/*
+			 * Zappings SPTEs in invalid roots doesn't require a TLB flush,
+			 * see kvm_tdp_mmu_zap_invalidated_roots() for details.
+			 */
+			if (!root->role.invalid)
+				flush = true;
+		}
+
+		tdp_iter_next_post_order(&iter, &event);
+		if (event.valid &&
+		    tdp_mmu_zap_completed_empty_leaf_pt(kvm, &event, start,
+							end, zapped_leaf_pt)) {
+			if (!root->role.invalid)
+				flush = true;
+			removed_sp = true;
+		}
+	}
 
 	/*
-	 * Because this flow zaps _only_ leaf SPTEs, the caller doesn't need
-	 * to provide RCU protection as no 'struct kvm_mmu_page' will be freed.
+	 * If any shadow pages were removed, service pending TLB flushes before
+	 * dropping RCU protection, as required by the TDP MMU iterator.
 	 */
+	if (removed_sp && flush) {
+		kvm_flush_remote_tlbs(kvm);
+		flush = false;
+	}
+
+	rcu_read_unlock();
 	return flush;
 }
 
-- 
2.15.0