fix disk metrics calculation for various block devices (#6226)

Signed-off-by: Dejan Zele Pejchev <[email protected]>

pkg/utilization/instrument.go

@@ -2,6 +2,7 @@ package utilization
 
 import (
 	errors2 "errors"
+	"regexp"
 	"sync"
 
 	"github.com/shirou/gopsutil/v4/disk"
@@ -11,6 +12,12 @@ import (
 	"github.com/shirou/gopsutil/v4/process"
 )
 
+// This regex matches:
+//   - sd + exactly one letter (e.g. sda, sdb)
+//   - vd + exactly one letter (e.g. vda, vdb)
+//   - nvme + digits + n + digits (e.g. nvme0n1, nvme1n1, not partitions)
+var reBaseDevice = regexp.MustCompile(`^(sd[a-z]|vd[a-z]|xvd[a-z]|nvme\d+n\d+)$`)
+
 type Metrics struct {
 	Memory  *process.MemoryInfoStat
 	CPU     float64
@@ -93,15 +100,47 @@ func aggregate(metrics []*Metrics) *Metrics {
 
 // recordSystemWideMetrics captures the network and disk system-wide metrics by using the global gopsutil packages instead of the process one.
 func (r *MetricRecorder) recordSystemWideMetrics(aggregated *Metrics) {
-	diskName := "vda"
-	io, _ := disk.IOCounters(diskName)
+	io, _ := disk.IOCounters()
 	if len(io) > 0 {
-		d := io[diskName]
-		aggregated.Disk = &d
+		aggregated.Disk = filterAndAggregateDiskStats(io)
 	}
 	n, _ := net.IOCounters(false)
 	if len(n) > 0 {
 		n := n[0]
 		aggregated.Network = &n
 	}
 }
+
+// filterAndAggregateDiskStats filters stats for disk devices (but not partitions).
+// It matches:
+//   - sda, sdb, sdc, etc. (SCSI/SATA/SAS)
+//   - vda, vdb, etc. (VirtIO)
+//   - xvda, xvdb, etc. (AWS Xen)
+//   - nvme0n1, nvme1n1, etc. (NVMe)
+//
+// It does NOT match sda1, sdb2, nvme0n1p1, etc.
+func filterAndAggregateDiskStats(stats map[string]disk.IOCountersStat) *disk.IOCountersStat {
+	aggregated := &disk.IOCountersStat{}
+
+	for diskName, stat := range stats {
+		if reBaseDevice.MatchString(diskName) {
+			aggregateDiskStats(aggregated, &stat)
+		}
+	}
+
+	return aggregated
+}
+
+func aggregateDiskStats(aggregate, stats *disk.IOCountersStat) {
+	aggregate.ReadCount += stats.ReadCount
+	aggregate.WriteCount += stats.WriteCount
+	aggregate.ReadBytes += stats.ReadBytes
+	aggregate.WriteBytes += stats.WriteBytes
+	aggregate.ReadTime += stats.ReadTime
+	aggregate.WriteTime += stats.WriteTime
+	aggregate.IopsInProgress += stats.IopsInProgress
+	aggregate.IoTime += stats.IoTime
+	aggregate.WeightedIO += stats.WeightedIO
+	aggregate.MergedReadCount += stats.MergedReadCount
+	aggregate.MergedWriteCount += stats.MergedWriteCount
+}

pkg/utilization/instrument_test.go

@@ -0,0 +1,74 @@
+package utilization
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/shirou/gopsutil/v4/disk"
+)
+
+func TestFilterAndAggregateDiskStats(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name  string
+		stats map[string]disk.IOCountersStat
+		want  disk.IOCountersStat
+	}{
+		{
+			name: "Include base sda, skip sda1",
+			stats: map[string]disk.IOCountersStat{
+				"sda":  {ReadBytes: 1000, WriteBytes: 2000},
+				"sda1": {ReadBytes: 3000, WriteBytes: 4000},
+			},
+			// Only sda is included.
+			want: disk.IOCountersStat{ReadBytes: 1000, WriteBytes: 2000},
+		},
+		{
+			name: "Include base nvme0n1, skip nvme0n1p1",
+			stats: map[string]disk.IOCountersStat{
+				"nvme0n1":   {ReadBytes: 5000, WriteBytes: 6000},
+				"nvme0n1p1": {ReadBytes: 7000, WriteBytes: 8000},
+			},
+			// Only nvme0n1 is included.
+			want: disk.IOCountersStat{ReadBytes: 5000, WriteBytes: 6000},
+		},
+		{
+			name: "Include multiple base devices (vda, vdb, sdz, nvme1n1) but skip their partitions",
+			stats: map[string]disk.IOCountersStat{
+				"vda":       {ReadBytes: 100, WriteBytes: 200},
+				"vdb":       {ReadBytes: 500, WriteBytes: 700},
+				"vda1":      {ReadBytes: 10, WriteBytes: 20},
+				"sdz":       {ReadBytes: 900, WriteBytes: 1100},
+				"sdz1":      {ReadBytes: 2000, WriteBytes: 3000},
+				"nvme1n1":   {ReadBytes: 4000, WriteBytes: 5000},
+				"nvme1n1p2": {ReadBytes: 100, WriteBytes: 100},
+				"xvda":      {ReadBytes: 4000, WriteBytes: 5000},
+				"xvda1":     {ReadBytes: 100, WriteBytes: 100},
+			},
+			// Summation of base devices only: vda + vdb + sdz + nvme1n1
+			want: disk.IOCountersStat{
+				ReadBytes:  100 + 500 + 900 + 4000 + 4000,
+				WriteBytes: 200 + 700 + 1100 + 5000 + 5000,
+			},
+		},
+		{
+			name: "No matching devices",
+			stats: map[string]disk.IOCountersStat{
+				"loop0": {ReadBytes: 1000, WriteBytes: 2000},
+				"dm-0":  {ReadBytes: 300, WriteBytes: 400},
+			},
+			// None match the regex, so the result should be zeroed.
+			want: disk.IOCountersStat{ReadBytes: 0, WriteBytes: 0},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := filterAndAggregateDiskStats(tt.stats)
+			assert.Equal(t, tt.want.ReadBytes, got.ReadBytes)
+			assert.Equal(t, tt.want.WriteBytes, got.WriteBytes)
+		})
+	}
+}