commit e6accb2757912d73f65a7b74cd1526d433e4b9f1
author Mike Frysinger <vapier@chromium.org> Tue Sep 26 13:04:04 2023 -0400
committer Chromeos LUCI <chromeos-scoped@luci-project-accounts.iam.gserviceaccount.com> Fri Oct 06 16:09:37 2023 +0000
tree 76741eafb58cc8280a13a034d103d23f6e5a460d
parent 6e88b03636a25f4ca1eddf8bc59987e1d2df19d7

utils: prctl: add {G,S}ET_NAME helpers

Just enough to help build out reading/writing strings.

BUG=None
TEST=CQ passes

Change-Id: Id46ce8449fa2f793fd7b895589a002fab4adb45a
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/4917268
Tested-by: Mike Frysinger <vapier@chromium.org>
Auto-Submit: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Cindy Lin <xcl@google.com>
Commit-Queue: Cindy Lin <xcl@google.com>

utils/prctl.py

diff --git a/utils/prctl.py b/utils/prctl.py
index d524f52..8a182c7 100644
--- a/utils/prctl.py
+++ b/utils/prctl.py
@@ -21,6 +21,12 @@
     # arg2 is int* output.
     GET_PDEATHSIG = 2
 
+    # arg2 is char* input.
+    SET_NAME = 15
+
+    # arg2 is char* output.
+    GET_NAME = 16
+
 
 class Error(Exception):
     """Base class for errors in this module."""
@@ -127,6 +133,23 @@
     return value.value
 
 
+def _set_str(option: Option, value: str) -> None:
+    """Helper for functions that have a single input string."""
+    c_str = ctypes.create_string_buffer(value.encode("utf-8"))
+    ret = _set_int(option, ctypes.byref(c_str))
+    if ret:
+        raise PrctlError(option, ret, [value])
+
+
+def _get_str(option: Option, length: int) -> str:
+    """Helper for functions that have a single output string."""
+    c_str = ctypes.create_string_buffer(length)
+    ret = prctl(Option.GET_NAME, ctypes.byref(c_str))
+    if ret:
+        raise PrctlError(option, ret)
+    return c_str.value.decode("utf-8")
+
+
 def set_pdeathsig(value: int) -> None:
     """SET_PDEATHSIG wrapper."""
     _set_int(Option.SET_PDEATHSIG, value)
@@ -135,3 +158,14 @@
 def get_pdeathsig() -> int:
     """GET_PDEATHSIG wrapper."""
     return _get_int(Option.GET_PDEATHSIG)
+
+
+def set_name(name: str) -> None:
+    """SET_NAME (thread name) wrapper."""
+    _set_str(Option.SET_NAME, name)
+
+
+def get_name() -> str:
+    """GET_NAME (thread name) wrapper."""
+    # Return is 16 bytes, and it's always NUL terminated.
+    return _get_str(Option.GET_NAME, 16)

utils/prctl_unittest.py

diff --git a/utils/prctl_unittest.py b/utils/prctl_unittest.py
index dd84089..a3e8ada 100644
--- a/utils/prctl_unittest.py
+++ b/utils/prctl_unittest.py
@@ -50,3 +50,13 @@
     assert prctl.get_pdeathsig() == signal.SIGINT
     # Restore the setting.
     prctl.set_pdeathsig(orig)
+
+
+def test_name():
+    """Check (thread) name helpers."""
+    assert prctl.set_name("foo") is None
+    assert prctl.get_name() == "foo"
+
+    # Check truncation.
+    assert prctl.set_name("1234567890" * 3) is None
+    assert prctl.get_name() == "123456789012345"

utils/proctitle_util.py

diff --git a/utils/proctitle_util.py b/utils/proctitle_util.py
index 2cd2801..7898824 100644
--- a/utils/proctitle_util.py
+++ b/utils/proctitle_util.py
@@ -15,9 +15,11 @@
     from setproctitle import getproctitle
     from setproctitle import setproctitle
 except ImportError:
-    # Module not available -> can't do anything.
-    getproctitle = lambda: None
-    setproctitle = lambda _x: None
+    # Module not available -> use basic prctl API.
+    from chromite.utils import prctl
+
+    getproctitle = prctl.get_name
+    setproctitle = prctl.set_name
 
 
 # Used with the settitle helper below.