[PATCH] iommupt/kunit: Read an exact NR_SECONDARY_PAGETABLE snapshot for the leak check

[fangyu . yu]

From: Fangyu Yu <fangyu.yu@xxxxxxxxxxxxxxxxx>

NR_SECONDARY_PAGETABLE is per-CPU batched and only folded into the global
counter about once per second. The tests complete within a few seconds, so
the init/exit checks can read stale values and randomly fail when a fold
lands mid-test, even though no page table is leaked:

    # test_map_simple: ASSERTION FAILED at .../kunit_iommu_pt.h:473
    Expected priv->orig_nr_secondary_pagetable ==
    global_node_page_state(NR_SECONDARY_PAGETABLE), but
        priv->orig_nr_secondary_pagetable == 68 (0x44)
        global_node_page_state(NR_SECONDARY_PAGETABLE) == 70 (0x46)

Sum the per-CPU batches together with the global counter to get an exact
snapshot before and after each test case, which makes the check
deterministic without any settling delay.

Fixes: bc5233c0904e ("iommupt: Add a kunit test for the IOMMU implementation")
Signed-off-by: Fangyu Yu <fangyu.yu@xxxxxxxxxxxxxxxxx>
---
 drivers/iommu/generic_pt/kunit_iommu_pt.h | 40 ++++++++++++++++++++---
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/drivers/iommu/generic_pt/kunit_iommu_pt.h b/drivers/iommu/generic_pt/kunit_iommu_pt.h
index ece1c9b8c55d..d20728b196be 100644
--- a/drivers/iommu/generic_pt/kunit_iommu_pt.h
+++ b/drivers/iommu/generic_pt/kunit_iommu_pt.h
@@ -4,6 +4,7 @@
  */
 #include "kunit_iommu.h"
 #include "pt_iter.h"
+#include <linux/cpu.h>
 #include <linux/generic_pt/iommu.h>
 #include <linux/iommu.h>
 
@@ -438,6 +439,36 @@ static struct kunit_case iommu_test_cases[] = {
 	{},
 };
 
+/*
+ * The global NR_SECONDARY_PAGETABLE counter is batched per-CPU and folded
+ * into the global value asynchronously, sum the global value and the per-CPU
+ * deltas to get an exact snapshot.
+ *
+ * The global value must be read without the clamp of global_node_page_state():
+ * the raw counter is transiently negative whenever a folded free batch
+ * outweighs allocation deltas still pending on other CPUs, and clamping
+ * before adding the deltas back would overestimate and report phantom leaks.
+ * global_node_page_state_monotonic() is normally reserved for monotonic
+ * counters because a raw negative read is garbage as a page count, but the
+ * sum with the pending deltas always restores the exact non-negative value.
+ */
+static unsigned long pt_kunit_nr_secondary_pagetable(void)
+{
+	unsigned long nr = global_node_page_state_monotonic(NR_SECONDARY_PAGETABLE);
+
+#ifdef CONFIG_SMP
+	int nid, cpu;
+
+	cpus_read_lock();
+	for_each_online_cpu(cpu)
+		for_each_online_node(nid)
+			nr += per_cpu_ptr(NODE_DATA(nid)->per_cpu_nodestats,
+					  cpu)->vm_node_stat_diff[NR_SECONDARY_PAGETABLE];
+	cpus_read_unlock();
+#endif
+	return nr;
+}
+
 static int pt_kunit_iommu_init(struct kunit *test)
 {
 	struct kunit_iommu_priv *priv;
@@ -447,8 +478,7 @@ static int pt_kunit_iommu_init(struct kunit *test)
 	if (!priv)
 		return -ENOMEM;
 
-	priv->orig_nr_secondary_pagetable =
-		global_node_page_state(NR_SECONDARY_PAGETABLE);
+	priv->orig_nr_secondary_pagetable = pt_kunit_nr_secondary_pagetable();
 	ret = pt_kunit_priv_init(test, priv);
 	if (ret) {
 		kunit_kfree(test, priv);
@@ -461,17 +491,19 @@ static int pt_kunit_iommu_init(struct kunit *test)
 static void pt_kunit_iommu_exit(struct kunit *test)
 {
 	struct kunit_iommu_priv *priv = test->priv;
+	unsigned long nr;
 
 	if (!test->priv)
 		return;
 
 	pt_iommu_deinit(priv->iommu);
+	nr = pt_kunit_nr_secondary_pagetable();
+
 	/*
 	 * Look for memory leaks, assumes kunit is running isolated and nothing
 	 * else is using secondary page tables.
 	 */
-	KUNIT_ASSERT_EQ(test, priv->orig_nr_secondary_pagetable,
-			global_node_page_state(NR_SECONDARY_PAGETABLE));
+	KUNIT_ASSERT_EQ(test, priv->orig_nr_secondary_pagetable, nr);
 	kunit_kfree(test, test->priv);
 }
 
-- 
2.50.1