From c7bf8bab38f8b89c1371eedb9229e00a29b5ca5b Mon Sep 17 00:00:00 2001
From: David CARLIER <devnexen@gmail.com>
Date: Sat, 13 May 2023 10:39:10 +0100
Subject: [PATCH] std.os: adding linux's sched_setaffinity and its wrapper

---
 lib/std/c/linux.zig  |  5 +++++
 lib/std/os.zig       | 21 +++++++++++++++++++++
 lib/std/os/linux.zig | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+)

diff --git a/lib/std/c/linux.zig b/lib/std/c/linux.zig
index 856d64b21c..eb31248683 100644
--- a/lib/std/c/linux.zig
+++ b/lib/std/c/linux.zig
@@ -13,6 +13,10 @@ pub const ARCH = linux.ARCH;
 pub const AT = linux.AT;
 pub const CLOCK = linux.CLOCK;
 pub const CPU_COUNT = linux.CPU_COUNT;
+pub const CPU_SET = linux.CPU_SET;
+pub const CPU_ISSET = linux.CPU_ISSET;
+pub const CPU_CLR = linux.CPU_CLR;
+pub const CPU_ZERO = linux.CPU_ZERO;
 pub const E = linux.E;
 pub const Elf_Symndx = linux.Elf_Symndx;
 pub const F = linux.F;
@@ -245,6 +249,7 @@ pub extern "c" fn setrlimit64(resource: rlimit_resource, rlim: *const rlimit) c_
 
 pub extern "c" fn getrandom(buf_ptr: [*]u8, buf_len: usize, flags: c_uint) isize;
 pub extern "c" fn sched_getaffinity(pid: c_int, size: usize, set: *cpu_set_t) c_int;
+pub extern "c" fn sched_setaffinity(pid: c_int, size: usize, set: *const cpu_set_t) c_int;
 pub extern "c" fn eventfd(initval: c_uint, flags: c_uint) c_int;
 pub extern "c" fn epoll_ctl(epfd: fd_t, op: c_uint, fd: fd_t, event: ?*epoll_event) c_int;
 pub extern "c" fn epoll_create1(flags: c_uint) c_int;
diff --git a/lib/std/os.zig b/lib/std/os.zig
index 13feb7b75d..779e913230 100644
--- a/lib/std/os.zig
+++ b/lib/std/os.zig
@@ -5530,6 +5530,7 @@ pub fn clock_getres(clk_id: i32, res: *timespec) ClockGetTimeError!void {
 }
 
 pub const SchedGetAffinityError = error{PermissionDenied} || UnexpectedError;
+pub const SchedSetAffinityError = error{ InvalidCpu, PermissionDenied } || UnexpectedError;
 
 pub fn sched_getaffinity(pid: pid_t) SchedGetAffinityError!cpu_set_t {
     var set: cpu_set_t = undefined;
@@ -5557,6 +5558,26 @@ pub fn sched_getaffinity(pid: pid_t) SchedGetAffinityError!cpu_set_t {
     }
 }
 
+pub fn sched_setaffinity(pid: pid_t, cpus: []usize) SchedSetAffinityError!cpu_set_t {
+    var set: cpu_set_t = undefined;
+    if (builtin.os.tag == .linux) {
+        system.CPU_ZERO(&set);
+        for (cpus) |cpu| {
+            system.CPU_SET(cpu, &set);
+        }
+        switch (errno(system.sched_setaffinity(pid, @sizeOf(cpu_set_t), &set))) {
+            .SUCCESS => return set,
+            .FAULT => unreachable,
+            .SRCH => unreachable,
+            .INVAL => return error.InvalidCpu,
+            .PERM => return error.PermissionDenied,
+            else => |err| return unexpectedErrno(err),
+        }
+    } else {
+        @compileError("unsupported platform");
+    }
+}
+
 /// Used to convert a slice to a null terminated slice on the stack.
 /// TODO https://github.com/ziglang/zig/issues/287
 pub fn toPosixPath(file_path: []const u8) ![MAX_PATH_BYTES - 1:0]u8 {
diff --git a/lib/std/os/linux.zig b/lib/std/os/linux.zig
index 896169ed57..c2b9dd7813 100644
--- a/lib/std/os/linux.zig
+++ b/lib/std/os/linux.zig
@@ -1535,6 +1535,12 @@ pub fn mbind(addr: ?*anyopaque, len: u32, mode: i32, nodemask: *const u32, maxno
     return syscall6(.mbind, addr, len, mode, nodemask, maxnode, flags);
 }
 
+pub fn sched_setaffinity(pid: pid_t, size: usize, set: *const cpu_set_t) usize {
+    const rc = syscall3(.sched_setaffinity, @bitCast(usize, @as(isize, pid)), size, @ptrToInt(set));
+    if (@bitCast(isize, rc) < 0) return rc;
+    return 0;
+}
+
 pub fn epoll_create() usize {
     return epoll_create1(0);
 }
@@ -3553,6 +3559,11 @@ pub const CPU_SETSIZE = 128;
 pub const cpu_set_t = [CPU_SETSIZE / @sizeOf(usize)]usize;
 pub const cpu_count_t = std.meta.Int(.unsigned, std.math.log2(CPU_SETSIZE * 8));
 
+fn cpu_mask(s: usize) cpu_count_t {
+    var x = s & (CPU_SETSIZE * 8);
+    return @intCast(cpu_count_t, 1) << @intCast(u4, x);
+}
+
 pub fn CPU_COUNT(set: cpu_set_t) cpu_count_t {
     var sum: cpu_count_t = 0;
     for (set) |x| {
@@ -3561,6 +3572,32 @@ pub fn CPU_COUNT(set: cpu_set_t) cpu_count_t {
     return sum;
 }
 
+pub fn CPU_ZERO(set: *cpu_set_t) void {
+    @memset(set, 0);
+}
+
+pub fn CPU_SET(cpu: usize, set: *cpu_set_t) void {
+    const x = cpu / @sizeOf(usize);
+    if (x < @sizeOf(cpu_set_t)) {
+        (set.*)[x] |= cpu_mask(x);
+    }
+}
+
+pub fn CPU_ISSET(cpu: usize, set: cpu_set_t) bool {
+    const x = cpu / @sizeOf(usize);
+    if (x < @sizeOf(cpu_set_t)) {
+        return set[x] & cpu_mask(x);
+    }
+    return false;
+}
+
+pub fn CPU_CLR(cpu: usize, set: *cpu_set_t) void {
+    const x = cpu / @sizeOf(usize);
+    if (x < @sizeOf(cpu_set_t)) {
+        (set.*)[x] &= !cpu_mask(x);
+    }
+}
+
 pub const MINSIGSTKSZ = switch (native_arch) {
     .x86, .x86_64, .arm, .mipsel => 2048,
     .aarch64 => 5120,