From c38455f9f647648d2f9adf9fdbd0dabe84eda062 Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Mon, 23 Sep 2024 15:46:52 +0800
Subject: [PATCH] Add tests for system call `sched_get/set_affinity`

---
 test/apps/cpu_affinity/cpu_affinity.c      | 69 ++++++++++++++++++++++
 test/apps/cpu_affinity/sched_getaffinity.c | 35 -----------
 test/apps/scripts/process.sh               |  1 +
 3 files changed, 70 insertions(+), 35 deletions(-)
 create mode 100644 test/apps/cpu_affinity/cpu_affinity.c
 delete mode 100644 test/apps/cpu_affinity/sched_getaffinity.c

diff --git a/test/apps/cpu_affinity/cpu_affinity.c b/test/apps/cpu_affinity/cpu_affinity.c
new file mode 100644
index 00000000..2e0bf4d2
--- /dev/null
+++ b/test/apps/cpu_affinity/cpu_affinity.c
@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: MPL-2.0
+
+#define _GNU_SOURCE
+#include <sched.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <pthread.h>
+
+int main()
+{
+	pid_t pid = getpid();
+	cpu_set_t mask;
+
+	// Get current affinity mask
+	if (sched_getaffinity(pid, sizeof(cpu_set_t), &mask) == -1) {
+		perror("sched_getaffinity");
+		exit(EXIT_FAILURE);
+	}
+
+	printf("Current CPU affinity:");
+	int cur_cpu_count = 0;
+	for (int i = 0; i < CPU_SETSIZE; i++) {
+		if (CPU_ISSET(i, &mask)) {
+			printf(" %d", i);
+			cur_cpu_count++;
+		}
+	}
+	printf("\n");
+	if (cur_cpu_count == 0) {
+		printf("Error: No CPU affinity set\n");
+		exit(EXIT_FAILURE);
+	}
+
+	// Set the process to run on CPU 0 only
+	CPU_ZERO(&mask);
+	CPU_SET(0, &mask);
+	if (sched_setaffinity(pid, sizeof(cpu_set_t), &mask) == -1) {
+		perror("sched_setaffinity");
+		exit(EXIT_FAILURE);
+	}
+	printf("Set CPU affinity to CPU 0\n");
+
+	// Verify the new CPU affinity
+	if (sched_getaffinity(pid, sizeof(cpu_set_t), &mask) == -1) {
+		perror("sched_getaffinity");
+		exit(EXIT_FAILURE);
+	}
+
+	printf("New CPU affinity:");
+	cur_cpu_count = 0;
+	for (int i = 0; i < CPU_SETSIZE; i++) {
+		if (CPU_ISSET(i, &mask)) {
+			printf(" %d", i);
+			cur_cpu_count++;
+			if (i != 0) {
+				printf("Error: CPU affinity not set to CPU 0\n");
+				exit(EXIT_FAILURE);
+			}
+		}
+	}
+	printf("\n");
+	if (cur_cpu_count != 1) {
+		printf("Error: CPU affinity not set to CPU 0\n");
+		exit(EXIT_FAILURE);
+	}
+
+	return 0;
+}
diff --git a/test/apps/cpu_affinity/sched_getaffinity.c b/test/apps/cpu_affinity/sched_getaffinity.c
deleted file mode 100644
index 31e6187c..00000000
--- a/test/apps/cpu_affinity/sched_getaffinity.c
+++ /dev/null
@@ -1,35 +0,0 @@
-// SPDX-License-Identifier: MPL-2.0
-
-#define _GNU_SOURCE
-#include <stdio.h>
-#include <unistd.h>
-#include <syscall.h>
-#include <errno.h>
-#include <sched.h> // Include sched.h for CPU_SETSIZE
-
-int main()
-{
-	// Create a mask for CPU_SETSIZE number of CPUs
-	unsigned long mask[CPU_SETSIZE / sizeof(unsigned long)];
-	int mask_size = sizeof(mask);
-
-	// Call the raw syscall to retrieve the CPU affinity mask of the current process
-	long res = syscall(__NR_sched_getaffinity, 0, mask_size, &mask);
-
-	if (res < 0) {
-		perror("Error calling sched_getaffinity");
-		return errno;
-	}
-
-	// Print the CPUs that are part of the current process's affinity mask
-	printf("Process can run on: ");
-	for (int i = 0; i < CPU_SETSIZE; ++i) {
-		if (mask[i / (8 * sizeof(long))] &
-		    (1UL << (i % (8 * sizeof(long))))) {
-			printf("%d ", i);
-		}
-	}
-	printf("\n");
-
-	return 0;
-}
diff --git a/test/apps/scripts/process.sh b/test/apps/scripts/process.sh
index f08c3c86..58617068 100755
--- a/test/apps/scripts/process.sh
+++ b/test/apps/scripts/process.sh
@@ -13,6 +13,7 @@ tests="
 clone3/clone_exit_signal
 clone3/clone_no_exit_signal
 clone3/clone_process
+cpu_affinity/cpu_affinity
 execve/execve
 eventfd2/eventfd2
 fork/fork