Merge pull request #2470 from jakkdl/named_threads

Add name parameter to to_thread.run_sync [SQUASH]

Author: Zac-HD

newsfragments/1148.feature.rst

@@ -0,0 +1 @@
+Added support for naming threads created with `trio.to_thread.run_sync`, requires pthreads so is only available on POSIX platforms with glibc installed.

trio/_core/_thread_cache.py

@@ -2,8 +2,69 @@
 import traceback
 from threading import Thread, Lock
 import outcome
+import ctypes
+import ctypes.util
 from itertools import count
 
+from typing import Callable, Optional, Tuple
+from functools import partial
+
+
+def _to_os_thread_name(name: str) -> bytes:
+    # ctypes handles the trailing \00
+    return name.encode("ascii", errors="replace")[:15]
+
+
+# used to construct the method used to set os thread name, or None, depending on platform.
+# called once on import
+def get_os_thread_name_func() -> Optional[Callable[[Optional[int], str], None]]:
+    def namefunc(setname: Callable[[int, bytes], int], ident: Optional[int], name: str):
+        # Thread.ident is None "if it has not been started". Unclear if that can happen
+        # with current usage.
+        if ident is not None:  # pragma: no cover
+            setname(ident, _to_os_thread_name(name))
+
+    # namefunc on mac also takes an ident, even if pthread_setname_np doesn't/can't use it
+    # so the caller don't need to care about platform.
+    def darwin_namefunc(
+        setname: Callable[[bytes], int], ident: Optional[int], name: str
+    ):
+        # I don't know if Mac can rename threads that hasn't been started, but default
+        # to no to be on the safe side.
+        if ident is not None:  # pragma: no cover
+            setname(_to_os_thread_name(name))
+
+    # find the pthread library
+    # this will fail on windows
+    libpthread_path = ctypes.util.find_library("pthread")
+    if not libpthread_path:
+        return None
+    libpthread = ctypes.CDLL(libpthread_path)
+
+    # get the setname method from it
+    # afaik this should never fail
+    pthread_setname_np = getattr(libpthread, "pthread_setname_np", None)
+    if pthread_setname_np is None:  # pragma: no cover
+        return None
+
+    # specify function prototype
+    pthread_setname_np.restype = ctypes.c_int
+
+    # on mac OSX pthread_setname_np does not take a thread id,
+    # it only lets threads name themselves, which is not a problem for us.
+    # Just need to make sure to call it correctly
+    if sys.platform == "darwin":
+        pthread_setname_np.argtypes = [ctypes.c_char_p]
+        return partial(darwin_namefunc, pthread_setname_np)
+
+    # otherwise assume linux parameter conventions. Should also work on *BSD
+    pthread_setname_np.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
+    return partial(namefunc, pthread_setname_np)
+
+
+# construct os thread name method
+set_os_thread_name = get_os_thread_name_func()
+
 # The "thread cache" is a simple unbounded thread pool, i.e., it automatically
 # spawns as many threads as needed to handle all the requests its given. Its
 # only purpose is to cache worker threads so that they don't have to be
@@ -44,7 +105,7 @@
 
 class WorkerThread:
     def __init__(self, thread_cache):
-        self._job = None
+        self._job: Optional[Tuple[Callable, Callable, str]] = None
         self._thread_cache = thread_cache
         # This Lock is used in an unconventional way.
         #
@@ -54,16 +115,34 @@ def __init__(self, thread_cache):
         # Initially we have no job, so it starts out in locked state.
         self._worker_lock = Lock()
         self._worker_lock.acquire()
-        thread = Thread(target=self._work, daemon=True)
-        thread.name = f"Trio worker thread {next(name_counter)}"
-        thread.start()
+        self._default_name = f"Trio thread {next(name_counter)}"
+
+        self._thread = Thread(target=self._work, name=self._default_name, daemon=True)
+
+        if set_os_thread_name:
+            set_os_thread_name(self._thread.ident, self._default_name)
+        self._thread.start()
 
     def _handle_job(self):
         # Handle job in a separate method to ensure user-created
         # objects are cleaned up in a consistent manner.
-        fn, deliver = self._job
+        assert self._job is not None
+        fn, deliver, name = self._job
         self._job = None
+
+        # set name
+        if name is not None:
+            self._thread.name = name
+            if set_os_thread_name:
+                set_os_thread_name(self._thread.ident, name)
         result = outcome.capture(fn)
+
+        # reset name if it was changed
+        if name is not None:
+            self._thread.name = self._default_name
+            if set_os_thread_name:
+                set_os_thread_name(self._thread.ident, self._default_name)
+
         # Tell the cache that we're available to be assigned a new
         # job. We do this *before* calling 'deliver', so that if
         # 'deliver' triggers a new job, it can be assigned to us
