Merge pull request #10003 from viriuwu/nt-thread-name

std.Thread.getName/setName: rework windows implementation

4 files changed, 126 insertions, 63 deletions

diff --git a/lib/std/Thread.zig b/lib/std/Thread.zig
index 1e54053146..07e5b9d233 100644
--- a/lib/std/Thread.zig
+++ b/lib/std/Thread.zig
@@ -4,6 +4,7 @@
 
 const std = @import("std.zig");
 const builtin = @import("builtin");
+const math = std.math;
 const os = std.os;
 const assert = std.debug.assert;
 const target = builtin.target;
@@ -86,20 +87,28 @@ pub fn setName(self: Thread, name: []const u8) SetNameError!void {
             try file.writer().writeAll(name);
             return;
         },
-        .windows => if (target.os.isAtLeast(.windows, .win10_rs1)) |res| {
-            // SetThreadDescription is only available since version 1607, which is 10.0.14393.795
-            // See https://en.wikipedia.org/wiki/Microsoft_Windows_SDK
-            if (!res) return error.Unsupported;
-
-            var name_buf_w: [max_name_len:0]u16 = undefined;
-            const length = try std.unicode.utf8ToUtf16Le(&name_buf_w, name);
-            name_buf_w[length] = 0;
+        .windows => {
+            var buf: [max_name_len]u16 = undefined;
+            const len = try std.unicode.utf8ToUtf16Le(&buf, name);
+            const byte_len = math.cast(c_ushort, len * 2) catch return error.NameTooLong;
+
+            // Note: NT allocates its own copy, no use-after-free here.
+            const unicode_string = os.windows.UNICODE_STRING{
+                .Length = byte_len,
+                .MaximumLength = byte_len,
+                .Buffer = &buf,
+            };
 
-            try os.windows.SetThreadDescription(
+            switch (os.windows.ntdll.NtSetInformationThread(
                 self.getHandle(),
-                @ptrCast(os.windows.LPWSTR, &name_buf_w),
-            );
-            return;
+                .ThreadNameInformation,
+                &unicode_string,
+                @sizeOf(os.windows.UNICODE_STRING),
+            )) {
+                .SUCCESS => return,
+                .NOT_IMPLEMENTED => return error.Unsupported,
+                else => |err| return os.windows.unexpectedStatus(err),
+            }
         },
         .macos, .ios, .watchos, .tvos => if (use_pthreads) {
             // There doesn't seem to be a way to set the name for an arbitrary thread, only the current one.
@@ -189,18 +198,25 @@ pub fn getName(self: Thread, buffer_ptr: *[max_name_len:0]u8) GetNameError!?[]co
             // musl doesn't provide pthread_getname_np and there's no way to retrieve the thread id of an arbitrary thread.
             return error.Unsupported;
         },
-        .windows => if (target.os.isAtLeast(.windows, .win10_rs1)) |res| {
-            // GetThreadDescription is only available since version 1607, which is 10.0.14393.795
-            // See https://en.wikipedia.org/wiki/Microsoft_Windows_SDK
-            if (!res) return error.Unsupported;
-
-            var name_w: os.windows.LPWSTR = undefined;
-            try os.windows.GetThreadDescription(self.getHandle(), &name_w);
-            defer os.windows.LocalFree(name_w);
+        .windows => {
+            const buf_capacity = @sizeOf(os.windows.UNICODE_STRING) + (@sizeOf(u16) * max_name_len);
+            var buf: [buf_capacity]u8 align(@alignOf(os.windows.UNICODE_STRING)) = undefined;
 
-            const data_len = try std.unicode.utf16leToUtf8(buffer, std.mem.sliceTo(name_w, 0));
-
-            return if (data_len >= 1) buffer[0..data_len] else null;
+            switch (os.windows.ntdll.NtQueryInformationThread(
+                self.getHandle(),
+                .ThreadNameInformation,
+                &buf,
+                buf_capacity,
+                null,
+            )) {
+                .SUCCESS => {
+                    const string = @ptrCast(*const os.windows.UNICODE_STRING, &buf);
+                    const len = try std.unicode.utf16leToUtf8(buffer, string.Buffer[0 .. string.Length / 2]);
+                    return if (len > 0) buffer[0..len] else null;
+                },
+                .NOT_IMPLEMENTED => return error.Unsupported,
+                else => |err| return os.windows.unexpectedStatus(err),
+            }
         },
         .macos, .ios, .watchos, .tvos => if (use_pthreads) {
             const err = std.c.pthread_getname_np(self.getHandle(), buffer.ptr, max_name_len + 1);
diff --git a/lib/std/os/windows.zig b/lib/std/os/windows.zig
index 31745cfe82..b1a2262083 100644
--- a/lib/std/os/windows.zig
+++ b/lib/std/os/windows.zig
@@ -2042,21 +2042,6 @@ pub fn unexpectedStatus(status: NTSTATUS) std.os.UnexpectedError {
     return error.Unexpected;
 }
 
-pub fn SetThreadDescription(hThread: HANDLE, lpThreadDescription: LPCWSTR) !void {
-    if (kernel32.SetThreadDescription(hThread, lpThreadDescription) == 0) {
-        switch (kernel32.GetLastError()) {
-            else => |err| return unexpectedError(err),
-        }
-    }
-}
-pub fn GetThreadDescription(hThread: HANDLE, ppszThreadDescription: *LPWSTR) !void {
-    if (kernel32.GetThreadDescription(hThread, ppszThreadDescription) == 0) {
-        switch (kernel32.GetLastError()) {
-            else => |err| return unexpectedError(err),
-        }
-    }
-}
-
 pub const Win32Error = @import("windows/win32error.zig").Win32Error;
 pub const NTSTATUS = @import("windows/ntstatus.zig").NTSTATUS;
 pub const LANG = @import("windows/lang.zig");
diff --git a/lib/std/os/windows/kernel32.zig b/lib/std/os/windows/kernel32.zig
index e947d1e505..f68ac87896 100644
--- a/lib/std/os/windows/kernel32.zig
+++ b/lib/std/os/windows/kernel32.zig
@@ -402,6 +402,3 @@ pub extern "kernel32" fn SleepConditionVariableSRW(
 pub extern "kernel32" fn TryAcquireSRWLockExclusive(s: *SRWLOCK) callconv(WINAPI) BOOLEAN;
 pub extern "kernel32" fn AcquireSRWLockExclusive(s: *SRWLOCK) callconv(WINAPI) void;
 pub extern "kernel32" fn ReleaseSRWLockExclusive(s: *SRWLOCK) callconv(WINAPI) void;
-
-pub extern "kernel32" fn SetThreadDescription(hThread: HANDLE, lpThreadDescription: LPCWSTR) callconv(WINAPI) HRESULT;
-pub extern "kernel32" fn GetThreadDescription(hThread: HANDLE, ppszThreadDescription: *LPWSTR) callconv(WINAPI) HRESULT;
diff --git a/lib/std/os/windows/ntdll.zig b/lib/std/os/windows/ntdll.zig
index 41c1a905ec..26cc19935f 100644
--- a/lib/std/os/windows/ntdll.zig
+++ b/lib/std/os/windows/ntdll.zig
@@ -23,23 +23,88 @@ const FILE_BASIC_INFORMATION = windows.FILE_BASIC_INFORMATION;
 const SIZE_T = windows.SIZE_T;
 const CURDIR = windows.CURDIR;
 
-pub extern "NtDll" fn RtlGetVersion(
+pub const THREADINFOCLASS = enum(c_int) {
+    ThreadBasicInformation,
+    ThreadTimes,
+    ThreadPriority,
+    ThreadBasePriority,
+    ThreadAffinityMask,
+    ThreadImpersonationToken,
+    ThreadDescriptorTableEntry,
+    ThreadEnableAlignmentFaultFixup,
+    ThreadEventPair_Reusable,
+    ThreadQuerySetWin32StartAddress,
+    ThreadZeroTlsCell,
+    ThreadPerformanceCount,
+    ThreadAmILastThread,
+    ThreadIdealProcessor,
+    ThreadPriorityBoost,
+    ThreadSetTlsArrayAddress,
+    ThreadIsIoPending,
+    // Windows 2000+ from here
+    ThreadHideFromDebugger,
+    // Windows XP+ from here
+    ThreadBreakOnTermination,
+    ThreadSwitchLegacyState,
+    ThreadIsTerminated,
+    // Windows Vista+ from here
+    ThreadLastSystemCall,
+    ThreadIoPriority,
+    ThreadCycleTime,
+    ThreadPagePriority,
+    ThreadActualBasePriority,
+    ThreadTebInformation,
+    ThreadCSwitchMon,
+    // Windows 7+ from here
+    ThreadCSwitchPmu,
+    ThreadWow64Context,
+    ThreadGroupInformation,
+    ThreadUmsInformation,
+    ThreadCounterProfiling,
+    ThreadIdealProcessorEx,
+    // Windows 8+ from here
+    ThreadCpuAccountingInformation,
+    // Windows 8.1+ from here
+    ThreadSuspendCount,
+    // Windows 10+ from here
+    ThreadHeterogeneousCpuPolicy,
+    ThreadContainerId,
+    ThreadNameInformation,
+    ThreadSelectedCpuSets,
+    ThreadSystemThreadInformation,
+    ThreadActualGroupAffinity,
+};
+pub extern "ntdll" fn NtQueryInformationThread(
+    ThreadHandle: HANDLE,
+    ThreadInformationClass: THREADINFOCLASS,
+    ThreadInformation: *anyopaque,
+    ThreadInformationLength: ULONG,
+    ReturnLength: ?*ULONG,
+) callconv(WINAPI) NTSTATUS;
+pub extern "ntdll" fn NtSetInformationThread(
+    ThreadHandle: HANDLE,
+    ThreadInformationClass: THREADINFOCLASS,
+    ThreadInformation: *const anyopaque,
+    ThreadInformationLength: ULONG,
+) callconv(WINAPI) NTSTATUS;
+
+pub extern "ntdll" fn RtlGetVersion(
     lpVersionInformation: *RTL_OSVERSIONINFOW,
 ) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn RtlCaptureStackBackTrace(
+pub extern "ntdll" fn RtlCaptureStackBackTrace(
     FramesToSkip: DWORD,
     FramesToCapture: DWORD,
     BackTrace: **anyopaque,
     BackTraceHash: ?*DWORD,
 ) callconv(WINAPI) WORD;
-pub extern "NtDll" fn NtQueryInformationFile(
+pub extern "ntdll" fn NtQueryInformationFile(
     FileHandle: HANDLE,
     IoStatusBlock: *IO_STATUS_BLOCK,
     FileInformation: *anyopaque,
     Length: ULONG,
     FileInformationClass: FILE_INFORMATION_CLASS,
 ) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn NtSetInformationFile(
+pub extern "ntdll" fn NtSetInformationFile(
     FileHandle: HANDLE,
     IoStatusBlock: *IO_STATUS_BLOCK,
     FileInformation: PVOID,
@@ -47,12 +112,12 @@ pub extern "NtDll" fn NtSetInformationFile(
     FileInformationClass: FILE_INFORMATION_CLASS,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtQueryAttributesFile(
+pub extern "ntdll" fn NtQueryAttributesFile(
     ObjectAttributes: *OBJECT_ATTRIBUTES,
     FileAttributes: *FILE_BASIC_INFORMATION,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtCreateFile(
+pub extern "ntdll" fn NtCreateFile(
     FileHandle: *HANDLE,
     DesiredAccess: ACCESS_MASK,
     ObjectAttributes: *OBJECT_ATTRIBUTES,
@@ -65,7 +130,7 @@ pub extern "NtDll" fn NtCreateFile(
     EaBuffer: ?*anyopaque,
     EaLength: ULONG,
 ) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn NtDeviceIoControlFile(
+pub extern "ntdll" fn NtDeviceIoControlFile(
     FileHandle: HANDLE,
     Event: ?HANDLE,
     ApcRoutine: ?IO_APC_ROUTINE,
@@ -77,7 +142,7 @@ pub extern "NtDll" fn NtDeviceIoControlFile(
     OutputBuffer: ?PVOID,
     OutputBufferLength: ULONG,
 ) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn NtFsControlFile(
+pub extern "ntdll" fn NtFsControlFile(
     FileHandle: HANDLE,
     Event: ?HANDLE,
     ApcRoutine: ?IO_APC_ROUTINE,
@@ -89,16 +154,16 @@ pub extern "NtDll" fn NtFsControlFile(
     OutputBuffer: ?PVOID,
     OutputBufferLength: ULONG,
 ) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn NtClose(Handle: HANDLE) callconv(WINAPI) NTSTATUS;
-pub extern "NtDll" fn RtlDosPathNameToNtPathName_U(
+pub extern "ntdll" fn NtClose(Handle: HANDLE) callconv(WINAPI) NTSTATUS;
+pub extern "ntdll" fn RtlDosPathNameToNtPathName_U(
     DosPathName: [*:0]const u16,
     NtPathName: *UNICODE_STRING,
     NtFileNamePart: ?*?[*:0]const u16,
     DirectoryInfo: ?*CURDIR,
 ) callconv(WINAPI) BOOL;
-pub extern "NtDll" fn RtlFreeUnicodeString(UnicodeString: *UNICODE_STRING) callconv(WINAPI) void;
+pub extern "ntdll" fn RtlFreeUnicodeString(UnicodeString: *UNICODE_STRING) callconv(WINAPI) void;
 
-pub extern "NtDll" fn NtQueryDirectoryFile(
+pub extern "ntdll" fn NtQueryDirectoryFile(
     FileHandle: HANDLE,
     Event: ?HANDLE,
     ApcRoutine: ?IO_APC_ROUTINE,
@@ -112,30 +177,30 @@ pub extern "NtDll" fn NtQueryDirectoryFile(
     RestartScan: BOOLEAN,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtCreateKeyedEvent(
+pub extern "ntdll" fn NtCreateKeyedEvent(
     KeyedEventHandle: *HANDLE,
     DesiredAccess: ACCESS_MASK,
     ObjectAttributes: ?PVOID,
     Flags: ULONG,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtReleaseKeyedEvent(
+pub extern "ntdll" fn NtReleaseKeyedEvent(
     EventHandle: ?HANDLE,
     Key: ?*const anyopaque,
     Alertable: BOOLEAN,
     Timeout: ?*const LARGE_INTEGER,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtWaitForKeyedEvent(
+pub extern "ntdll" fn NtWaitForKeyedEvent(
     EventHandle: ?HANDLE,
     Key: ?*const anyopaque,
     Alertable: BOOLEAN,
     Timeout: ?*const LARGE_INTEGER,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn RtlSetCurrentDirectory_U(PathName: *UNICODE_STRING) callconv(WINAPI) NTSTATUS;
+pub extern "ntdll" fn RtlSetCurrentDirectory_U(PathName: *UNICODE_STRING) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtQueryObject(
+pub extern "ntdll" fn NtQueryObject(
     Handle: HANDLE,
     ObjectInformationClass: OBJECT_INFORMATION_CLASS,
     ObjectInformation: PVOID,
@@ -143,22 +208,22 @@ pub extern "NtDll" fn NtQueryObject(
     ReturnLength: ?*ULONG,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn RtlWakeAddressAll(
+pub extern "ntdll" fn RtlWakeAddressAll(
     Address: ?*const anyopaque,
 ) callconv(WINAPI) void;
 
-pub extern "NtDll" fn RtlWakeAddressSingle(
+pub extern "ntdll" fn RtlWakeAddressSingle(
     Address: ?*const anyopaque,
 ) callconv(WINAPI) void;
 
-pub extern "NtDll" fn RtlWaitOnAddress(
+pub extern "ntdll" fn RtlWaitOnAddress(
     Address: ?*const anyopaque,
     CompareAddress: ?*const anyopaque,
     AddressSize: SIZE_T,
     Timeout: ?*const LARGE_INTEGER,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtLockFile(
+pub extern "ntdll" fn NtLockFile(
     FileHandle: HANDLE,
     Event: ?HANDLE,
     ApcRoutine: ?*IO_APC_ROUTINE,
@@ -171,7 +236,7 @@ pub extern "NtDll" fn NtLockFile(
     ExclusiveLock: BOOLEAN,
 ) callconv(WINAPI) NTSTATUS;
 
-pub extern "NtDll" fn NtUnlockFile(
+pub extern "ntdll" fn NtUnlockFile(
     FileHandle: HANDLE,
     IoStatusBlock: *IO_STATUS_BLOCK,
     ByteOffset: *const LARGE_INTEGER,