+/*
+ * thread cpu affinity routines
+ */
+
+typedef struct _THREAD_BASIC_INFORMATION {
+ NTSTATUS ExitStatus;
+ PVOID TebBaseAddress;
+ CLIENT_ID ClientId;
+ ULONG_PTR AffinityMask;
+ KPRIORITY Priority;
+ LONG BasePriority;
+} THREAD_BASIC_INFORMATION;
+
+typedef THREAD_BASIC_INFORMATION *PTHREAD_BASIC_INFORMATION;
+
+#define THREAD_QUERY_INFORMATION (0x0040)
+
+NTSYSAPI
+NTSTATUS
+NTAPI
+ZwOpenThread (
+ __out PHANDLE ThreadHandle,
+ __in ACCESS_MASK DesiredAccess,
+ __in POBJECT_ATTRIBUTES ObjectAttributes,
+ __in_opt PCLIENT_ID ClientId
+ );
+
+NTSYSAPI
+NTSTATUS
+NTAPI
+ZwQueryInformationThread (
+ __in HANDLE ThreadHandle,
+ __in THREADINFOCLASS ThreadInformationClass,
+ __out_bcount(ThreadInformationLength) PVOID ThreadInformation,
+ __in ULONG ThreadInformationLength,
+ __out_opt PULONG ReturnLength
+ );
+
+NTSYSAPI
+NTSTATUS
+NTAPI
+ZwSetInformationThread (
+ __in HANDLE ThreadHandle,
+ __in THREADINFOCLASS ThreadInformationClass,
+ __in_bcount(ThreadInformationLength) PVOID ThreadInformation,
+ __in ULONG ThreadInformationLength
+ );
+
+HANDLE
+cfs_open_current_thread()
+{
+ NTSTATUS status;
+ HANDLE handle = NULL;
+ OBJECT_ATTRIBUTES oa;
+ CLIENT_ID cid;
+
+ /* initialize object attributes */
+ InitializeObjectAttributes( &oa, NULL, OBJ_KERNEL_HANDLE |
+ OBJ_CASE_INSENSITIVE, NULL, NULL);
+
+ /* initialize client id */
+ cid.UniqueProcess = PsGetCurrentProcessId();
+ cid.UniqueThread = PsGetCurrentThreadId();
+
+ /* get thread handle */
+ status = ZwOpenThread( &handle, THREAD_QUERY_INFORMATION |
+ THREAD_SET_INFORMATION, &oa, &cid);
+ if (!NT_SUCCESS(status)) {
+ handle = NULL;
+ }
+
+ return handle;
+}
+
+void
+cfs_close_thread_handle(HANDLE handle)
+{
+ if (handle)
+ ZwClose(handle);
+}
+
+KAFFINITY
+cfs_query_thread_affinity()
+{
+ NTSTATUS status;
+ HANDLE handle = NULL;
+ DWORD size;
+ THREAD_BASIC_INFORMATION TBI = {0};
+
+ /* open current thread */
+ handle = cfs_open_current_thread();
+ if (!handle) {
+ goto errorout;
+ }
+
+ /* query thread cpu affinity */
+ status = ZwQueryInformationThread(handle, ThreadBasicInformation,
+ &TBI, sizeof(THREAD_BASIC_INFORMATION), &size);
+ if (!NT_SUCCESS(status)) {
+ goto errorout;
+ }
+
+errorout:
+
+ cfs_close_thread_handle(handle);
+ return TBI.AffinityMask;
+}
+
+int
+cfs_set_thread_affinity(KAFFINITY affinity)
+{
+ NTSTATUS status;
+ HANDLE handle = NULL;
+
+ /* open current thread */
+ handle = cfs_open_current_thread();
+ if (!handle) {
+ goto errorout;
+ }
+
+ /* set thread cpu affinity */
+ status = ZwSetInformationThread(handle, ThreadAffinityMask,
+ &affinity, sizeof(KAFFINITY));
+ if (!NT_SUCCESS(status)) {
+ goto errorout;
+ }
+
+errorout:
+
+ cfs_close_thread_handle(handle);
+ return NT_SUCCESS(status);
+}
+
+int
+cfs_tie_thread_to_cpu(int cpu)
+{
+ return cfs_set_thread_affinity((KAFFINITY) (1 << cpu));
+}
+
+int
+cfs_set_thread_priority(KPRIORITY priority)
+{
+ NTSTATUS status;
+ HANDLE handle = NULL;
+
+ /* open current thread */
+ handle = cfs_open_current_thread();
+ if (!handle) {
+ goto errorout;
+ }
+
+ /* set thread cpu affinity */
+ status = ZwSetInformationThread(handle, ThreadPriority,
+ &priority, sizeof(KPRIORITY));
+ if (!NT_SUCCESS(status)) {
+ KdPrint(("set_thread_priority failed: %xh\n", status));
+ goto errorout;
+ }
+
+errorout:
+
+ cfs_close_thread_handle(handle);
+ return NT_SUCCESS(status);
+}
+
+int need_resched(void)
+{
+ return 0;
+}
+
+void cond_resched(void)
+{
+}
+