Add proc-title-template option. (#8397)

Make it possible to customize the process title, i.e. include custom
strings, immutable configuration like port, tls-port, unix socket name,
etc.

7 files changed, 244 insertions, 12 deletions

diff --git a/redis.conf b/redis.conf
index 370542344..465d56fc0 100644
--- a/redis.conf
+++ b/redis.conf
@@ -330,6 +330,23 @@ always-show-logo no
 # the process name as executed by setting the following to no.
 set-proc-title yes
 
+# When changing the process title, Redis uses the following template to construct
+# the modified title.
+#
+# Template variables are specified in curly brackets. The following variables are
+# supported:
+#
+# {title}           Name of process as executed if parent, or type of child process.
+# {listen-addr}     Bind address or '*' followed by TCP or TLS port listening on, or
+#                   Unix socket if only that's available.
+# {server-mode}     Special mode, i.e. "[sentinel]" or "[cluster]".
+# {port}            TCP port listening on, or 0.
+# {tls-port}        TLS port listening on, or 0.
+# {unixsocket}      Unix domain socket listening on, or "".
+# {config-file}     Name of configuration file used.
+#
+proc-title-template "{title} {listen-addr} {server-mode}"
+
 ################################ SNAPSHOTTING  ################################
 
 # Save the DB to disk.
diff --git a/src/config.c b/src/config.c
index 20d432573..0bd89c2b9 100644
--- a/src/config.c
+++ b/src/config.c
@@ -2236,6 +2236,25 @@ static int isValidAOFfilename(char *val, const char **err) {
     return 1;
 }
 
+/* Validate specified string is a valid proc-title-template */
+static int isValidProcTitleTemplate(char *val, const char **err) {
+    if (!validateProcTitleTemplate(val)) {
+        *err = "template format is invalid or contains unknown variables";
+        return 0;
+    }
+    return 1;
+}
+
+static int updateProcTitleTemplate(char *val, char *prev, const char **err) {
+    UNUSED(val);
+    UNUSED(prev);
+    if (redisSetProcTitle(NULL) == C_ERR) {
+        *err = "failed to set process title";
+        return 0;
+    }
+    return 1;
+}
+
 static int updateHZ(long long val, long long prev, const char **err) {
     UNUSED(prev);
     UNUSED(err);
@@ -2435,6 +2454,7 @@ standardConfig configs[] = {
     createStringConfig("aof_rewrite_cpulist", NULL, IMMUTABLE_CONFIG, EMPTY_STRING_IS_NULL, server.aof_rewrite_cpulist, NULL, NULL, NULL),
     createStringConfig("bgsave_cpulist", NULL, IMMUTABLE_CONFIG, EMPTY_STRING_IS_NULL, server.bgsave_cpulist, NULL, NULL, NULL),
     createStringConfig("ignore-warnings", NULL, MODIFIABLE_CONFIG, ALLOW_EMPTY_STRING, server.ignore_warnings, "", NULL, NULL),
+    createStringConfig("proc-title-template", NULL, MODIFIABLE_CONFIG, ALLOW_EMPTY_STRING, server.proc_title_template, CONFIG_DEFAULT_PROC_TITLE_TEMPLATE, isValidProcTitleTemplate, updateProcTitleTemplate),
 
     /* SDS Configs */
     createSDSConfig("masterauth", NULL, MODIFIABLE_CONFIG, EMPTY_STRING_IS_NULL, server.masterauth, NULL, NULL, NULL),
diff --git a/src/sds.c b/src/sds.c
index f16114471..ad30e2ad4 100644
--- a/src/sds.c
+++ b/src/sds.c
@@ -1157,12 +1157,80 @@ void *sds_malloc(size_t size) { return s_malloc(size); }
 void *sds_realloc(void *ptr, size_t size) { return s_realloc(ptr,size); }
 void sds_free(void *ptr) { s_free(ptr); }
 
+/* Perform expansion of a template string and return the result as a newly
+ * allocated sds.
+ *
+ * Template variables are specified using curly brackets, e.g. {variable}.
+ * An opening bracket can be quoted by repeating it twice.
+ */
+sds sdstemplate(const char *template, sdstemplate_callback_t cb_func, void *cb_arg)
+{
+    sds res = sdsempty();
+    const char *p = template;
+
+    while (*p) {
+        /* Find next variable, copy everything until there */
+        const char *sv = strchr(p, '{');
+        if (!sv) {
+            /* Not found: copy till rest of template and stop */
+            res = sdscat(res, p);
+            break;
+        } else if (sv > p) {
+            /* Found: copy anything up to the begining of the variable */
+            res = sdscatlen(res, p, sv - p);
+        }
+
+        /* Skip into variable name, handle premature end or quoting */
+        sv++;
+        if (!*sv) goto error;       /* Premature end of template */
+        if (*sv == '{') {
+            /* Quoted '{' */
+            p = sv + 1;
+            res = sdscat(res, "{");
+            continue;
+        }
+
+        /* Find end of variable name, handle premature end of template */
+        const char *ev = strchr(sv, '}');
+        if (!ev) goto error;
+
+        /* Pass variable name to callback and obtain value. If callback failed,
+         * abort. */
+        sds varname = sdsnewlen(sv, ev - sv);
+        sds value = cb_func(varname, cb_arg);
+        sdsfree(varname);
+        if (!value) goto error;
+
+        /* Append value to result and continue */
+        res = sdscat(res, value);
+        sdsfree(value);
+        p = ev + 1;
+    }
+
+    return res;
+
+error:
+    sdsfree(res);
+    return NULL;
+}
+
 #ifdef REDIS_TEST
 #include <stdio.h>
 #include <limits.h>
 #include "testhelp.h"
 
 #define UNUSED(x) (void)(x)
+
+static sds sdsTestTemplateCallback(sds varname, void *arg) {
+    UNUSED(arg);
+    static const char *_var1 = "variable1";
+    static const char *_var2 = "variable2";
+
+    if (!strcmp(varname, _var1)) return sdsnew("value1");
+    else if (!strcmp(varname, _var2)) return sdsnew("value2");
+    else return NULL;
+}
+
 int sdsTest(int argc, char **argv) {
     UNUSED(argc);
     UNUSED(argv);
@@ -1342,6 +1410,30 @@ int sdsTest(int argc, char **argv) {
 
             sdsfree(x);
         }
+
+        /* Simple template */
+        x = sdstemplate("v1={variable1} v2={variable2}", sdsTestTemplateCallback, NULL);
+        test_cond("sdstemplate() normal flow",
+                  memcmp(x,"v1=value1 v2=value2",19) == 0);
+        sdsfree(x);
+
+        /* Template with callback error */
+        x = sdstemplate("v1={variable1} v3={doesnotexist}", sdsTestTemplateCallback, NULL);
+        test_cond("sdstemplate() with callback error", x == NULL);
+
+        /* Template with empty var name */
+        x = sdstemplate("v1={", sdsTestTemplateCallback, NULL);
+        test_cond("sdstemplate() with empty var name", x == NULL);
+
+        /* Template with truncated var name */
+        x = sdstemplate("v1={start", sdsTestTemplateCallback, NULL);
+        test_cond("sdstemplate() with truncated var name", x == NULL);
+
+        /* Template with quoting */
+        x = sdstemplate("v1={{{variable1}} {{} v2={variable2}", sdsTestTemplateCallback, NULL);
+        test_cond("sdstemplate() with quoting",
+                  memcmp(x,"v1={value1} {} v2=value2",24) == 0);
+        sdsfree(x);
     }
     test_report();
     return 0;
diff --git a/src/sds.h b/src/sds.h
index 3a9e4cefe..85dc0b680 100644
--- a/src/sds.h
+++ b/src/sds.h
@@ -253,6 +253,14 @@ sds sdsmapchars(sds s, const char *from, const char *to, size_t setlen);
 sds sdsjoin(char **argv, int argc, char *sep);
 sds sdsjoinsds(sds *argv, int argc, const char *sep, size_t seplen);
 
+/* Callback for sdstemplate. The function gets called by sdstemplate
+ * every time a variable needs to be expanded. The variable name is
+ * provided as variable, and the callback is expected to return a
+ * substitution value. Returning a NULL indicates an error.
+ */
+typedef sds (*sdstemplate_callback_t)(const sds variable, void *arg);
+sds sdstemplate(const char *template, sdstemplate_callback_t cb_func, void *cb_arg);
+
 /* Low level functions exposed to the user API */
 sds sdsMakeRoomFor(sds s, size_t addlen);
 void sdsIncrLen(sds s, ssize_t incr);
diff --git a/src/server.c b/src/server.c
index 9135c0867..8b8da8c89 100644
--- a/src/server.c
+++ b/src/server.c
@@ -5648,20 +5648,68 @@ void redisOutOfMemoryHandler(size_t allocation_size) {
         allocation_size);
 }
 
-void redisSetProcTitle(char *title) {
+/* Callback for sdstemplate on proc-title-template. See redis.conf for
+ * supported variables.
+ */
+static sds redisProcTitleGetVariable(const sds varname, void *arg)
+{
+    if (!strcmp(varname, "title")) {
+        return sdsnew(arg);
+    } else if (!strcmp(varname, "listen-addr")) {
+        if (server.port || server.tls_port)
+            return sdscatprintf(sdsempty(), "%s:%u",
+                                server.bindaddr_count ? server.bindaddr[0] : "*",
+                                server.port ? server.port : server.tls_port);
+        else
+            return sdscatprintf(sdsempty(), "unixsocket:%s", server.unixsocket);
+    } else if (!strcmp(varname, "server-mode")) {
+        if (server.cluster_enabled) return sdsnew("[cluster]");
+        else if (server.sentinel_mode) return sdsnew("[sentinel]");
+        else return sdsempty();
+    } else if (!strcmp(varname, "config-file")) {
+        return sdsnew(server.configfile ? server.configfile : "-");
+    } else if (!strcmp(varname, "port")) {
+        return sdscatprintf(sdsempty(), "%u", server.port);
+    } else if (!strcmp(varname, "tls-port")) {
+        return sdscatprintf(sdsempty(), "%u", server.tls_port);
+    } else if (!strcmp(varname, "unixsocket")) {
+        return sdsnew(server.unixsocket);
+    } else
+        return NULL;    /* Unknown variable name */
+}
+
+/* Expand the specified proc-title-template string and return a newly
+ * allocated sds, or NULL. */
+static sds expandProcTitleTemplate(const char *template, const char *title) {
+    sds res = sdstemplate(template, redisProcTitleGetVariable, (void *) title);
+    if (!res)
+        return NULL;
+    return sdstrim(res, " ");
+}
+/* Validate the specified template, returns 1 if valid or 0 otherwise. */
+int validateProcTitleTemplate(const char *template) {
+    int ok = 1;
+    sds res = expandProcTitleTemplate(template, "");
+    if (!res)
+        return 0;
+    if (sdslen(res) == 0) ok = 0;
+    sdsfree(res);
+    return ok;
+}
+
+int redisSetProcTitle(char *title) {
 #ifdef USE_SETPROCTITLE
-    char *server_mode = "";
-    if (server.cluster_enabled) server_mode = " [cluster]";
-    else if (server.sentinel_mode) server_mode = " [sentinel]";
-
-    setproctitle("%s %s:%d%s",
-        title,
-        server.bindaddr_count ? server.bindaddr[0] : "*",
-        server.port ? server.port : server.tls_port,
-        server_mode);
+    if (!title) title = server.exec_argv[0];
+    sds proc_title = expandProcTitleTemplate(server.proc_title_template, title);
+    if (!proc_title) return C_ERR;  /* Not likely, proc_title_template is validated */
+
+    setproctitle("%s", proc_title);
+    sdsfree(proc_title);
 #else
     UNUSED(title);
 #endif
+
+    return C_OK;
 }
 
 void redisSetCpuAffinity(const char *cpulist) {
@@ -5925,7 +5973,7 @@ int main(int argc, char **argv) {
     readOOMScoreAdj();
     initServer();
     if (background || server.pidfile) createPidFile();
-    if (server.set_proc_title) redisSetProcTitle(argv[0]);
+    if (server.set_proc_title) redisSetProcTitle(NULL);
     redisAsciiArt();
     checkTcpBacklogSettings();
 
diff --git a/src/server.h b/src/server.h
index b72d7bd0d..349c887cb 100644
--- a/src/server.h
+++ b/src/server.h
@@ -115,6 +115,7 @@ typedef long long ustime_t; /* microsecond time type. */
 #define NET_ADDR_STR_LEN (NET_IP_STR_LEN+32) /* Must be enough for ip:port */
 #define CONFIG_BINDADDR_MAX 16
 #define CONFIG_MIN_RESERVED_FDS 32
+#define CONFIG_DEFAULT_PROC_TITLE_TEMPLATE "{title} {listen-addr} {server-mode}"
 
 #define ACTIVE_EXPIRE_CYCLE_SLOW 0
 #define ACTIVE_EXPIRE_CYCLE_FAST 1
@@ -1298,6 +1299,7 @@ struct redisServer {
     int supervised_mode;            /* See SUPERVISED_* */
     int daemonize;                  /* True if running as a daemon */
     int set_proc_title;             /* True if change proc title */
+    char *proc_title_template;      /* Process title template format */
     clientBufferLimitsConfig client_obuf_limits[CLIENT_TYPE_OBUF_COUNT];
     /* AOF persistence */
     int aof_enabled;                /* AOF configuration */
@@ -1749,7 +1751,8 @@ void getRandomBytes(unsigned char *p, size_t len);
 uint64_t crc64(uint64_t crc, const unsigned char *s, uint64_t l);
 void exitFromChild(int retcode);
 size_t redisPopcount(void *s, long count);
-void redisSetProcTitle(char *title);
+int redisSetProcTitle(char *title);
+int validateProcTitleTemplate(const char *template);
 int redisCommunicateSystemd(const char *sd_notify_msg);
 void redisSetCpuAffinity(const char *cpulist);
 
diff --git a/tests/unit/other.tcl b/tests/unit/other.tcl
index d98dc1bd4..a35ac1752 100644
--- a/tests/unit/other.tcl
+++ b/tests/unit/other.tcl
@@ -321,3 +321,47 @@ start_server {tags {"other"}} {
         assert_match "*table size: 8192*" [r debug HTSTATS 9]
     }
 }
+
+proc read_proc_title {pid} {
+    set fd [open "/proc/$pid/cmdline" "r"]
+    set cmdline [read $fd 1024]
+    close $fd
+
+    return $cmdline
+}
+
+start_server {tags {"other"}} {
+    test {Process title set as expected} {
+        # Test only on Linux where it's easy to get cmdline without relying on tools.
+        # Skip valgrind as it messes up the arguments.
+        set os [exec uname]
+        if {$os == "Linux" && !$::valgrind} {
+            # Set a custom template
+            r config set "proc-title-template" "TEST {title} {listen-addr} {port} {tls-port} {unixsocket} {config-file}"
+            set cmdline [read_proc_title [srv 0 pid]]
+
+            assert_equal "TEST" [lindex $cmdline 0]
+            assert_match "*/redis-server" [lindex $cmdline 1]
+            
+            if {$::tls} {
+                set expect_port 0
+                set expect_tls_port [srv 0 port]
+            } else {
+                set expect_port [srv 0 port]
+                set expect_tls_port 0
+            }
+            set port [srv 0 port]
+
+            assert_equal "$::host:$port" [lindex $cmdline 2]
+            assert_equal $expect_port [lindex $cmdline 3]
+            assert_equal $expect_tls_port [lindex $cmdline 4]
+            assert_match "*/tests/tmp/server.*/socket" [lindex $cmdline 5]
+            assert_match "*/tests/tmp/redis.conf.*" [lindex $cmdline 6]
+
+            # Try setting a bad template
+            catch {r config set "proc-title-template" "{invalid-var}"} err
+            assert_match {*template format is invalid*} $err
+        }
+    }
+}
+