set and get a thread's CPU affinity mask lab

gistfile1.c

#define _GNU_SOURCE
#include <sched.h>
#include <stdio.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdlib.h> /* abort() */
#include <stdarg.h>
#include <string.h>
#include <errno.h>

#define ERR(msg) fprintf(stderr, "[%s] [%s] at %s line %d\n", msg, strerror(errno), __FILE__, __LINE__)

void print_cpuset(const char *msg);
void use_cpu(int n, ...);


int main() {
	print_cpuset("origin");

	use_cpu(1, 1);
	print_cpuset("set 1");

	use_cpu(1, 0);
	print_cpuset("set 0");

	use_cpu(2, 0, 1);
	print_cpuset("set 0,1");
	return 0;
}


void print_cpuset(const char *msg) {
	pid_t pid;
	cpu_set_t *cpusetp = NULL;
	size_t cpuset_siz = 0;
	int ncpu = 0;
	int err = 0;

	cpuset_siz = CPU_ALLOC_SIZE((int)sysconf(_SC_NPROCESSORS_ONLN));
	cpusetp = CPU_ALLOC((int)sysconf(_SC_NPROCESSORS_ONLN));
	if (!cpusetp) {
		abort();
	}

	pid = getpid();

	CPU_ZERO_S(cpuset_siz, cpusetp);
	err = sched_getaffinity(pid, cpuset_siz, cpusetp);
	if (err) {
		ERR("sched_getaffinity");
		abort();
	}
	printf("%s pid %ld, available cpu num %d\n",  msg, (long)pid, CPU_COUNT_S(cpuset_siz, cpusetp));
	CPU_FREE(cpusetp);
}

void use_cpu(int n, ...) {
	cpu_set_t *cpusetp = NULL;
	int i;
	va_list ap;

	cpusetp = CPU_ALLOC(n);
	if (!cpusetp) {
		abort();
	}
	CPU_ZERO_S(CPU_ALLOC_SIZE(n), cpusetp);

	va_start (ap, n);
	for (i = 0; i < n; i++) {
		CPU_SET(va_arg(ap, int), cpusetp);
	}
	va_end (ap);
	if (sched_setaffinity(0, CPU_ALLOC_SIZE(n), cpusetp)) {
		ERR("sched_setaffinity");
		abort();
	}
	CPU_FREE(cpusetp);
}