Windows affinity fix for thread based jobs

Signed-off-by: Jens Axboe <jaxboe@fusionio.com>

diff --git a/os/os-windows.h b/os/os-windows.h
index 9f2243eb34dadf5e2255a506bf0d0d050b3b3004..db4127b3a2e35eed0fa796eec671e7334ad2f21e 100644 (file)
--- a/os/os-windows.h
+++ b/os/os-windows.h
@@ -16,6 +16,7 @@
 #define FIO_HAVE_FALLOCATE\r
 #define FIO_HAVE_FDATASYNC\r
 #define FIO_HAVE_WINDOWSAIO\r
+#define FIO_HAVE_GETTID\r
 \r
 #define FIO_USE_GENERIC_RAND\r
 \r
@@ -92,6 +93,11 @@ static inline void os_get_tmpdir(char *path, int len)
 \r
 typedef DWORD_PTR os_cpu_mask_t;\r
 \r
+static inline int gettid(void)\r
+{\r
+       return GetCurrentThreadId();\r
+}\r
+\r
 static inline int pid_to_winpid(int pid)\r
 {\r
        int winpid = 0;\r
@@ -120,31 +126,57 @@ static inline int pid_to_winpid(int pid)
        return winpid;\r
 }\r
 \r
+HANDLE WINAPI OpenThread(\r
+    DWORD dwDesiredAccess,\r
+    BOOL bInheritHandle,\r
+    DWORD dwThreadId);\r
+    \r
+DWORD WINAPI GetProcessIdOfThread(HANDLE Thread);\r
+\r
 static inline int fio_setaffinity(int pid, os_cpu_mask_t cpumask)\r
 {\r
-       int winpid = pid_to_winpid(pid);\r
-       HANDLE h = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION, TRUE, winpid);\r
-       if (h == NULL)\r
-               return -1;\r
+       HANDLE h;\r
+       BOOL bSuccess;\r
+       int winpid;\r
        \r
-       BOOL bSuccess = SetProcessAffinityMask(h, cpumask);\r
+       h = OpenThread(THREAD_QUERY_INFORMATION | THREAD_SET_INFORMATION, TRUE, pid);\r
+       if (h != NULL) {\r
+               bSuccess = SetThreadAffinityMask(h, cpumask);\r
+       } else {\r
+               // then we might have a process id instead of a thread id\r
+               winpid = pid_to_winpid(pid);\r
+               h = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION, TRUE, winpid);\r
+               if (h == NULL)\r
+                       return -1;\r
+\r
+               bSuccess = SetProcessAffinityMask(h, cpumask);\r
+       }\r
+\r
        CloseHandle(h);\r
+\r
        return (bSuccess)? 0 : -1;\r
 }\r
 \r
 static inline void fio_getaffinity(int pid, os_cpu_mask_t *mask)\r
 {\r
        os_cpu_mask_t systemMask;\r
-       int winpid = pid_to_winpid(pid);\r
-       HANDLE h = OpenProcess(PROCESS_QUERY_INFORMATION, TRUE, winpid);\r
-       if (h == NULL)\r
-       {\r
-               printf("OpenProcess failed!\n");\r
-               return;\r
-  }\r
-  \r
-       GetProcessAffinityMask(h, mask, &systemMask);\r
-       CloseHandle(h);\r
+       int winpid;\r
+       \r
+       HANDLE h = OpenThread(THREAD_QUERY_INFORMATION, TRUE, pid);\r
+       if (h != NULL)\r
+               winpid = GetProcessIdOfThread(h);\r
+       else\r
+               winpid = pid_to_winpid(pid);\r
+       \r
+       h = OpenProcess(PROCESS_QUERY_INFORMATION, TRUE, winpid);\r
+\r
+       if (h != NULL) {\r
+               GetProcessAffinityMask(h, mask, &systemMask);\r
+               CloseHandle(h);\r
+       } else {\r
+               fprintf(stderr, "fio_getaffinity failed: failed to get handle for pid %d\n", pid);\r
+       }\r
+       \r
 }\r
 \r
 static inline void fio_cpu_clear(os_cpu_mask_t *mask, int cpu)\r