[PATCH 2 of 2 V3] osutil: implement setprocname to set process title for some platforms

[Jun Wu]

# HG changeset patch
# User Jun Wu <quark at fb.com>
# Date 1478898677 0
#      Fri Nov 11 21:11:17 2016 +0000
# Node ID c4af4fec293bed4e5105442137a24c34d58e740f
# Parent  98761d64eaaf67f3bdb99f3f80a57910e2624b78
# Available At https://bitbucket.org/quark-zju/hg-draft
#              hg pull https://bitbucket.org/quark-zju/hg-draft -r c4af4fec293b
osutil: implement setprocname to set process title for some platforms

This patch adds a simple setprocname method to osutil. The operation is not
defined by any standard and is platform-specific, the current implementation
tries to cover some major platforms (ex. Linux, OS X, FreeBSD) that is
relatively easy to support. Other platforms (Windows [4], other BSDs, ...)
can be added in the future.

The current implementation supports two methods to change process title:
  a. setproctitle if available (works in FreeBSD).
  b. rewrite argv in place (works in Linux [1] and Mac OS X). [2] [3]

[1]: Linux has "prctl(PR_SET_NAME, ...)" but 1) it has 16-byte limit, which
is too small; 2) it is not quite equivalent to what we want - it changes
"/proc/self/comm", not "/proc/self/cmdline" - "comm" change won't show up
in "ps" output unless "-o comm" is used.

[2]: The implementation does not rewrite the **environ buffer like some
other implementations do, just to make the code simpler and safer. However,
this also means the buffer size we can rewrite is significantly shorter. If
we are really greedy and want the "environ" space, we can change the
implementation later.

[3]: It requires a CPython private API: Py_GetArgcArgv to get the original
argv. Unfortunately Python 3 makes a copy of argv and returns the wchar_t
version, so it is not supported for now. (if we really want to, we could
count backwards from "char **environ", given known argc and argv, not sure
if that's a good idea - probably not)

[4]: The feature is aimed to make it easier for forked command server
processes to show what they are doing. Since Windows does not support
fork(), despite it's a major platform, its support is not added in this
patch.

diff --git a/mercurial/osutil.c b/mercurial/osutil.c
--- a/mercurial/osutil.c
+++ b/mercurial/osutil.c
@@ -728,4 +728,57 @@ bail:
 
 #endif /* CMSG_LEN */
+
+#if defined(HAVE_SETPROCTITLE)
+/* setproctitle is the first choice - available in FreeBSD */
+#define SETPROCNAME_USE_SETPROCTITLE
+#elif (defined(__linux__) || defined(__APPLE__)) && PY_MAJOR_VERSION == 2
+/* rewrite the argv buffer in place - works in Linux and OS X. Py_GetArgcArgv
+ * in Python 3 returns the copied wchar_t **argv, thus unsupported. */
+#define SETPROCNAME_USE_ARGVREWRITE
+#else
+#define SETPROCNAME_USE_NONE
+#endif
+
+#ifndef SETPROCNAME_USE_NONE
+static PyObject *setprocname(PyObject *self, PyObject *args)
+{
+	const char *name = NULL;
+	if (!PyArg_ParseTuple(args, "s", &name))
+		return NULL;
+
+#if defined(SETPROCNAME_USE_SETPROCTITLE)
+	setproctitle("%s", name);
+#elif defined(SETPROCNAME_USE_ARGVREWRITE)
+	{
+		static char *argvstart = NULL;
+		static size_t argvsize = 0;
+		if (argvstart == NULL) {
+			int argc = 0, i;
+			char **argv = NULL;
+			extern void Py_GetArgcArgv(int *argc, char ***argv);
+			Py_GetArgcArgv(&argc, &argv);
+
+			/* Check the memory we can use. Typically, argv[i] and
+			 * argv[i + 1] are continuous. */
+			for (argvstart = argv[0], i = 0; i < argc; ++i) {
+				if (argv[i] > argvstart + argvsize)
+					break; /* not continuous */
+				size_t len = strlen(argv[i]);
+				argvsize = argv[i] - argv[0] + len + 1;
+			}
+		}
+
+		if (argvstart && argvsize > 1) {
+			int n = snprintf(argvstart, argvsize, "%s", name);
+			if (n >= 0 && (size_t)n < argvsize)
+				memset(argvstart + n, 0, argvsize - n);
+		}
+	}
+#endif
+
+	Py_RETURN_NONE;
+}
+#endif /* ndef SETPROCNAME_USE_NONE */
+
 #endif /* ndef _WIN32 */
 
@@ -900,5 +953,9 @@ static PyMethodDef methods[] = {
 	 "receive list of file descriptors via socket\n"},
 #endif
+#ifndef SETPROCNAME_USE_NONE
+	{"setprocname", (PyCFunction)setprocname, METH_VARARGS,
+	 "set process title (best-effort)\n"},
 #endif
+#endif /* ndef _WIN32 */
 #ifdef __APPLE__
 	{