KVM: selftests: Expand set of APIs for pinning tasks to a single CPU

Expand kvm_pin_this_task_to_pcpu() into a set of APIs to allow pinning a
task (or self) to a CPU (any or specific).  This will allow deduplicating
code throughout a variety of selftests.

Opportunistically use "self" instead of "this_task" as it is both more
concise and less ambiguous.

Link: https://lore.kernel.org/r/20250626001225.744268-4-seanjc@google.com
Signed-off-by: Sean Christopherson <seanjc@google.com>

diff --git a/tools/testing/selftests/kvm/include/kvm_util.h b/tools/testing/selftests/kvm/include/kvm_util.h
index bee65ca087217675b3e1d125486c3cff316b16e9..492f0edc795597435e489739050151d58586ddca 100644 (file)
--- a/tools/testing/selftests/kvm/include/kvm_util.h
+++ b/tools/testing/selftests/kvm/include/kvm_util.h
@@ -20,6 +20,8 @@
 
 #include <sys/ioctl.h>
 
+#include <pthread.h>
+
 #include "kvm_util_arch.h"
 #include "kvm_util_types.h"
 #include "sparsebit.h"
@@ -1013,7 +1015,34 @@ struct kvm_vcpu *vm_recreate_with_one_vcpu(struct kvm_vm *vm);
 
 void kvm_set_files_rlimit(uint32_t nr_vcpus);
 
-void kvm_pin_this_task_to_pcpu(uint32_t pcpu);
+int __pin_task_to_cpu(pthread_t task, int cpu);
+
+static inline void pin_task_to_cpu(pthread_t task, int cpu)
+{
+       int r;
+
+       r = __pin_task_to_cpu(task, cpu);
+       TEST_ASSERT(!r, "Failed to set thread affinity to pCPU '%u'", cpu);
+}
+
+static inline int pin_task_to_any_cpu(pthread_t task)
+{
+       int cpu = sched_getcpu();
+
+       pin_task_to_cpu(task, cpu);
+       return cpu;
+}
+
+static inline void pin_self_to_cpu(int cpu)
+{
+       pin_task_to_cpu(pthread_self(), cpu);
+}
+
+static inline int pin_self_to_any_cpu(void)
+{
+       return pin_task_to_any_cpu(pthread_self());
+}
+
 void kvm_print_vcpu_pinning_help(void);
 void kvm_parse_vcpu_pinning(const char *pcpus_string, uint32_t vcpu_to_pcpu[],
                            int nr_vcpus);

diff --git a/tools/testing/selftests/kvm/lib/kvm_util.c b/tools/testing/selftests/kvm/lib/kvm_util.c
index a055343a7bf7586a2f43d9e5072b1a032d3093fc..1a59d69edc6ee877117807939679619af824f0d5 100644 (file)
--- a/tools/testing/selftests/kvm/lib/kvm_util.c
+++ b/tools/testing/selftests/kvm/lib/kvm_util.c
@@ -605,15 +605,14 @@ struct kvm_vcpu *vm_recreate_with_one_vcpu(struct kvm_vm *vm)
        return vm_vcpu_recreate(vm, 0);
 }
 
-void kvm_pin_this_task_to_pcpu(uint32_t pcpu)
+int __pin_task_to_cpu(pthread_t task, int cpu)
 {
-       cpu_set_t mask;
-       int r;
+       cpu_set_t cpuset;
 
-       CPU_ZERO(&mask);
-       CPU_SET(pcpu, &mask);
-       r = sched_setaffinity(0, sizeof(mask), &mask);
-       TEST_ASSERT(!r, "sched_setaffinity() failed for pCPU '%u'.", pcpu);
+       CPU_ZERO(&cpuset);
+       CPU_SET(cpu, &cpuset);
+
+       return pthread_setaffinity_np(task, sizeof(cpuset), &cpuset);
 }
 
 static uint32_t parse_pcpu(const char *cpu_str, const cpu_set_t *allowed_mask)
@@ -667,7 +666,7 @@ void kvm_parse_vcpu_pinning(const char *pcpus_string, uint32_t vcpu_to_pcpu[],
 
        /* 2. Check if the main worker needs to be pinned. */
        if (cpu) {
-               kvm_pin_this_task_to_pcpu(parse_pcpu(cpu, &allowed_mask));
+               pin_self_to_cpu(parse_pcpu(cpu, &allowed_mask));
                cpu = strtok(NULL, delim);
        }
 

diff --git a/tools/testing/selftests/kvm/lib/memstress.c b/tools/testing/selftests/kvm/lib/memstress.c
index 313277486a1dee8f5b015cbe69ccfe7ede6beeaa..557c0a0a56580c9705729d3209fcce501742522e 100644 (file)
--- a/tools/testing/selftests/kvm/lib/memstress.c
+++ b/tools/testing/selftests/kvm/lib/memstress.c
@@ -265,7 +265,7 @@ static void *vcpu_thread_main(void *data)
        int vcpu_idx = vcpu->vcpu_idx;
 
        if (memstress_args.pin_vcpus)
-               kvm_pin_this_task_to_pcpu(memstress_args.vcpu_to_pcpu[vcpu_idx]);
+               pin_self_to_cpu(memstress_args.vcpu_to_pcpu[vcpu_idx]);
 
        WRITE_ONCE(vcpu->running, true);