Add native Go method for finding pids to procstat (#3559)

Makefile

@@ -68,6 +68,7 @@ test-windows:
 	go test ./plugins/inputs/ping/...
 	go test ./plugins/inputs/win_perf_counters/...
 	go test ./plugins/inputs/win_services/...
+	go test ./plugins/inputs/procstat/...
 
 # vet runs the Go source code static analysis tool `vet` to find
 # any common errors.

plugins/inputs/procstat/README.md

@@ -27,8 +27,28 @@ Additionally the plugin will tag processes by their PID (pid_tag = true in the c
 * pid
 * process_name
 
+
+### Windows
+On windows we only support exe and pattern. Both of these are implemented using WMI queries. exe is on the Name field and pattern is on the CommandLine field.
+
+Windows Support:
+* exe  (WMI Name)
+* pattern (WMI CommandLine)
+
+this allows you to do fuzzy matching but only what is supported by [WMI query patterns](https://msdn.microsoft.com/en-us/library/aa392263(v=vs.85).aspx).
+
 Example:
 
+Windows fuzzy matching:
+```[[inputs.procstat]]
+  exe = "%influx%"
+  process_name="influxd"
+  prefix = "influxd"
+
+```
+
+### Linux
+
 ```
 [[inputs.procstat]]
   exe = "influxd"
@@ -48,7 +68,6 @@ The above configuration would result in output like:
 # Measurements
 Note: prefix can be set by the user, per process.
 
-
 Threads related measurement names:
 - procstat_[prefix_]num_threads value=5
 

plugins/inputs/procstat/native_finder.go

@@ -0,0 +1,57 @@
+package procstat
+
+import (
+	"fmt"
+	"io/ioutil"
+	"strconv"
+	"strings"
+
+	"github.com/shirou/gopsutil/process"
+)
+
+//NativeFinder uses gopsutil to find processes
+type NativeFinder struct {
+}
+
+//NewNativeFinder ...
+func NewNativeFinder() (PIDFinder, error) {
+	return &NativeFinder{}, nil
+}
+
+//Uid will return all pids for the given user
+func (pg *NativeFinder) Uid(user string) ([]PID, error) {
+	var dst []PID
+	procs, err := process.Processes()
+	if err != nil {
+		return dst, err
+	}
+	for _, p := range procs {
+		username, err := p.Username()
+		if err != nil {
+			//skip, this can happen if we don't have permissions or
+			//the pid no longer exists
+			continue
+		}
+		if username == user {
+			dst = append(dst, PID(p.Pid))
+		}
+	}
+	return dst, nil
+}
+
+//PidFile returns the pid from the pid file given.
+func (pg *NativeFinder) PidFile(path string) ([]PID, error) {
+	var pids []PID
+	pidString, err := ioutil.ReadFile(path)
+	if err != nil {
+		return pids, fmt.Errorf("Failed to read pidfile '%s'. Error: '%s'",
+			path, err)
+	}
+	pid, err := strconv.Atoi(strings.TrimSpace(string(pidString)))
+	if err != nil {
+		return pids, err
+	}
+	pids = append(pids, PID(pid))
+	return pids, nil
+
+}

plugins/inputs/procstat/native_finder_notwindows.go

@@ -0,0 +1,59 @@
+// +build !windows
+
+package procstat
+
+import (
+	"regexp"
+
+	"github.com/shirou/gopsutil/process"
+)
+
+//Pattern matches on the process name
+func (pg *NativeFinder) Pattern(pattern string) ([]PID, error) {
+	var pids []PID
+	regxPattern, err := regexp.Compile(pattern)
+	if err != nil {
+		return pids, err
+	}
+	procs, err := process.Processes()
+	if err != nil {
+		return pids, err
+	}
+	for _, p := range procs {
+		name, err := p.Exe()
+		if err != nil {
+			//skip, this can be caused by the pid no longer existing
+			//or you having no permissions to access it
+			continue
+		}
+		if regxPattern.MatchString(name) {
+			pids = append(pids, PID(p.Pid))
+		}
+	}
+	return pids, err
+}
+
+//FullPattern matches on the command line when the proccess was executed
+func (pg *NativeFinder) FullPattern(pattern string) ([]PID, error) {
+	var pids []PID
+	regxPattern, err := regexp.Compile(pattern)
+	if err != nil {
+		return pids, err
+	}
+	procs, err := process.Processes()
+	if err != nil {
+		return pids, err
+	}
+	for _, p := range procs {
+		cmd, err := p.Cmdline()
+		if err != nil {
+			//skip, this can be caused by the pid no longer existing
+			//or you having no permissions to access it
+			continue
+		}
+		if regxPattern.MatchString(cmd) {
+			pids = append(pids, PID(p.Pid))
+		}
+	}
+	return pids, err
+}

plugins/inputs/procstat/native_finder_windows.go

@@ -0,0 +1,91 @@
+package procstat
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"time"
+
+	"github.com/StackExchange/wmi"
+	"github.com/shirou/gopsutil/process"
+)
+
+//Timeout is the timeout used when making wmi calls
+var Timeout = 5 * time.Second
+
+type queryType string
+
+const (
+	like     = queryType("LIKE")
+	equals   = queryType("=")
+	notEqual = queryType("!=")
+)
+
+//Pattern matches on the process name
+func (pg *NativeFinder) Pattern(pattern string) ([]PID, error) {
+	var pids []PID
+	regxPattern, err := regexp.Compile(pattern)
+	if err != nil {
+		return pids, err
+	}
+	procs, err := process.Processes()
+	if err != nil {
+		return pids, err
+	}
+	for _, p := range procs {
+		name, err := p.Name()
+		if err != nil {
+			//skip, this can be caused by the pid no longer existing
+			//or you having no permissions to access it
+			continue
+		}
+		if regxPattern.MatchString(name) {
+			pids = append(pids, PID(p.Pid))
+		}
+	}
+	return pids, err
+}
+
+//FullPattern matches the cmdLine on windows and will find a pattern using a WMI like query
+func (pg *NativeFinder) FullPattern(pattern string) ([]PID, error) {
+	var pids []PID
+	procs, err := getWin32ProcsByVariable("CommandLine", like, pattern, Timeout)
+	if err != nil {
+		return pids, err
+	}
+	for _, p := range procs {
+		pids = append(pids, PID(p.ProcessID))
+	}
+	return pids, nil
+}
+
+//GetWin32ProcsByVariable allows you to query any variable with a like query
+func getWin32ProcsByVariable(variable string, qType queryType, value string, timeout time.Duration) ([]process.Win32_Process, error) {
+	var dst []process.Win32_Process
+	var query string
+	// should look like "WHERE CommandLine LIKE "procstat"
+	query = fmt.Sprintf("WHERE %s %s %q", variable, qType, value)
+	q := wmi.CreateQuery(&dst, query)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	err := WMIQueryWithContext(ctx, q, &dst)
+	if err != nil {
+		return []process.Win32_Process{}, fmt.Errorf("could not get win32Proc: %s", err)
+	}
+	return dst, nil
+}
+
+// WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging
+func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error {
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- wmi.Query(query, dst, connectServerArgs...)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-errChan:
+		return err
+	}
+}

plugins/inputs/procstat/native_finder_windows_test.go

@@ -0,0 +1,40 @@
+package procstat
+
+import (
+	"fmt"
+	"testing"
+
+	"os/user"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGather_RealPattern(t *testing.T) {
+	pg, err := NewNativeFinder()
+	require.NoError(t, err)
+	pids, err := pg.Pattern(`procstat`)
+	require.NoError(t, err)
+	fmt.Println(pids)
+	assert.Equal(t, len(pids) > 0, true)
+}
+
+func TestGather_RealFullPattern(t *testing.T) {
+	pg, err := NewNativeFinder()
+	require.NoError(t, err)
+	pids, err := pg.FullPattern(`%procstat%`)
+	require.NoError(t, err)
+	fmt.Println(pids)
+	assert.Equal(t, len(pids) > 0, true)
+}
+
+func TestGather_RealUser(t *testing.T) {
+	user, err := user.Current()
+	require.NoError(t, err)
+	pg, err := NewNativeFinder()
+	require.NoError(t, err)
+	pids, err := pg.Uid(user.Username)
+	require.NoError(t, err)
+	fmt.Println(pids)
+	assert.Equal(t, len(pids) > 0, true)
+}

plugins/inputs/procstat/pgrep.go

@@ -8,13 +8,6 @@ import (
 	"strings"
 )
 
-type PIDFinder interface {
-	PidFile(path string) ([]PID, error)
-	Pattern(pattern string) ([]PID, error)
-	Uid(user string) ([]PID, error)
-	FullPattern(path string) ([]PID, error)
-}
-
 // Implemention of PIDGatherer that execs pgrep to find processes
 type Pgrep struct {
 	path string

plugins/inputs/procstat/process.go

@@ -23,6 +23,13 @@ type Process interface {
 	RlimitUsage(bool) ([]process.RlimitStat, error)
 }
 
+type PIDFinder interface {
+	PidFile(path string) ([]PID, error)
+	Pattern(pattern string) ([]PID, error)
+	Uid(user string) ([]PID, error)
+	FullPattern(path string) ([]PID, error)
+}
+
 type Proc struct {
 	hasCPUTimes bool
 	tags        map[string]string

plugins/inputs/procstat/procstat.go

@@ -22,6 +22,7 @@ var (
 type PID int32
 
 type Procstat struct {
+	PidFinder   string `toml:"pid_finder"`
 	PidFile     string `toml:"pid_file"`
 	Exe         string
 	Pattern     string
@@ -32,13 +33,19 @@ type Procstat struct {
 	CGroup      string `toml:"cgroup"`
 	PidTag      bool
 
-	pidFinder       PIDFinder
+	finder PIDFinder
+
 	createPIDFinder func() (PIDFinder, error)
 	procs           map[PID]Process
 	createProcess   func(PID) (Process, error)
 }
 
 var sampleConfig = `
+  ## pidFinder can be pgrep or native
+  ## pgrep tries to exec pgrep
+  ## native will work on all platforms, unix systems will use regexp. 
+  ## Windows will use WMI calls with like queries
+  pid_finder = "native"
   ## Must specify one of: pid_file, exe, or pattern
   ## PID file to monitor process
   pid_file = "/var/run/nginx.pid"
@@ -74,7 +81,15 @@ func (_ *Procstat) Description() string {
 
 func (p *Procstat) Gather(acc telegraf.Accumulator) error {
 	if p.createPIDFinder == nil {
-		p.createPIDFinder = defaultPIDFinder
+		switch p.PidFinder {
+		case "native":
+			p.createPIDFinder = NewNativeFinder
+		case "pgrep":
+			p.createPIDFinder = NewPgrep
+		default:
+			p.createPIDFinder = defaultPIDFinder
+		}
+
 	}
 	if p.createProcess == nil {
 		p.createProcess = defaultProcess
@@ -252,14 +267,15 @@ func (p *Procstat) updateProcesses(prevInfo map[PID]Process) (map[PID]Process, e
 
 // Create and return PIDGatherer lazily
 func (p *Procstat) getPIDFinder() (PIDFinder, error) {
-	if p.pidFinder == nil {
+
+	if p.finder == nil {
 		f, err := p.createPIDFinder()
 		if err != nil {
 			return nil, err
 		}
-		p.pidFinder = f
+		p.finder = f
 	}
-	return p.pidFinder, nil
+	return p.finder, nil
 }
 
 // Get matching PIDs and their initial tags