From e3ef92b21ab805a5c9cabd600e1e86af6dd4491a Mon Sep 17 00:00:00 2001
From: Gyuho Lee <gyuhox@gmail.com>
Date: Thu, 5 Dec 2024 12:58:06 +0800
Subject: [PATCH] add disk render

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>
---
 cmd/gpud/command/command.go    |  6 ++++++
 cmd/gpud/command/scan.go       |  1 +
 components/diagnose/options.go |  9 ++++++++-
 components/diagnose/scan.go    | 11 +++++++++++
 pkg/disk/disk.go               | 21 +++++++++++++++++++++
 pkg/disk/disk_test.go          |  7 ++++++-
 6 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/cmd/gpud/command/command.go b/cmd/gpud/command/command.go
index bde1f756..a493e06d 100644
--- a/cmd/gpud/command/command.go
+++ b/cmd/gpud/command/command.go
@@ -41,6 +41,7 @@ var (
 	pollXidEvents bool
 	pollGPMEvents bool
 	netcheck      bool
+	diskcheck     bool
 
 	enableAutoUpdate   bool
 	autoUpdateExitCode int
@@ -530,6 +531,11 @@ cat summary.txt
 					Usage:       "enable network connectivity checks to global edge/derp servers (default: true)",
 					Destination: &netcheck,
 				},
+				&cli.BoolTFlag{
+					Name:        "diskcheck",
+					Usage:       "enable disk checks (default: true)",
+					Destination: &diskcheck,
+				},
 			},
 		},
 		{
diff --git a/cmd/gpud/command/scan.go b/cmd/gpud/command/scan.go
index bf433951..94b09d06 100644
--- a/cmd/gpud/command/scan.go
+++ b/cmd/gpud/command/scan.go
@@ -29,6 +29,7 @@ func cmdScan(cliContext *cli.Context) error {
 		diagnose.WithPollXidEvents(pollXidEvents),
 		diagnose.WithPollGPMEvents(pollGPMEvents),
 		diagnose.WithNetcheck(netcheck),
+		diagnose.WithDiskcheck(diskcheck),
 	}
 	if zapLvl.Level() <= zap.DebugLevel { // e.g., info, warn, error
 		diagnoseOpts = append(diagnoseOpts, diagnose.WithDebug(true))
diff --git a/components/diagnose/options.go b/components/diagnose/options.go
index d4238fc3..69326bf3 100644
--- a/components/diagnose/options.go
+++ b/components/diagnose/options.go
@@ -8,7 +8,8 @@ type Op struct {
 	pollXidEvents bool
 	pollGPMEvents bool
 
-	netcheck bool
+	netcheck  bool
+	diskcheck bool
 }
 
 type OpOption func(*Op)
@@ -59,3 +60,9 @@ func WithNetcheck(b bool) OpOption {
 		op.netcheck = b
 	}
 }
+
+func WithDiskcheck(b bool) OpOption {
+	return func(op *Op) {
+		op.diskcheck = b
+	}
+}
diff --git a/components/diagnose/scan.go b/components/diagnose/scan.go
index b04e7917..41597e96 100644
--- a/components/diagnose/scan.go
+++ b/components/diagnose/scan.go
@@ -15,6 +15,7 @@ import (
 	query_log_common "github.com/leptonai/gpud/components/query/log/common"
 	query_log_tail "github.com/leptonai/gpud/components/query/log/tail"
 	"github.com/leptonai/gpud/log"
+	"github.com/leptonai/gpud/pkg/disk"
 	pkg_dmesg "github.com/leptonai/gpud/pkg/dmesg"
 	"github.com/leptonai/gpud/pkg/file"
 	latency_edge "github.com/leptonai/gpud/pkg/latency/edge"
@@ -238,6 +239,16 @@ func Scan(ctx context.Context, opts ...OpOption) error {
 		}
 	}
 
+	if op.diskcheck {
+		fmt.Printf("\n%s checking disk\n", inProgress)
+		partitions, err := disk.GetPartitions()
+		if err != nil {
+			log.Logger.Warnw("error getting partitions", "error", err)
+		} else {
+			partitions.RenderTable(os.Stdout)
+		}
+	}
+
 	fmt.Printf("\n\n%s scan complete\n\n", checkMark)
 	return nil
 }
diff --git a/pkg/disk/disk.go b/pkg/disk/disk.go
index 82d82dce..7e665f79 100644
--- a/pkg/disk/disk.go
+++ b/pkg/disk/disk.go
@@ -3,10 +3,12 @@ package disk
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 
 	"github.com/dustin/go-humanize"
+	"github.com/olekukonko/tablewriter"
 	"github.com/shirou/gopsutil/v4/disk"
 	"sigs.k8s.io/yaml"
 )
@@ -89,6 +91,25 @@ func (parts Partitions) TotalBytes() uint64 {
 	return total
 }
 
+func (parts Partitions) RenderTable(wr io.Writer) {
+	table := tablewriter.NewWriter(wr)
+	table.SetHeader([]string{"Device", "Mount", "Mounted", "Fstype", "Total", "Free", "Used"})
+
+	for _, part := range parts {
+		table.Append([]string{
+			part.Device,
+			part.MountPoint,
+			strconv.FormatBool(part.Mounted),
+			part.Fstype,
+			part.Usage.TotalHumanized,
+			part.Usage.FreeHumanized,
+			part.Usage.UsedHumanized,
+		})
+	}
+
+	table.Render()
+}
+
 type Partition struct {
 	Device     string `json:"device"`
 	MountPoint string `json:"mount_point"`
diff --git a/pkg/disk/disk_test.go b/pkg/disk/disk_test.go
index a7a346e1..0e7320d5 100644
--- a/pkg/disk/disk_test.go
+++ b/pkg/disk/disk_test.go
@@ -1,6 +1,9 @@
 package disk
 
-import "testing"
+import (
+	"os"
+	"testing"
+)
 
 func TestGetPartitions(t *testing.T) {
 	t.Parallel()
@@ -14,4 +17,6 @@ func TestGetPartitions(t *testing.T) {
 		t.Fatalf("failed to marshal partitions to yaml: %v", err)
 	}
 	t.Logf("partitions:\n%s\n", string(yb))
+
+	partitions.RenderTable(os.Stdout)
 }