@@ -102,19 +181,19 @@ class ThreadCache:
     def __init__(self):
         self._idle_workers = {}
 
-    def start_thread_soon(self, fn, deliver):
+    def start_thread_soon(self, fn, deliver, name: Optional[str] = None):
         try:
             worker, _ = self._idle_workers.popitem()
         except KeyError:
             worker = WorkerThread(self)
-        worker._job = (fn, deliver)
+        worker._job = (fn, deliver, name)
         worker._worker_lock.release()
 
 
 THREAD_CACHE = ThreadCache()
 
 
-def start_thread_soon(fn, deliver):
+def start_thread_soon(fn, deliver, name: Optional[str] = None):
     """Runs ``deliver(outcome.capture(fn))`` in a worker thread.
 
     Generally ``fn`` does some blocking work, and ``deliver`` delivers the
@@ -174,4 +253,4 @@ def start_thread_soon(fn, deliver):
         limit how many threads they're using then it's polite to respect that.
 
     """
-    THREAD_CACHE.start_thread_soon(fn, deliver)
+    THREAD_CACHE.start_thread_soon(fn, deliver, name)

trio/_threads.py

@@ -1,24 +1,25 @@
 import contextvars
-import threading
-import queue as stdlib_queue
 import functools
+import inspect
+import queue as stdlib_queue
+import threading
 from itertools import count
+from typing import Optional
 
 import attr
-import inspect
 import outcome
 from sniffio import current_async_library_cvar
 
 import trio
 
-from ._sync import CapacityLimiter
 from ._core import (
-    enable_ki_protection,
-    disable_ki_protection,
     RunVar,
     TrioToken,
+    disable_ki_protection,
+    enable_ki_protection,
     start_thread_soon,
 )
+from ._sync import CapacityLimiter
 from ._util import coroutine_or_error
 
 # Global due to Threading API, thread local storage for trio token
@@ -57,7 +58,9 @@ class ThreadPlaceholder:
 
 
 @enable_ki_protection
