gh-59705: Implement _thread.set_name() on Windows (#128675)

Implement set_name() with SetThreadDescription() and _get_name() with
GetThreadDescription(). If SetThreadDescription() or
GetThreadDescription() is not available in kernelbase.dll, delete the
method when the _thread module is imported.

Truncate the thread name to 32766 characters.

Co-authored-by: Eryk Sun <eryksun@gmail.com>

Modules/_threadmodule.c

@@ -47,6 +47,14 @@ get_thread_state(PyObject *module)
 }
 
 
+#ifdef MS_WINDOWS
+typedef HRESULT (WINAPI *PF_GET_THREAD_DESCRIPTION)(HANDLE, PCWSTR*);
+typedef HRESULT (WINAPI *PF_SET_THREAD_DESCRIPTION)(HANDLE, PCWSTR);
+static PF_GET_THREAD_DESCRIPTION pGetThreadDescription = NULL;
+static PF_SET_THREAD_DESCRIPTION pSetThreadDescription = NULL;
+#endif
+
+
 /*[clinic input]
 module _thread
 [clinic start generated code]*/
@@ -2368,7 +2376,7 @@ Internal only. Return a non-zero integer that uniquely identifies the main threa
 of the main interpreter.");
 
 
-#ifdef HAVE_PTHREAD_GETNAME_NP
+#if defined(HAVE_PTHREAD_GETNAME_NP) || defined(MS_WINDOWS)
 /*[clinic input]
 _thread._get_name
 
@@ -2379,6 +2387,7 @@ static PyObject *
 _thread__get_name_impl(PyObject *module)
 /*[clinic end generated code: output=20026e7ee3da3dd7 input=35cec676833d04c8]*/
 {
+#ifndef MS_WINDOWS
     // Linux and macOS are limited to respectively 16 and 64 bytes
     char name[100];
     pthread_t thread = pthread_self();
@@ -2393,11 +2402,26 @@ _thread__get_name_impl(PyObject *module)
 #else
     return PyUnicode_DecodeFSDefault(name);
 #endif
+#else
+    // Windows implementation
+    assert(pGetThreadDescription != NULL);
+
+    wchar_t *name;
+    HRESULT hr = pGetThreadDescription(GetCurrentThread(), &name);
+    if (FAILED(hr)) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+
+    PyObject *name_obj = PyUnicode_FromWideChar(name, -1);
+    LocalFree(name);
+    return name_obj;
+#endif
 }
 #endif  // HAVE_PTHREAD_GETNAME_NP
 
 
-#ifdef HAVE_PTHREAD_SETNAME_NP
+#if defined(HAVE_PTHREAD_SETNAME_NP) || defined(MS_WINDOWS)
 /*[clinic input]
 _thread.set_name
 
@@ -2410,6 +2434,7 @@ static PyObject *
 _thread_set_name_impl(PyObject *module, PyObject *name_obj)
 /*[clinic end generated code: output=402b0c68e0c0daed input=7e7acd98261be82f]*/
 {
+#ifndef MS_WINDOWS
 #ifdef __sun
     // Solaris always uses UTF-8
     const char *encoding = "utf-8";
@@ -2455,6 +2480,35 @@ _thread_set_name_impl(PyObject *module, PyObject *name_obj)
         return PyErr_SetFromErrno(PyExc_OSError);
     }
     Py_RETURN_NONE;
+#else
+    // Windows implementation
+    assert(pSetThreadDescription != NULL);
+
+    Py_ssize_t len;
+    wchar_t *name = PyUnicode_AsWideCharString(name_obj, &len);
+    if (name == NULL) {
+        return NULL;
+    }
+
+    if (len > PYTHREAD_NAME_MAXLEN) {
+        // Truncate the name
+        Py_UCS4 ch = name[PYTHREAD_NAME_MAXLEN-1];
+        if (Py_UNICODE_IS_HIGH_SURROGATE(ch)) {
+            name[PYTHREAD_NAME_MAXLEN-1] = 0;
+        }
+        else {
+            name[PYTHREAD_NAME_MAXLEN] = 0;
+        }
+    }
+
+    HRESULT hr = pSetThreadDescription(GetCurrentThread(), name);
+    PyMem_Free(name);
+    if (FAILED(hr)) {
+        PyErr_SetFromWindowsErr((int)hr);
+        return NULL;
+    }
+    Py_RETURN_NONE;
+#endif
 }
 #endif  // HAVE_PTHREAD_SETNAME_NP
 
@@ -2598,6 +2652,31 @@ thread_module_exec(PyObject *module)
     }
 #endif
 
+#ifdef MS_WINDOWS
+    HMODULE kernelbase = GetModuleHandleW(L"kernelbase.dll");
+    if (kernelbase != NULL) {
+        if (pGetThreadDescription == NULL) {
+            pGetThreadDescription = (PF_GET_THREAD_DESCRIPTION)GetProcAddress(
+                                        kernelbase, "GetThreadDescription");
+        }
+        if (pSetThreadDescription == NULL) {
+            pSetThreadDescription = (PF_SET_THREAD_DESCRIPTION)GetProcAddress(
+                                        kernelbase, "SetThreadDescription");
+        }
+    }
+
+    if (pGetThreadDescription == NULL) {
+        if (PyObject_DelAttrString(module, "_get_name") < 0) {
+            return -1;
+        }
+    }
+    if (pSetThreadDescription == NULL) {
+        if (PyObject_DelAttrString(module, "set_name") < 0) {
+            return -1;
+        }
+    }
+#endif
+
     return 0;
 }