[PATCH 4/4] selftests/ksm: add test for only_merge_zero_pages

["Xu Xin (vivo)" <xu.xin@linux.dev>]

Add a functional test that enables use_zero_pages and
only_merge_zero_pages, maps a 2 MiB region with the first half filled
with non-zero data and the second half with zeroes, and verifies that
after KSM scanning reaches steady state:

  - ksm_zero_pages equals the zero half (merged with the kernel zero
    page)
  - ksm_merging_pages is 0 (non-zero pages never reach the stable tree)
  - ksm_rmap_items is 0 (non-zero pages are skipped before rmap_item
    allocation, and the zero pages' rmap_items are freed once a scan
    finds the merged pages are no longer anonymous)

The test skips if the only_merge_zero_pages attribute is absent.

Signed-off-by: Xu Xin (vivo) <xu.xin@linux.dev>
---
 .../selftests/mm/ksm_functional_tests.c       | 117 ++++++++++++++++++
 1 file changed, 117 insertions(+)

diff --git a/tools/testing/selftests/mm/ksm_functional_tests.c b/tools/testing/selftests/mm/ksm_functional_tests.c
index 2ffb97ae3120..764ead3425cb 100644
--- a/tools/testing/selftests/mm/ksm_functional_tests.c
+++ b/tools/testing/selftests/mm/ksm_functional_tests.c
@@ -257,6 +257,122 @@ static void test_unmerge_zero_pages(void)
 	munmap(map, size);
 }
 
+/*
+ * Enable or disable /sys/kernel/mm/ksm/only_merge_zero_pages. Returns 0 on
+ * success, -errno on failure.
+ */
+static int ksm_set_only_merge_zero_pages(int enable)
+{
+	int fd;
+	ssize_t ret;
+
+	fd = open("/sys/kernel/mm/ksm/only_merge_zero_pages", O_RDWR);
+	if (fd < 0)
+		return -errno;
+
+	ret = write(fd, enable ? "1" : "0", 1);
+	close(fd);
+	return ret == 1 ? 0 : -errno;
+}
+
+static void test_only_merge_zero_pages(void)
+{
+	const unsigned int size = 2 * MiB;
+	const unsigned int half = size / 2;
+	const unsigned int half_pages = half / pagesize;
+	char *map;
+	int i;
+	long zero_pages, merging_pages, rmap_items;
+
+	ksft_print_msg("[RUN] %s\n", __func__);
+
+	if (ksm_use_zero_pages() < 0) {
+		ksft_test_result_skip("write /sys/kernel/mm/ksm/use_zero_pages failed\n");
+		return;
+	}
+	if (ksm_set_only_merge_zero_pages(1) < 0) {
+		ksft_test_result_skip("write /sys/kernel/mm/ksm/only_merge_zero_pages failed\n");
+		return;
+	}
+
+	/* Stabilize accounting by disabling KSM completely. */
+	if (ksm_stop() < 0) {
+		ksft_test_result_fail("Disabling (unmerging) KSM failed\n");
+		goto restore;
+	}
+
+	map = mmap(NULL, size, PROT_READ | PROT_WRITE,
+		   MAP_PRIVATE | MAP_ANON, -1, 0);
+	if (map == MAP_FAILED) {
+		ksft_test_result_fail("mmap() failed\n");
+		goto restore;
+	}
+
+	/* Don't use THP. Ignore if THP are not around on a kernel. */
+	if (madvise(map, size, MADV_NOHUGEPAGE) && errno != EINVAL) {
+		ksft_test_result_fail("MADV_NOHUGEPAGE failed\n");
+		goto unmap;
+	}
+
+	/*
+	 * Fill the first half with non-zero data (all 0x01) and the second
+	 * half with zeroes. Under only_merge_zero_pages the zero half should
+	 * be merged with the kernel zero page, while the non-zero half must be
+	 * skipped entirely (no rmap_item allocated, no stable-tree insertion).
+	 */
+	memset(map, 0x01, half);
+	memset(map + half, 0x00, half);
+
+	if (madvise(map, size, MADV_MERGEABLE)) {
+		ksft_test_result_fail("MADV_MERGEABLE failed\n");
+		goto unmap;
+	}
+
+	/*
+	 * Run KSM until steady state. Each ksm_start() waits for two full
+	 * scans: the first records the checksums, the second merges the zero
+	 * half with the kernel zero page. The second ksm_start() waits for
+	 * more scans so that the rmap_items of the (now merged, no longer
+	 * anonymous) zero pages are freed and ksm_rmap_items settles at 0.
+	 */
+	for (i = 0; i < 2; i++) {
+		if (ksm_start() < 0) {
+			ksft_test_result_fail("Running KSM failed\n");
+			goto unmap;
+		}
+	}
+
+	/*
+	 * Expected results (steady state):
+	 *  - ksm_zero_pages: the zero half (size/2 pages) is merged with the
+	 *    kernel zero page.
+	 *  - ksm_merging_pages: 0, because the non-zero half is skipped before
+	 *    it can reach the stable tree, and zero pages are accounted in
+	 *    ksm_zero_pages, not ksm_merging_pages.
+	 *  - ksm_rmap_items: 0. The non-zero half is skipped before
+	 *    get_next_rmap_item(), so it never allocates an rmap_item; and the
+	 *    zero half's rmap_items are freed once a scan finds the (now
+	 *    merged) pages are no longer anonymous.
+	 */
+	zero_pages = ksm_get_self_stat("ksm_zero_pages");
+	merging_pages = ksm_get_self_stat("ksm_merging_pages");
+	rmap_items = ksm_get_self_stat("ksm_rmap_items");
+
+	ksft_print_msg("ksm_zero_pages=%ld ksm_merging_pages=%ld ksm_rmap_items=%ld"
+		       " (expected %u/0/0)\n", zero_pages, merging_pages, rmap_items,
+		       half_pages);
+
+	ksft_test_result(zero_pages == half_pages && merging_pages == 0 &&
+			 rmap_items == 0, "only_merge_zero_pages: zero half merged, "
+			 "non-zero half skipped\n");
+
+unmap:
+	ksm_stop();
+	munmap(map, size);
+restore:
+	ksm_set_only_merge_zero_pages(0);
+}
+
 static void test_unmerge_discarded(void)
 {
 	const unsigned int size = 2 * MiB;
@@ -827,6 +943,7 @@ int main(int argc, char **argv)
 
 	test_unmerge();
 	test_unmerge_zero_pages();
+	test_only_merge_zero_pages();
 	test_unmerge_discarded();
 #ifdef __NR_userfaultfd
 	test_unmerge_uffd_wp();
-- 
2.34.1