Merge pull request #32 from bluewave-labs/multi_storage

Multi partition support

internal/metric/disk.go

@@ -1,12 +1,16 @@
 package metric
 
 import (
-	disk2 "github.com/shirou/gopsutil/v4/disk"
+	"slices"
+	"strings"
+
+	"github.com/shirou/gopsutil/v4/disk"
 )
 
 func CollectDiskMetrics() (MetricsSlice, []CustomErr) {
 	defaultDiskData := []*DiskData{
 		{
+			Device:          "unknown",
 			ReadSpeedBytes:  nil,
 			WriteSpeedBytes: nil,
 			TotalBytes:      nil,
@@ -15,35 +19,56 @@ func CollectDiskMetrics() (MetricsSlice, []CustomErr) {
 		},
 	}
 	var diskErrors []CustomErr
-	diskUsage, diskUsageErr := disk2.Usage("/")
+	var metricsSlice MetricsSlice
+	var checkedSlice []string // To keep track of checked partitions
+
+	// Set all flag to "false" to get only necessary partitions
+	// Avoiding unnecessary partitions like /run/user/1000, /run/credentials
+	partitions, partErr := disk.Partitions(false)
 
-	if diskUsageErr != nil {
+	if partErr != nil {
 		diskErrors = append(diskErrors, CustomErr{
-			Metric: []string{"disk.usage_percent", "disk.total_bytes", "disk.free_bytes"},
-			Error:  diskUsageErr.Error(),
+			Metric: []string{"disk.partitions"},
+			Error:  partErr.Error(),
 		})
-		return MetricsSlice{defaultDiskData[0]}, diskErrors
 	}
 
-	// diskMetrics, diskErr := disk1.Get()
-	// if diskErr != nil {
-	// 	log.Fatalf("Unable to get disk metrics")
-	// }
+	for _, p := range partitions {
+		// Filter out partitions that are already checked or not a device
+		// Also, exclude '/dev/loop' devices to avoid unnecessary partitions
+		// * /dev/loop devices are used for mounting snap packages
+		if slices.Contains(checkedSlice, p.Device) || !strings.HasPrefix(p.Device, "/dev") || strings.HasPrefix(p.Device, "/dev/loop") {
+			continue
+		}
 
-	// for _, p := range diskMetrics {
-	// 	fmt.Println(p.Name, p.ReadsCompleted)
-	// }
+		diskUsage, diskUsageErr := disk.Usage(p.Mountpoint)
+		if diskUsageErr != nil {
+			diskErrors = append(diskErrors, CustomErr{
+				Metric: []string{"disk.usage_percent", "disk.total_bytes", "disk.free_bytes"},
+				Error:  diskUsageErr.Error() + " " + p.Mountpoint,
+			})
+			continue
+		}
 
-	// var a uint64 = 2e+12
-	var metricsSlice MetricsSlice
+		checkedSlice = append(checkedSlice, p.Device)
+		metricsSlice = append(metricsSlice, &DiskData{
+			Device:          p.Device,
+			ReadSpeedBytes:  nil, // TODO: Implement
+			WriteSpeedBytes: nil, // TODO: Implement
+			TotalBytes:      &diskUsage.Total,
+			FreeBytes:       &diskUsage.Free,
+			UsagePercent:    RoundFloatPtr(diskUsage.UsedPercent/100, 4),
+		})
+	}
+
+	if len(diskErrors) == 0 {
+		return metricsSlice, nil
+	}
+
+	if len(metricsSlice) == 0 {
+		return MetricsSlice{defaultDiskData[0]}, diskErrors
+	}
 
-	metricsSlice = append(metricsSlice, &DiskData{
-		ReadSpeedBytes:  nil, // TODO: Implement
-		WriteSpeedBytes: nil, // TODO: Implement
-		TotalBytes:      &diskUsage.Total,
-		FreeBytes:       &diskUsage.Free,
-		UsagePercent:    RoundFloatPtr(diskUsage.UsedPercent/100, 4),
-	})
 	return metricsSlice, diskErrors
 }
 

internal/metric/metric.go

@@ -49,11 +49,12 @@ type MemoryData struct {
 func (m MemoryData) isMetric() {}
 
 type DiskData struct {
+	Device          string   `json:"device"`            // Device
 	ReadSpeedBytes  *uint64  `json:"read_speed_bytes"`  // TODO: Implement
 	WriteSpeedBytes *uint64  `json:"write_speed_bytes"` // TODO: Implement
-	TotalBytes      *uint64  `json:"total_bytes"`       // Total space of "/" in bytes
-	FreeBytes       *uint64  `json:"free_bytes"`        // Free space of "/" in bytes
-	UsagePercent    *float64 `json:"usage_percent"`     // Usage Percent of "/"
+	TotalBytes      *uint64  `json:"total_bytes"`       // Total space of device in bytes
+	FreeBytes       *uint64  `json:"free_bytes"`        // Free space of device in bytes
+	UsagePercent    *float64 `json:"usage_percent"`     // Usage Percent of device
 }
 
 func (d DiskData) isMetric() {}

internal/metric/metric_math.go

@@ -9,9 +9,15 @@ import (
 
 // Round a float to a given precision and return the pointer of the result
 func RoundFloatPtr(val float64, precision uint) *float64 {
+	r := RoundFloat(val, precision)
+	return &r
+}
+
+// Round a float to a given precision and return the result
+func RoundFloat(val float64, precision uint) float64 {
 	ratio := math.Pow(10, float64(precision))
 	prc := math.Round(val*ratio) / ratio
-	return &prc
+	return prc
 }
 
 func RandomIntPtr(max int64) *int {