From cf40c67808f56904f703406060f8448e2d8844fe Mon Sep 17 00:00:00 2001
From: zeim839 <zeim839@users.noreply.github.com>
Date: Wed, 11 Oct 2023 11:03:28 -0400
Subject: [PATCH] add CaptureCPUStats, CaptureDiskStats

---
 cpu.go  | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 disk.go | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/cpu.go b/cpu.go
index 831730e..ed584ee 100644
--- a/cpu.go
+++ b/cpu.go
@@ -16,6 +16,21 @@
 
 package metrics
 
+import (
+	"time"
+	"sync"
+)
+
+var (
+	cpuStats CPUStats
+	cpuMetrics struct {
+		GlobalTime GaugeFloat64
+		GlobalWait GaugeFloat64
+		LocalTime GaugeFloat64
+	}
+	registerCPUMetricsOnce = sync.Once{}
+)
+
 // CPUStats is the system and process CPU stats.
 // All values are in seconds.
 type CPUStats struct {
@@ -23,3 +38,41 @@ type CPUStats struct {
 	GlobalWait float64 // Time spent by waiting on disk for all processes.
 	LocalTime  float64 // Time spent by the CPU working on this process.
 }
+
+// CaptureCPUStats captures new values for the Go process CPU usage
+// statistics exported in cpu.CPUStats. This is designed to be called as a
+// goroutine.
+func CaptureCPUStats(d time.Duration) {
+	for range time.Tick(d) {
+		CaptureCPUStatsOnce()
+	}
+}
+
+// CaptureCPUStatsOnce captures new values for the Go process CPU usage
+// statistics exported in cpu.CPUStats. This is designed to be called in a
+// background goroutine.
+func CaptureCPUStatsOnce() {
+	err := ReadCPUStats(&cpuStats)
+	if err != nil {
+		panic(err)
+	}
+	cpuMetrics.GlobalTime.Update(cpuStats.GlobalTime)
+	cpuMetrics.GlobalWait.Update(cpuStats.GlobalWait)
+	cpuMetrics.LocalTime.Update(cpuStats.LocalTime)
+}
+
+// RegisterCPUStats registers metrics for the Go process CPU usage statistics
+// exported in cpu.CPUStats.
+func RegisterCPUStats(r Registry) {
+	if r == nil {
+		r = DefaultRegistry
+	}
+	registerCPUMetricsOnce.Do(func() {
+		cpuMetrics.GlobalTime = NewGaugeFloat64(nil)
+		cpuMetrics.GlobalWait = NewGaugeFloat64(nil)
+		cpuMetrics.LocalTime = NewGaugeFloat64(nil)
+		r.Register("cpu.CPUStats.GlobalTime", cpuMetrics.GlobalTime)
+		r.Register("cpu.CPUStats.GlobalWait", cpuMetrics.GlobalWait)
+		r.Register("cpu.CPUStats.LocalTime", cpuMetrics.LocalTime)
+	})
+}
diff --git a/disk.go b/disk.go
index 25142d2..224114a 100644
--- a/disk.go
+++ b/disk.go
@@ -16,6 +16,22 @@
 
 package metrics
 
+import (
+	"time"
+	"sync"
+)
+
+var (
+	diskStats  DiskStats
+	diskMetrics struct {
+		ReadCount  Gauge
+		ReadBytes  Gauge
+		WriteCount Gauge
+		WriteBytes Gauge
+	}
+	registerDiskMetricsOnce = sync.Once{}
+)
+
 // DiskStats is the per process disk io stats.
 type DiskStats struct {
 	ReadCount  int64 // Number of read operations executed
@@ -23,3 +39,44 @@ type DiskStats struct {
 	WriteCount int64 // Number of write operations executed
 	WriteBytes int64 // Total number of byte written
 }
+
+// CaptureDiskStats captures new values for the Go process disk usage
+// statistics exported in disk.DiskStats. This is designed to be called as a
+// goroutine.
+func CaptureDiskStats(d time.Duration) {
+	for range time.Tick(d) {
+		CaptureDiskStatsOnce()
+	}
+}
+
+// CaptureDiskStatsOnce captures new values for the Go process disk usage
+// statistics exported in disk.DiskStats. This is designed to be called in a
+// background goroutine.
+func CaptureDiskStatsOnce() {
+	err := ReadDiskStats(&diskStats)
+	if err != nil {
+		panic(err)
+	}
+	diskMetrics.ReadCount.Update(diskStats.ReadCount)
+	diskMetrics.ReadBytes.Update(diskStats.ReadBytes)
+	diskMetrics.WriteCount.Update(diskStats.WriteCount)
+	diskMetrics.WriteBytes.Update(diskStats.WriteBytes)
+}
+
+// RegisterDiskStats registers metrics for the Go process disk usage statistics
+// exported in disk.DiskStats.
+func RegisterDiskStats(r Registry) {
+	if r == nil {
+		r = DefaultRegistry
+	}
+	registerDiskMetricsOnce.Do(func() {
+		diskMetrics.ReadCount = NewGauge(nil)
+		diskMetrics.ReadBytes = NewGauge(nil)
+		diskMetrics.WriteCount = NewGauge(nil)
+		diskMetrics.WriteBytes = NewGauge(nil)
+		r.Register("disk.DiskStats.ReadCount", diskMetrics.ReadCount)
+		r.Register("disk.DiskStats.ReadBytes", diskMetrics.ReadBytes)
+		r.Register("disk.DiskStats.WriteCount", diskMetrics.WriteCount)
+		r.Register("disk.DiskStats.WriteBytes", diskMetrics.WriteBytes)
+	})
+}