From 2323d9ae48f983f8f9423250daf406fd5db55ea6 Mon Sep 17 00:00:00 2001
From: Patrick Hemmer <phemmer@users.noreply.github.com>
Date: Mon, 13 Nov 2017 13:59:27 -0500
Subject: [PATCH] Add systemd unit pid and cgroup matching to procstat (#3459)

---
 plugins/inputs/procstat/README.md   | 12 ++---
 plugins/inputs/procstat/procstat.go | 68 +++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/plugins/inputs/procstat/README.md b/plugins/inputs/procstat/README.md
index 57ede8391..00820be9a 100644
--- a/plugins/inputs/procstat/README.md
+++ b/plugins/inputs/procstat/README.md
@@ -6,11 +6,11 @@ The procstat plugin can be used to monitor system resource usage by an
 individual process using their /proc data.
 
 Processes can be specified either by pid file, by executable name, by command
-line pattern matching, or by username (in this order or priority. Procstat
-plugin will use `pgrep` when executable name is provided to obtain the pid.
-Procstat plugin will transmit IO, memory, cpu, file descriptor related
-measurements for every process specified. A prefix can be set to isolate
-individual process specific measurements.
+line pattern matching, by username, by systemd unit name, or by cgroup name/path
+(in this order or priority). Procstat plugin will use `pgrep` when executable
+name is provided to obtain the pid. Procstat plugin will transmit IO, memory,
+cpu, file descriptor related measurements for every process specified. A prefix
+can be set to isolate individual process specific measurements.
 
 The plugin will tag processes according to how they are specified in the configuration. If a pid file is used, a "pidfile" tag will be generated.
 On the other hand, if an executable is used an "exe" tag will be generated. Possible tag names:
@@ -19,6 +19,8 @@ On the other hand, if an executable is used an "exe" tag will be generated. Poss
 * exe
 * pattern
 * user
+* systemd_unit
+* cgroup
 
 Additionally the plugin will tag processes by their PID (pid_tag = true in the config) and their process name:
 
diff --git a/plugins/inputs/procstat/procstat.go b/plugins/inputs/procstat/procstat.go
index a216e0e38..5df7d05e4 100644
--- a/plugins/inputs/procstat/procstat.go
+++ b/plugins/inputs/procstat/procstat.go
@@ -1,7 +1,10 @@
 package procstat
 
 import (
+	"bytes"
 	"fmt"
+	"io/ioutil"
+	"os/exec"
 	"strconv"
 	"time"
 
@@ -24,6 +27,8 @@ type Procstat struct {
 	Prefix      string
 	ProcessName string
 	User        string
+	SystemdUnit string
+	CGroup      string `toml:"cgroup"`
 	PidTag      bool
 
 	pidFinder       PIDFinder
@@ -42,6 +47,10 @@ var sampleConfig = `
   # pattern = "nginx"
   ## user as argument for pgrep (ie, pgrep -u <user>)
   # user = "nginx"
+  ## Systemd unit name
+  # systemd_unit = "nginx.service"
+	## CGroup name or path
+	# cgroup = "systemd/system.slice/nginx.service"
 
   ## override for process_name
   ## This is optional; default is sourced from /proc/<pid>/status
@@ -275,6 +284,12 @@ func (p *Procstat) findPids() ([]PID, map[string]string, error) {
 	} else if p.User != "" {
 		pids, err = f.Uid(p.User)
 		tags = map[string]string{"user": p.User}
+	} else if p.SystemdUnit != "" {
+		pids, err = p.systemdUnitPIDs()
+		tags = map[string]string{"systemd_unit": p.SystemdUnit}
+	} else if p.CGroup != "" {
+		pids, err = p.cgroupPIDs()
+		tags = map[string]string{"cgroup": p.CGroup}
 	} else {
 		err = fmt.Errorf("Either exe, pid_file, user, or pattern has to be specified")
 	}
@@ -282,6 +297,59 @@ func (p *Procstat) findPids() ([]PID, map[string]string, error) {
 	return pids, tags, err
 }
 
+func (p *Procstat) systemdUnitPIDs() ([]PID, error) {
+	var pids []PID
+	cmd := exec.Command("systemctl", "show", p.SystemdUnit)
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, err
+	}
+	for _, line := range bytes.Split(out, []byte{'\n'}) {
+		kv := bytes.SplitN(line, []byte{'='}, 2)
+		if len(kv) != 2 {
+			continue
+		}
+		if !bytes.Equal(kv[0], []byte("MainPID")) {
+			continue
+		}
+		if len(kv[1]) == 0 {
+			return nil, nil
+		}
+		pid, err := strconv.Atoi(string(kv[1]))
+		if err != nil {
+			return nil, fmt.Errorf("invalid pid '%s'", kv[1])
+		}
+		pids = append(pids, PID(pid))
+	}
+	return pids, nil
+}
+
+func (p *Procstat) cgroupPIDs() ([]PID, error) {
+	var pids []PID
+
+	procsPath := p.CGroup
+	if procsPath[0] != '/' {
+		procsPath = "/sys/fs/cgroup/" + procsPath
+	}
+	procsPath = procsPath + "/cgroup.procs"
+	out, err := ioutil.ReadFile(procsPath)
+	if err != nil {
+		return nil, err
+	}
+	for _, pidBS := range bytes.Split(out, []byte{'\n'}) {
+		if len(pidBS) == 0 {
+			continue
+		}
+		pid, err := strconv.Atoi(string(pidBS))
+		if err != nil {
+			return nil, fmt.Errorf("invalid pid '%s'", pidBS)
+		}
+		pids = append(pids, PID(pid))
+	}
+
+	return pids, nil
+}
+
 func init() {
 	inputs.Add("procstat", func() telegraf.Input {
 		return &Procstat{}