From 8fcc6320a2316d5d179b2a9a4a78d7281891c44d Mon Sep 17 00:00:00 2001
From: rushilenekar20 <rushilenekar20@gmail.com>
Date: Tue, 7 Jun 2022 01:17:31 -0500
Subject: [PATCH] Add diskstats include and exclude device flags

Use standard include/exclude pattern for device include/exclude in the
diskstats collector.

Signed-off-by: Ben Kochie <superq@gmail.com>
Co-authored-by: rushilenekar20 <rushilenekar20@gmail.com>
---
 collector/diskstats_linux.go      | 44 ++++++++++++++++++++++---------
 collector/diskstats_linux_test.go |  2 +-
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/collector/diskstats_linux.go b/collector/diskstats_linux.go
index f0de40776b..8f25b587f1 100644
--- a/collector/diskstats_linux.go
+++ b/collector/diskstats_linux.go
@@ -17,9 +17,8 @@
 package collector
 
 import (
+	"errors"
 	"fmt"
-	"regexp"
-
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
@@ -36,7 +35,9 @@ const (
 )
 
 var (
-	ignoredDevices = kingpin.Flag("collector.diskstats.ignored-devices", "Regexp of devices to ignore for diskstats.").Default("^(ram|loop|fd|(h|s|v|xv)d[a-z]|nvme\\d+n\\d+p)\\d+$").String()
+	diskstatsDeviceExclude    = kingpin.Flag("collector.diskstats.device-exclude", "Regexp of diskstats devices to exclude (mutually exclusive to device-include).").String()
+	oldDiskstatsDeviceExclude = kingpin.Flag("collector.diskstats.ignored-devices", "DEPRECATED: Use collector.diskstats.device-exclude").String()
+	diskstatsDeviceInclude    = kingpin.Flag("collector.diskstats.device-include", "Regexp of diskstats devices to include (mutually exclusive to device-exclude).").String()
 )
 
 type typedFactorDesc struct {
@@ -49,11 +50,11 @@ func (d *typedFactorDesc) mustNewConstMetric(value float64, labels ...string) pr
 }
 
 type diskstatsCollector struct {
-	ignoredDevicesPattern *regexp.Regexp
-	fs                    blockdevice.FS
-	infoDesc              typedFactorDesc
-	descs                 []typedFactorDesc
-	logger                log.Logger
+	deviceFilter deviceFilter
+	fs           blockdevice.FS
+	infoDesc     typedFactorDesc
+	descs        []typedFactorDesc
+	logger       log.Logger
 }
 
 func init() {
@@ -68,10 +69,29 @@ func NewDiskstatsCollector(logger log.Logger) (Collector, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to open sysfs: %w", err)
 	}
+	if *oldDiskstatsDeviceExclude != "" {
+		if *diskstatsDeviceExclude == "" {
+			level.Warn(logger).Log("msg", "--collector.diskstats.ignored-devices is DEPRECATED and will be removed in 2.0.0, use --collector.diskstats.device-exclude")
+			*diskstatsDeviceExclude = *oldDiskstatsDeviceExclude
+		} else {
+			return nil, errors.New("--collector.diskstats.ignored-devices and --collector.diskstats.device-exclude are mutually exclusive")
+		}
+	}
+
+	if *diskstatsDeviceExclude != "" && *diskstatsDeviceInclude != "" {
+		return nil, errors.New("device-exclude & device-include are mutually exclusive")
+	}
 
+	if *diskstatsDeviceExclude != "" {
+		level.Info(logger).Log("msg", "Parsed flag --collector.diskstats.device-exclude", "flag", *diskstatsDeviceExclude)
+	}
+
+	if *diskstatsDeviceInclude != "" {
+		level.Info(logger).Log("msg", "Parsed Flag --collector.diskstats.device-include", "flag", *diskstatsDeviceInclude)
+	}
 	return &diskstatsCollector{
-		ignoredDevicesPattern: regexp.MustCompile(*ignoredDevices),
-		fs:                    fs,
+		deviceFilter: newDeviceFilter(*diskstatsDeviceExclude, *diskstatsDeviceInclude),
+		fs:           fs,
 		infoDesc: typedFactorDesc{
 			desc: prometheus.NewDesc(prometheus.BuildFQName(namespace, diskSubsystem, "info"),
 				"Info of /sys/block/<block_device>.",
@@ -194,11 +214,9 @@ func (c *diskstatsCollector) Update(ch chan<- prometheus.Metric) error {
 
 	for _, stats := range diskStats {
 		dev := stats.DeviceName
-		if c.ignoredDevicesPattern.MatchString(dev) {
-			level.Debug(c.logger).Log("msg", "Ignoring device", "device", dev, "pattern", c.ignoredDevicesPattern)
+		if c.deviceFilter.ignored(dev) {
 			continue
 		}
-
 		ch <- c.infoDesc.mustNewConstMetric(1.0, dev, fmt.Sprint(stats.MajorNumber), fmt.Sprint(stats.MinorNumber))
 
 		statCount := stats.IoStatsCount - 3 // Total diskstats record count, less MajorNumber, MinorNumber and DeviceName
diff --git a/collector/diskstats_linux_test.go b/collector/diskstats_linux_test.go
index 0e8dc84ff2..4aeb025e7a 100644
--- a/collector/diskstats_linux_test.go
+++ b/collector/diskstats_linux_test.go
@@ -49,7 +49,7 @@ func NewTestDiskStatsCollector(logger log.Logger) (prometheus.Collector, error)
 func TestDiskStats(t *testing.T) {
 	*sysPath = "fixtures/sys"
 	*procPath = "fixtures/proc"
-	*ignoredDevices = "^(ram|loop|fd|(h|s|v|xv)d[a-z]|nvme\\d+n\\d+p)\\d+$"
+	*diskstatsDeviceExclude = "^(ram|loop|fd|(h|s|v|xv)d[a-z]|nvme\\d+n\\d+p)\\d+$"
 	testcase := `# HELP node_disk_discard_time_seconds_total This is the total number of seconds spent by all discards.
 # TYPE node_disk_discard_time_seconds_total counter
 node_disk_discard_time_seconds_total{device="sdb"} 11.13