From e890743f7978abec2de95b4e116745486f73782c Mon Sep 17 00:00:00 2001
From: Calle Pettersson <cpettsson@gmail.com>
Date: Tue, 16 Aug 2016 08:39:49 +0200
Subject: [PATCH] Add white/blacklisting of units for systemd collector

---
 collector/systemd_linux.go      | 44 +++++++++++++++++++++++++----
 collector/systemd_linux_test.go | 49 +++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 6 deletions(-)

diff --git a/collector/systemd_linux.go b/collector/systemd_linux.go
index c2fae13c..93b25d49 100644
--- a/collector/systemd_linux.go
+++ b/collector/systemd_linux.go
@@ -18,14 +18,23 @@ package collector
 import (
 	"flag"
 	"fmt"
+	"regexp"
 
 	"github.com/coreos/go-systemd/dbus"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/log"
+)
+
+var (
+	unitWhitelist = flag.String("collector.systemd.unit-whitelist", ".+", "Regexp of systemd units to whitelist. Units must both match whitelist and not match blacklist to be included.")
+	unitBlacklist = flag.String("collector.systemd.unit-blacklist", "", "Regexp of systemd units to blacklist. Units must both match whitelist and not match blacklist to be included.")
 )
 
 type systemdCollector struct {
-	unitDesc          *prometheus.Desc
-	systemRunningDesc *prometheus.Desc
+	unitDesc             *prometheus.Desc
+	systemRunningDesc    *prometheus.Desc
+	unitWhitelistPattern *regexp.Regexp
+	unitBlacklistPattern *regexp.Regexp
 }
 
 var unitStatesName = []string{"active", "activating", "deactivating", "inactive", "failed"}
@@ -56,10 +65,14 @@ func NewSystemdCollector() (Collector, error) {
 		"Whether the system is operational (see 'systemctl is-system-running')",
 		nil, nil,
 	)
+	unitWhitelistPattern := regexp.MustCompile(fmt.Sprintf("^(?:%s)$", *unitWhitelist))
+	unitBlacklistPattern := regexp.MustCompile(fmt.Sprintf("^(?:%s)$", *unitBlacklist))
 
 	return &systemdCollector{
-		unitDesc:          unitDesc,
-		systemRunningDesc: systemRunningDesc,
+		unitDesc:             unitDesc,
+		systemRunningDesc:    systemRunningDesc,
+		unitWhitelistPattern: unitWhitelistPattern,
+		unitBlacklistPattern: unitBlacklistPattern,
 	}, nil
 }
 
@@ -113,9 +126,28 @@ func (c *systemdCollector) listUnits() ([]dbus.UnitStatus, error) {
 	if err != nil {
 		return nil, fmt.Errorf("couldn't get dbus connection: %s", err)
 	}
-	units, err := conn.ListUnits()
+	allUnits, err := conn.ListUnits()
 	conn.Close()
-	return units, err
+
+	if err != nil {
+		return []dbus.UnitStatus{}, err
+	}
+
+	units := filterUnits(allUnits, c.unitWhitelistPattern, c.unitBlacklistPattern)
+	return units, nil
+}
+
+func filterUnits(units []dbus.UnitStatus, whitelistPattern, blacklistPattern *regexp.Regexp) []dbus.UnitStatus {
+	filtered := make([]dbus.UnitStatus, 0, len(units))
+	for _, unit := range units {
+		if whitelistPattern.MatchString(unit.Name) && !blacklistPattern.MatchString(unit.Name) {
+			filtered = append(filtered, unit)
+		} else {
+			log.Debugf("Ignoring unit: %s", unit.Name)
+		}
+	}
+
+	return filtered
 }
 
 func (c *systemdCollector) getSystemState() (state string, err error) {
diff --git a/collector/systemd_linux_test.go b/collector/systemd_linux_test.go
index c6217203..07e013a7 100644
--- a/collector/systemd_linux_test.go
+++ b/collector/systemd_linux_test.go
@@ -14,6 +14,7 @@
 package collector
 
 import (
+	"regexp"
 	"testing"
 
 	"github.com/coreos/go-systemd/dbus"
@@ -47,6 +48,30 @@ func getUnitListFixtures() [][]dbus.UnitStatus {
 			JobType:     "",
 			JobPath:     "/",
 		},
+		dbus.UnitStatus{
+			Name:        "foobar",
+			Description: "bar desc",
+			LoadState:   "not-found",
+			ActiveState: "inactive",
+			SubState:    "dead",
+			Followed:    "",
+			Path:        "/org/freedesktop/systemd1/unit/bar",
+			JobId:       0,
+			JobType:     "",
+			JobPath:     "/",
+		},
+		dbus.UnitStatus{
+			Name:        "baz",
+			Description: "bar desc",
+			LoadState:   "not-found",
+			ActiveState: "inactive",
+			SubState:    "dead",
+			Followed:    "",
+			Path:        "/org/freedesktop/systemd1/unit/bar",
+			JobId:       0,
+			JobType:     "",
+			JobPath:     "/",
+		},
 	}
 
 	fixture2 := []dbus.UnitStatus{}
@@ -72,3 +97,27 @@ func TestSystemdCollectorDoesntCrash(t *testing.T) {
 		collector.collectUnitStatusMetrics(sink, units)
 	}
 }
+
+func TestSystemdIgnoreFilter(t *testing.T) {
+	fixtures := getUnitListFixtures()
+	whitelistPattern := regexp.MustCompile("foo")
+	blacklistPattern := regexp.MustCompile("bar")
+	filtered := filterUnits(fixtures[0], whitelistPattern, blacklistPattern)
+	for _, unit := range filtered {
+		if blacklistPattern.MatchString(unit.Name) || !whitelistPattern.MatchString(unit.Name) {
+			t.Error(unit.Name, "should not be in the filtered list")
+		}
+	}
+}
+func TestSystemdIgnoreFilterDefaultKeepsAll(t *testing.T) {
+	c, err := NewSystemdCollector()
+	if err != nil {
+		t.Fatal(err)
+	}
+	fixtures := getUnitListFixtures()
+	collector := c.(*systemdCollector)
+	filtered := filterUnits(fixtures[0], collector.unitWhitelistPattern, collector.unitBlacklistPattern)
+	if len(filtered) != len(fixtures[0]) {
+		t.Error("Default filters removed units")
+	}
+}