Process affinity implementation (#1041)

CHANGELOG.md

@@ -2,6 +2,7 @@
 ================
 * [#1038](https://github.com/oshi/oshi/pull/1038): More Battery Statistics. - [@dbwiddis](https://github.com/dbwiddis).
 * [#1039](https://github.com/oshi/oshi/pull/1039): JNA 5.5.0. - [@dbwiddis](https://github.com/dbwiddis).
+* [#1041](https://github.com/oshi/oshi/pull/1041): Process Affinity. - [@dbwiddis](https://github.com/dbwiddis).
 * Your contribution here.
 
 4.1.0 (10/16/2019), 4.1.1 (10/24/2019)

oshi-core/src/main/java/oshi/jna/platform/linux/LinuxLibc.java

@@ -34,7 +34,5 @@
  */
 public interface LinuxLibc extends LibC, CLibrary {
 
-    /** Constant <code>INSTANCE</code> */
     LinuxLibc INSTANCE = Native.load("c", LinuxLibc.class);
-
 }

oshi-core/src/main/java/oshi/jna/platform/windows/Kernel32.java

@@ -0,0 +1,45 @@
+package oshi.jna.platform.windows;
+
+import com.sun.jna.Native; // NOSONAR squid:S1192
+import com.sun.jna.platform.win32.WinNT;
+import com.sun.jna.win32.W32APIOptions;
+
+public interface Kernel32 extends com.sun.jna.platform.win32.Kernel32 {
+
+    Kernel32 INSTANCE = Native.load("kernel32", Kernel32.class, W32APIOptions.DEFAULT_OPTIONS);
+
+    /**
+     * Retrieves the process affinity mask for the specified process and the system
+     * affinity mask for the system.
+     *
+     * @param hProcess
+     *            A handle to the process whose affinity mask is desired.
+     *            <p>
+     *            This handle must have the {@link WinNT#PROCESS_QUERY_INFORMATION}
+     *            or {@link WinNT#PROCESS_QUERY_LIMITED_INFORMATION} access right.
+     * @param lpProcessAffinityMask
+     *            A pointer to a variable that receives the affinity mask for the
+     *            specified process.
+     * @param lpSystemAffinityMask
+     *            A pointer to a variable that receives the affinity mask for the
+     *            system.
+     * @return If the function succeeds, returns {@code true} and the function sets
+     *         the variables pointed to by {@code lpProcessAffinityMask} and
+     *         {@code lpSystemAffinityMask} to the appropriate affinity masks.
+     *         <p>
+     *         On a system with more than 64 processors, if the threads of the
+     *         calling process are in a single processor group, the function sets
+     *         the variables pointed to by {@code lpProcessAffinityMask} and
+     *         {@code lpSystemAffinityMask} to the process affinity mask and the
+     *         processor mask of active logical processors for that group. If the
+     *         calling process contains threads in multiple groups, the function
+     *         returns zero for both affinity masks.
+     *         <p>
+     *         If the function fails, the return value is {@code false}, and the
+     *         values of the variables pointed to by {@code lpProcessAffinityMask}
+     *         and {@code lpSystemAffinityMask} are undefined. To get extended error
+     *         information, call {@link #GetLastError()}.
+     */
+    boolean GetProcessAffinityMask(HANDLE hProcess, ULONG_PTRByReference lpProcessAffinityMask,
+            ULONG_PTRByReference lpSystemAffinityMask);
+}

oshi-core/src/main/java/oshi/software/os/OperatingSystem.java

@@ -208,6 +208,25 @@ enum ProcessSort {
      */
     OSProcess[] getChildProcesses(int parentPid, int limit, ProcessSort sort);
 
+    /**
+     * Retrieves the process affinity mask for the specified process.
+     * <p>
+     * On Windows systems with more than 64 processors, if the threads of the
+     * calling process are in a single processor group, returns the process affinity
+     * mask for that group (which may be zero if the specified process is running in
+     * a different group). If the calling process contains threads in multiple
+     * groups, returns zero.
+     * <p>
+     * If the Operating System fails to retrieve an affinity mask (e.g., the process
+     * has terminated), returns zero.
+     *
+     * @param processId
+     *            The process ID for which to retrieve the affinity.
+     * @return a bit vector in which each bit represents the processors that a
+     *         process is allowed to run on.
+     */
+    long getProcessAffinityMask(int processId);
+
     /**
      * Gets the current process ID
      *

oshi-core/src/main/java/oshi/software/os/linux/LinuxOperatingSystem.java

@@ -30,6 +30,7 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -354,6 +355,22 @@ private static int getParentPidFromProcFile(int pid) {
         return (int) statArray[ProcPidStat.PPID.ordinal()];
     }
 
+    @Override
+    public long getProcessAffinityMask(int processId) {
+        // Would prefer to use native sched_getaffinity call but variable sizing is
+        // kernel-dependent and requires C macros, so we use command line instead.
+        String mask = ExecutingCommand.getFirstAnswer("taskset -p " + processId);
+        // Output:
+        // pid 3283's current affinity mask: 3
+        // pid 9726's current affinity mask: f
+        String[] split = ParseUtil.whitespaces.split(mask);
+        try {
+            return new BigInteger(split[split.length - 1], 16).longValue();
+        } catch (NumberFormatException e) {
+            return 0;
+        }
+    }
+
     @Override
     public int getProcessId() {
         return LinuxLibc.INSTANCE.getpid();

oshi-core/src/main/java/oshi/software/os/mac/MacOperatingSystem.java

@@ -401,6 +401,13 @@ private String getCommandLine(int pid) {
         return String.join("\0", args);
     }
 
+    @Override
+    public long getProcessAffinityMask(int processId) {
+        // macOS doesn't do affinity. Return a bitmask of the current processors.
+        int logicalProcessorCount = SysctlUtil.sysctl("hw.logicalcpu", 1);
+        return logicalProcessorCount < 64 ? (1L << logicalProcessorCount) - 1 : -1L;
+    }
+
     @Override
     public int getProcessId() {
         return SystemB.INSTANCE.getpid();

oshi-core/src/main/java/oshi/software/os/unix/freebsd/FreeBsdOperatingSystem.java

@@ -230,6 +230,28 @@ private List<OSProcess> getProcessListFromPS(String psCommand, int pid, boolean
         return procs;
     }
 
+    @Override
+    public long getProcessAffinityMask(int processId) {
+        long bitMask = 0L;
+        // Would prefer to use native cpuset_getaffinity call but variable sizing is
+        // kernel-dependent and requires C macros, so we use commandline instead.
+        String cpuset = ExecutingCommand.getFirstAnswer("cpuset -gp " + processId);
+        // Sample output:
+        // pid 8 mask: 0, 1
+        // cpuset: getaffinity: No such process
+        String[] split = cpuset.split(":");
+        if (split.length > 1) {
+            String[] bits = split[1].split(",");
+            for (String bit : bits) {
+                int bitToSet = ParseUtil.parseIntOrDefault(bit.trim(), -1);
+                if (bitToSet >= 0) {
+                    bitMask |= (1L << bitToSet);
+                }
+            }
+        }
+        return bitMask;
+    }
+
     @Override
     public int getProcessId() {
         return FreeBsdLibc.INSTANCE.getpid();

oshi-core/src/main/java/oshi/software/os/unix/solaris/SolarisOperatingSystem.java

@@ -206,6 +206,39 @@ private List<OSProcess> getProcessListFromPS(String psCommand, int pid, boolean
         return procs;
     }
 
+    @Override
+    public long getProcessAffinityMask(int processId) {
+        long bitMask = 0L;
+        String cpuset = ExecutingCommand.getFirstAnswer("pbind -q " + processId);
+        // Sample output:
+        // <empty string if no binding>
+        // pid 101048 strongly bound to processor(s) 0 1 2 3.
+        if (cpuset.isEmpty()) {
+            List<String> allProcs = ExecutingCommand.runNative("psrinfo");
+            for (String proc : allProcs) {
+                String[] split = ParseUtil.whitespaces.split(proc);
+                int bitToSet = ParseUtil.parseIntOrDefault(split[0], -1);
+                if (bitToSet >= 0) {
+                    bitMask |= (1L << bitToSet);
+                }
+            }
+            return bitMask;
+        } else if (cpuset.endsWith(".") && cpuset.contains("strongly bound to processor(s)")) {
+            String parse = cpuset.substring(0, cpuset.length() - 1);
+            String[] split = ParseUtil.whitespaces.split(parse);
+            for (int i = split.length - 1; i >= 0; i--) {
+                int bitToSet = ParseUtil.parseIntOrDefault(split[i], -1);
+                if (bitToSet >= 0) {
+                    bitMask |= (1L << bitToSet);
+                } else {
+                    // Once we run into the word processor(s) we're done
+                    break;
+                }
+            }
+        }
+        return bitMask;
+    }
+
     @Override
     public int getProcessId() {
         return SolarisLibc.INSTANCE.getpid();

oshi-core/src/main/java/oshi/software/os/windows/WindowsOperatingSystem.java

@@ -48,7 +48,7 @@
 import com.sun.jna.platform.win32.Advapi32Util.Account;
 import com.sun.jna.platform.win32.Advapi32Util.EventLogIterator;
 import com.sun.jna.platform.win32.Advapi32Util.EventLogRecord;
-import com.sun.jna.platform.win32.Kernel32;
+import com.sun.jna.platform.win32.BaseTSD.ULONG_PTRByReference;
 import com.sun.jna.platform.win32.Kernel32Util;
 import com.sun.jna.platform.win32.Psapi;
 import com.sun.jna.platform.win32.Psapi.PERFORMANCE_INFORMATION;
@@ -78,6 +78,7 @@
 import com.sun.jna.ptr.IntByReference;
 import com.sun.jna.ptr.PointerByReference;
 
+import oshi.jna.platform.windows.Kernel32;
 import oshi.software.common.AbstractOperatingSystem;
 import oshi.software.os.FileSystem;
 import oshi.software.os.NetworkParams;
@@ -578,6 +579,19 @@ private Map<Integer, OSProcess> buildProcessMapFromPerfCounters(Collection<Integ
         return processMap;
     }
 
+    @Override
+    public long getProcessAffinityMask(int processId) {
+        final HANDLE pHandle = Kernel32.INSTANCE.OpenProcess(WinNT.PROCESS_QUERY_INFORMATION, false, processId);
+        if (pHandle != null) {
+            ULONG_PTRByReference processAffinity = new ULONG_PTRByReference();
+            ULONG_PTRByReference systemAffinity = new ULONG_PTRByReference();
+            if (Kernel32.INSTANCE.GetProcessAffinityMask(pHandle, processAffinity, systemAffinity)) {
+                return Pointer.nativeValue(processAffinity.getValue().toPointer());
+            }
+        }
+        return 0L;
+    }
+
     @Override
     public int getProcessId() {
         return Kernel32.INSTANCE.GetCurrentProcessId();

oshi-core/src/test/java/oshi/SystemInfoTest.java

@@ -238,6 +238,8 @@ private static void printCpu(CentralProcessor processor) {
     }
 
     private static void printProcesses(OperatingSystem os, GlobalMemory memory) {
+        oshi.add("My PID: " + os.getProcessId() + " with affinity "
+                + Long.toBinaryString(os.getProcessAffinityMask(os.getProcessId())));
         oshi.add("Processes: " + os.getProcessCount() + ", Threads: " + os.getThreadCount());
         // Sort by highest CPU
         List<OSProcess> procs = Arrays.asList(os.getProcesses(5, ProcessSort.CPU));