pythongh-59705: Implement _thread.set_name() on Windows (python#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 <[email protected]>

Author: vstinner

Lib/test/test_threading.py

@@ -2130,6 +2130,15 @@ def test_set_name(self):
 
             # Test long non-ASCII name (truncated)
             "x" * (limit - 1) + "é€",
+
+            # Test long non-BMP names (truncated) creating surrogate pairs
+            # on Windows
+            "x" * (limit - 1) + "\U0010FFFF",
+            "x" * (limit - 2) + "\U0010FFFF" * 2,
+            "x" + "\U0001f40d" * limit,
+            "xx" + "\U0001f40d" * limit,
+            "xxx" + "\U0001f40d" * limit,
+            "xxxx" + "\U0001f40d" * limit,
         ]
         if os_helper.FS_NONASCII:
             tests.append(f"nonascii:{os_helper.FS_NONASCII}")
@@ -2146,15 +2155,31 @@ def work():
             work_name = _thread._get_name()
 
         for name in tests:
-            encoded = name.encode(encoding, "replace")
-            if b'\0' in encoded:
-                encoded = encoded.split(b'\0', 1)[0]
-            if truncate is not None:
-                encoded = encoded[:truncate]
-            if sys.platform.startswith("solaris"):
-                expected = encoded.decode("utf-8", "surrogateescape")
+            if not support.MS_WINDOWS:
+                encoded = name.encode(encoding, "replace")
+                if b'\0' in encoded:
+                    encoded = encoded.split(b'\0', 1)[0]
+                if truncate is not None:
+                    encoded = encoded[:truncate]
+                if sys.platform.startswith("solaris"):
+                    expected = encoded.decode("utf-8", "surrogateescape")
+                else:
+                    expected = os.fsdecode(encoded)
             else:
-                expected = os.fsdecode(encoded)
+                size = 0
+                chars = []
+                for ch in name:
+                    if ord(ch) > 0xFFFF:
+                        size += 2
+                    else:
+                        size += 1
+                    if size > truncate:
+                        break
+                    chars.append(ch)
+                expected = ''.join(chars)
+
+                if '\0' in expected:
+                    expected = expected.split('\0', 1)[0]
 
             with self.subTest(name=name, expected=expected):
                 work_name = None

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;
 }
 

Modules/clinic/_threadmodule.c.h

@@ -753,4 +753,8 @@ Py_NO_ENABLE_SHARED to find out.  Also support MS_NO_COREDLL for b/w compat */
 /* Define if libssl has X509_VERIFY_PARAM_set1_host and related function */
 #define HAVE_X509_VERIFY_PARAM_SET1_HOST 1
 
+// Truncate the thread name to 64 characters. The OS limit is 32766 wide
+// characters, but long names aren't of practical use.
+#define PYTHREAD_NAME_MAXLEN 32766
+
 #endif /* !Py_CONFIG_H */