From 29f052acbec1cc013277e68b2e939b6acbc0d730 Mon Sep 17 00:00:00 2001
From: MOCHIZUKI Natsune <me@natsuneko.cat>
Date: Mon, 12 Aug 2024 20:27:04 +0900
Subject: [PATCH] feat: get savedata (.dat) size

---
 main.go | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index bc31cc3..bc3ff90 100644
--- a/main.go
+++ b/main.go
@@ -4,12 +4,16 @@ import (
 	"flag"
 	mp "github.com/mackerelio/go-mackerel-plugin"
 	"github.com/sch8ill/mclib"
+	"os"
+	filepath2 "path/filepath"
+	"strings"
 )
 
 // Mackerel Agent
 type MinecraftPlugin struct {
-	Prefix string
-	Server string
+	Prefix   string
+	Server   string
+	SaveData string
 }
 
 func (mc MinecraftPlugin) FetchMetrics() (map[string]float64, error) {
@@ -23,10 +27,57 @@ func (mc MinecraftPlugin) FetchMetrics() (map[string]float64, error) {
 		return nil, err
 	}
 
+	var overworldSize int64 = 0
+	var netherSize int64 = 0
+	var theEndSize int64 = 0
+
+	if mc.SaveData != "" {
+		path, _ := filepath2.Abs(mc.SaveData)
+		if _, err := os.Stat(path); os.IsNotExist(err) {
+			return nil, err
+		}
+
+		err = filepath2.Walk(path, func(data string, info os.FileInfo, err error) error {
+			if filepath2.Ext(data) == ".dat" {
+				rel, err := filepath2.Rel(path, data)
+				if err != nil {
+					return err
+				}
+
+				// Dimension: -1 is Nether
+				if strings.HasPrefix(rel, "DIM-1") {
+					netherSize += info.Size()
+					return nil
+				}
+
+				// Dimension: 1 is The End
+				if strings.HasPrefix(rel, "DIM1") {
+					theEndSize += info.Size()
+					return nil
+				}
+
+				// Dimension: 0 is Overworld
+				if strings.HasPrefix(rel, "data") {
+					overworldSize += info.Size()
+					return nil
+				}
+			}
+
+			return nil
+		})
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return map[string]float64{
 		"max_players":    float64(res.Players.Max),
 		"online_players": float64(res.Players.Online),
 		"latency":        float64(res.Latency),
+		"overworldSize":  float64(overworldSize) / float64(1024),
+		"netherSize":     float64(netherSize) / float64(1024),
+		"the_end_size":   float64(theEndSize) / float64(1024),
 	}, nil
 }
 
@@ -41,6 +92,9 @@ func (mc MinecraftPlugin) GraphDefinition() map[string]mp.Graphs {
 				{Name: "max_players", Label: "Limit"},
 				{Name: "online_players", Label: "Current Players"},
 				{Name: "latency", Label: "Latency"},
+				{Name: "overworld_size", Label: "Overworld DataSize"},
+				{Name: "nether_size", Label: "Nether DataSize"},
+				{Name: "the_end_size", Label: "TheEnd DataSize"},
 			},
 		},
 	}
@@ -65,13 +119,15 @@ func (mc MinecraftPlugin) GetServerAddress() string {
 func main() {
 	prefix := flag.String("metric-key-prefix", "minecraft", "Metric key prefix")
 	server := flag.String("server", "localhost:25565", "Server address")
+	savedata := flag.String("savedata", "", "Minecraft saved data dir")
 	tempfile := flag.String("tempfile", "", "Temp file name")
 
 	flag.Parse()
 
 	mc := MinecraftPlugin{
-		Prefix: *prefix,
-		Server: *server,
+		Prefix:   *prefix,
+		Server:   *server,
+		SaveData: *savedata,
 	}
 
 	helper := mp.NewMackerelPlugin(mc)