From fa6ca8f7175e3d6db34706e361e448bade201280 Mon Sep 17 00:00:00 2001
From: usbharu <64310155+usbharu@users.noreply.github.com>
Date: Wed, 13 Mar 2024 10:54:28 +0900
Subject: [PATCH] =?UTF-8?q?watchdog-go:=20disk=E6=83=85=E5=A0=B1=E3=82=92?=
 =?UTF-8?q?=E5=8F=96=E5=BE=97=E3=81=A7=E3=81=8D=E3=82=8B=E3=82=88=E3=81=86?=
 =?UTF-8?q?=E3=81=AB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 watchdog-go/watchdog/disk_collector.go | 81 ++++++++++++++++++++++++++
 watchdog-go/watchdog/watchdog_go.go    | 18 ++++++
 2 files changed, 99 insertions(+)
 create mode 100644 watchdog-go/watchdog/disk_collector.go

diff --git a/watchdog-go/watchdog/disk_collector.go b/watchdog-go/watchdog/disk_collector.go
new file mode 100644
index 0000000..f95c655
--- /dev/null
+++ b/watchdog-go/watchdog/disk_collector.go
@@ -0,0 +1,81 @@
+package watchdog
+
+import (
+	"fmt"
+	"github.com/shirou/gopsutil/v3/disk"
+	"google.golang.org/protobuf/types/known/timestamppb"
+	"time"
+)
+
+type UsedDisk struct {
+	config        Config
+	partitionStat disk.PartitionStat
+}
+
+func NewUsedDisk(config Config, stat disk.PartitionStat) MetricsCollector {
+	return &UsedDisk{
+		config:        config,
+		partitionStat: stat,
+	}
+}
+
+func (u UsedDisk) Collect() (*Metric, error) {
+	usage, err := disk.Usage(u.partitionStat.Mountpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	name := u.config.ClientName + " Used Disk " + u.partitionStat.Device
+	usedDisk := fmt.Sprintf("%g", usage.UsedPercent)
+
+	return &Metric{
+		Name:      name,
+		ObjectId:  toId(name),
+		Domain:    u.config.ClientDomain,
+		Status:    Status_OK,
+		Value:     usedDisk,
+		Timestamp: timestamppb.Now(),
+		Message:   "",
+	}, nil
+}
+
+func (u UsedDisk) Timer() time.Duration {
+	return 1 * time.Hour
+}
+
+type TotalDisk struct {
+	Config
+	disk.PartitionStat
+}
+
+func NewTotalDisk(config Config, stat disk.PartitionStat) MetricsCollector {
+	return &TotalDisk{
+		Config:        config,
+		PartitionStat: stat,
+	}
+}
+
+func (t TotalDisk) Collect() (*Metric, error) {
+	usage, err := disk.Usage(t.Mountpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	name := t.ClientName + " Total Disk " + t.Device
+
+	totalDisk := fmt.Sprintf("%d", usage.Total)
+
+	return &Metric{
+		Name:      name,
+		ObjectId:  toId(name),
+		Domain:    t.ClientDomain,
+		Status:    Status_OK,
+		Value:     totalDisk,
+		Timestamp: timestamppb.Now(),
+		Message:   "",
+	}, nil
+}
+
+func (t TotalDisk) Timer() time.Duration {
+	return 24 * time.Hour
+}
diff --git a/watchdog-go/watchdog/watchdog_go.go b/watchdog-go/watchdog/watchdog_go.go
index b6cc742..165d3aa 100644
--- a/watchdog-go/watchdog/watchdog_go.go
+++ b/watchdog-go/watchdog/watchdog_go.go
@@ -2,6 +2,7 @@ package watchdog
 
 import (
 	"fmt"
+	"github.com/shirou/gopsutil/v3/disk"
 	"github.com/shirou/gopsutil/v3/host"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
@@ -37,6 +38,8 @@ func Run() {
 		NewUsedCpu(config),
 	}
 
+	collectors = append(collectors, diskCollector(config)...)
+
 	collectorMap := map[time.Duration][]MetricsCollector{}
 
 	for _, collector := range collectors {
@@ -106,3 +109,18 @@ func buildConfig() Config {
 		ClientDomain: load.Section("Watch").Key("domain").MustString(hostinfo.Hostname),
 	}
 }
+
+func diskCollector(config Config) []MetricsCollector {
+	partitions, err := disk.Partitions(true)
+	if err != nil {
+		log.Println(err)
+		return []MetricsCollector{}
+	}
+
+	var collectors []MetricsCollector
+	for _, partition := range partitions {
+		collectors = append(collectors, NewUsedDisk(config, partition))
+		collectors = append(collectors, NewTotalDisk(config, partition))
+	}
+	return collectors
+}