-async def to_thread_run_sync(sync_fn, *args, cancellable=False, limiter=None):
+async def to_thread_run_sync(
+    sync_fn, *args, thread_name: Optional[str] = None, cancellable=False, limiter=None
+):
     """Convert a blocking operation into an async operation using a thread.
 
     These two lines are equivalent::
@@ -79,6 +82,12 @@ async def to_thread_run_sync(sync_fn, *args, cancellable=False, limiter=None):
           arguments, use :func:`functools.partial`.
       cancellable (bool): Whether to allow cancellation of this operation. See
           discussion below.
+      thread_name (str): Optional string to set the name of the thread.
+          Will always set `threading.Thread.name`, but only set the os name
+          if pthread.h is available (i.e. most POSIX installations).
+          pthread names are limited to 15 characters, and can be read from
+          ``/proc/<PID>/task/<SPID>/comm`` or with ``ps -eT``, among others.
+          Defaults to ``{sync_fn.__name__|None} from {trio.lowlevel.current_task().name}``.
       limiter (None, or CapacityLimiter-like object):
           An object used to limit the number of simultaneous threads. Most
           commonly this will be a `~trio.CapacityLimiter`, but it could be
@@ -166,6 +175,9 @@ def do_release_then_return_result():
 
     current_trio_token = trio.lowlevel.current_trio_token()
 
+    if thread_name is None:
+        thread_name = f"{getattr(sync_fn, '__name__', None)} from {trio.lowlevel.current_task().name}"
+
     def worker_fn():
         current_async_library_cvar.set(None)
         TOKEN_LOCAL.token = current_trio_token
@@ -198,7 +210,9 @@ def deliver_worker_fn_result(result):
 
     await limiter.acquire_on_behalf_of(placeholder)
     try:
-        start_thread_soon(contextvars_aware_worker_fn, deliver_worker_fn_result)
+        start_thread_soon(
+            contextvars_aware_worker_fn, deliver_worker_fn_result, thread_name
+        )
     except:
         limiter.release_on_behalf_of(placeholder)
         raise

trio/tests/test_threads.py

@@ -1,25 +1,28 @@
 import contextvars
-import threading
 import queue as stdlib_queue
+import re
+import sys
+import threading
 import time
 import weakref
+from functools import partial
+from typing import Callable, Optional
 
 import pytest
 from sniffio import current_async_library_cvar
+
 from trio._core import TrioToken, current_trio_token
 
-from .. import _core
-from .. import Event, CapacityLimiter, sleep
-from ..testing import wait_all_tasks_blocked
+from .. import CapacityLimiter, Event, _core, sleep
+from .._core.tests.test_ki import ki_self
 from .._core.tests.tutil import buggy_pypy_asyncgens
 from .._threads import (
-    to_thread_run_sync,
     current_default_thread_limiter,
     from_thread_run,
     from_thread_run_sync,
+    to_thread_run_sync,
 )
-
-from .._core.tests.test_ki import ki_self
+from ..testing import wait_all_tasks_blocked
 
 
 async def test_do_in_trio_thread():
@@ -164,6 +167,117 @@ async def main():
     assert record == ["sleeping", "cancelled"]
 
 
+async def test_named_thread():
+    ending = " from trio.tests.test_threads.test_named_thread"
+
+    def inner(name="inner" + ending) -> threading.Thread:
+        assert threading.current_thread().name == name
+        return threading.current_thread()
+
+    def f(name: str) -> Callable[[None], threading.Thread]:
+        return partial(inner, name)
+
+    # test defaults
+    await to_thread_run_sync(inner)
+    await to_thread_run_sync(inner, thread_name=None)
+
+    # functools.partial doesn't have __name__, so defaults to None
+    await to_thread_run_sync(f("None" + ending))
+
+    # test that you can set a custom name, and that it's reset afterwards
+    async def test_thread_name(name: str):
+        thread = await to_thread_run_sync(f(name), thread_name=name)
+        assert re.match("Trio thread [0-9]*", thread.name)
+
+    await test_thread_name("")
+    await test_thread_name("fobiedoo")
+    await test_thread_name("name_longer_than_15_characters")
+
+    await test_thread_name("💙")
+
+
+def _get_thread_name(ident: Optional[int] = None) -> Optional[str]:
+    import ctypes
+    import ctypes.util
+
+    libpthread_path = ctypes.util.find_library("pthread")
+    if not libpthread_path:
+        print(f"no pthread on {sys.platform})")
+        return None
+    libpthread = ctypes.CDLL(libpthread_path)
+
+    pthread_getname_np = getattr(libpthread, "pthread_getname_np", None)
+
+    # this should never fail on any platforms afaik
+    assert pthread_getname_np
+
+    # thankfully getname signature doesn't differ between platforms
+    pthread_getname_np.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_size_t,
+    ]
+    pthread_getname_np.restype = ctypes.c_int
+
+    name_buffer = ctypes.create_string_buffer(b"", size=16)
+    if ident is None:
+        ident = threading.get_ident()
+    assert pthread_getname_np(ident, name_buffer, 16) == 0
+    try:
+        return name_buffer.value.decode()
+    except UnicodeDecodeError as e:  # pragma: no cover
+        # used for debugging when testing via CI
+        pytest.fail(f"value: {name_buffer.value!r}, exception: {e}")
+
+
+# test os thread naming
+# this depends on pthread being available, which is the case on 99.9% of linux machines
+# and most mac machines. So unless the platform is linux it will just skip
+# in case it fails to fetch the os thread name.
+async def test_named_thread_os():
+    def inner(name) -> threading.Thread:
+        os_thread_name = _get_thread_name()
+        if os_thread_name is None and sys.platform != "linux":
+            pytest.skip(f"no pthread OS support on {sys.platform}")
+        else:
+            assert os_thread_name == name[:15]
+
+        return threading.current_thread()
+
+    def f(name: str) -> Callable[[None], threading.Thread]:
+        return partial(inner, name)
+
+    # test defaults
+    default = "None from trio.tests.test_threads.test_named_thread"
+    await to_thread_run_sync(f(default))
+    await to_thread_run_sync(f(default), thread_name=None)
+
+    # test that you can set a custom name, and that it's reset afterwards
+    async def test_thread_name(name: str, expected: Optional[str] = None):
+        if expected is None:
+            expected = name
+        thread = await to_thread_run_sync(f(expected), thread_name=name)
+
+        os_thread_name = _get_thread_name(thread.ident)
+        assert os_thread_name is not None, "should skip earlier if this is the case"
+        assert re.match("Trio thread [0-9]*", os_thread_name)
+
+    await test_thread_name("")
+    await test_thread_name("fobiedoo")
+    await test_thread_name("name_longer_than_15_characters")
+
+    await test_thread_name("💙", expected="?")
+
+
+async def test_has_pthread_setname_np():
+    from trio._core._thread_cache import get_os_thread_name_func
+
+    k = get_os_thread_name_func()
+    if k is None:
+        assert sys.platform != "linux"
+        pytest.skip(f"no pthread_setname_np on {sys.platform}")
+
+
 async def test_run_in_worker_thread():
     trio_thread = threading.current_thread()