From 4660ab927cebf7c900fb57ec5d4e2c2e341a1ea9 Mon Sep 17 00:00:00 2001
From: Natalya McKay <zea_64@proton.me>
Date: Sun, 25 May 2025 00:31:44 -0400
Subject: [PATCH] Implement setThreadAffinity for Linux.

---
 .../cortex/voxy/common/util/ThreadUtils.java  | 27 +++++++++++++++++++
 .../voxy/common/util/cpu/CpuLayout.java       | 12 +++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/main/java/me/cortex/voxy/common/util/ThreadUtils.java b/src/main/java/me/cortex/voxy/common/util/ThreadUtils.java
index 4db33e7a..34f80835 100644
--- a/src/main/java/me/cortex/voxy/common/util/ThreadUtils.java
+++ b/src/main/java/me/cortex/voxy/common/util/ThreadUtils.java
@@ -5,6 +5,7 @@ import org.lwjgl.system.MemoryStack;
 import org.lwjgl.system.MemoryUtil;
 import org.lwjgl.system.Platform;
 import org.lwjgl.system.windows.Kernel32;
+import org.lwjgl.system.APIUtil;
 
 //Platform specific code to assist in thread utilities
 public class ThreadUtils {
@@ -15,6 +16,7 @@ public class ThreadUtils {
     private static final boolean isWindows = Platform.get() == Platform.WINDOWS;
     private static final long SetThreadPriority;
     private static final long SetThreadSelectedCpuSetMasks;
+    private static final long schedSetaffinity;
     static {
         if (isWindows) {
             SetThreadPriority = Kernel32.getLibrary().getFunctionAddress("SetThreadPriority");
@@ -23,6 +25,13 @@ public class ThreadUtils {
             SetThreadPriority = 0;
             SetThreadSelectedCpuSetMasks = 0;
         }
+
+        if (Platform.get() == Platform.LINUX) {
+            var libc = APIUtil.apiCreateLibrary("libc.so.6");
+            schedSetaffinity = APIUtil.apiGetFunctionAddress(libc, "sched_setaffinity");
+        } else {
+            schedSetaffinity = 0;
+        }
     }
 
     public static boolean SetThreadSelectedCpuSetMasksWin32(long mask) {
@@ -70,4 +79,22 @@ public class ThreadUtils {
         }
         return true;
     }
+
+    public static boolean schedSetaffinityLinux(long masks[]) {
+        if (schedSetaffinity == 0 || isWindows) {
+            return false;
+        }
+        try (var stack = MemoryStack.stackPush()) {
+            long ptr = stack.ncalloc(8, masks.length, 8);
+            for (int i=0; i<masks.length; i++) {
+                MemoryUtil.memPutLong(ptr+i*8L, masks[i]);
+            }
+
+            int retVal = JNI.invokePPI(0, (long)masks.length*8, ptr, schedSetaffinity);
+            if (retVal != 0) {
+                throw new IllegalStateException();
+            }
+            return true;
+        }
+    }
 }
diff --git a/src/main/java/me/cortex/voxy/common/util/cpu/CpuLayout.java b/src/main/java/me/cortex/voxy/common/util/cpu/CpuLayout.java
index 77fa3bf6..a57b59db 100644
--- a/src/main/java/me/cortex/voxy/common/util/cpu/CpuLayout.java
+++ b/src/main/java/me/cortex/voxy/common/util/cpu/CpuLayout.java
@@ -25,7 +25,8 @@ public class CpuLayout {
     }
 
     public static void setThreadAffinity(Affinity... affinities) {
-        if (Platform.get() == Platform.WINDOWS) {
+        var platform = Platform.get();
+        if (platform == Platform.WINDOWS) {
             long[] msks = new long[affinities.length];
             short[] groups = new short[affinities.length];Arrays.fill(groups, (short) -1);
             int i = 0;
@@ -36,8 +37,15 @@ public class CpuLayout {
                 msks[idx] |= a.msk;
             }
             ThreadUtils.SetThreadSelectedCpuSetMasksWin32(Arrays.copyOf(msks, i), Arrays.copyOf(groups, i));
+        } else if (platform == Platform.LINUX) {
+            Arrays.sort(affinities, (a, b) -> a.group - b.group);
+            long[] msks = new long[affinities.length];
+            for (int i=0; i<affinities.length; i++) {
+                msks[i] = affinities[i].msk;
+            }
+            ThreadUtils.schedSetaffinityLinux(msks);
         } else {
-            Logger.error("TODO: THIS");
+            Logger.error("Don't know how to set thread affinity on this platform.");
         }
     }