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>

diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
index 3e164a12581dd13e2b12bd08686f98306f210353..214e1ba0b53dd29dbb1b3a33aae7d27ddcd24d84 100644 (file)
--- a/Lib/test/test_threading.py
+++ b/Lib/test/test_threading.py
@@ -2130,6 +2130,15 @@ class MiscTestCase(unittest.TestCase):
 
             # 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 @@ class MiscTestCase(unittest.TestCase):
             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

diff --git a/Modules/_threadmodule.c b/Modules/_threadmodule.c
index d19ae326bd6b4822a7b3332dc88f6bf6e02969cd..586ed368024fb12ad41eed1992048d419dd14231 100644 (file)
--- a/Modules/_threadmodule.c
+++ b/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;
 }
 

diff --git a/Modules/clinic/_threadmodule.c.h b/Modules/clinic/_threadmodule.c.h
index 8f0507d40285b3048c160ee4209381087ee9a398..09b7afebd6d8d981211cffc81829c02ef3abb1b3 100644 (file)
--- a/Modules/clinic/_threadmodule.c.h
+++ b/Modules/clinic/_threadmodule.c.h
@@ -8,7 +8,7 @@ preserve
 #endif
 #include "pycore_modsupport.h"    // _PyArg_UnpackKeywords()
 
-#if defined(HAVE_PTHREAD_GETNAME_NP)
+#if (defined(HAVE_PTHREAD_GETNAME_NP) || defined(MS_WINDOWS))
 
 PyDoc_STRVAR(_thread__get_name__doc__,
 "_get_name($module, /)\n"
@@ -28,9 +28,9 @@ _thread__get_name(PyObject *module, PyObject *Py_UNUSED(ignored))
     return _thread__get_name_impl(module);
 }
 
-#endif /* defined(HAVE_PTHREAD_GETNAME_NP) */
+#endif /* (defined(HAVE_PTHREAD_GETNAME_NP) || defined(MS_WINDOWS)) */
 
-#if defined(HAVE_PTHREAD_SETNAME_NP)
+#if (defined(HAVE_PTHREAD_SETNAME_NP) || defined(MS_WINDOWS))
 
 PyDoc_STRVAR(_thread_set_name__doc__,
 "set_name($module, /, name)\n"
@@ -92,7 +92,7 @@ exit:
     return return_value;
 }
 
-#endif /* defined(HAVE_PTHREAD_SETNAME_NP) */
+#endif /* (defined(HAVE_PTHREAD_SETNAME_NP) || defined(MS_WINDOWS)) */
 
 #ifndef _THREAD__GET_NAME_METHODDEF
     #define _THREAD__GET_NAME_METHODDEF
@@ -101,4 +101,4 @@ exit:
 #ifndef _THREAD_SET_NAME_METHODDEF
     #define _THREAD_SET_NAME_METHODDEF
 #endif /* !defined(_THREAD_SET_NAME_METHODDEF) */
-/*[clinic end generated code: output=b5cb85aaccc45bf6 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=6e88ef6b126cece8 input=a9049054013a1b77]*/

diff --git a/PC/pyconfig.h.in b/PC/pyconfig.h.in
index 010f5fe56466303c034030f82ec95931ea782c16..837461d0e884bcad4a84c2d9197eaad6ae15c962 100644 (file)
--- a/PC/pyconfig.h.in
+++ b/PC/pyconfig.h.in
@@ -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 */