Refactor size formatting for readability

1 files changed, 9 insertions, 13 deletions

diff --git a/util/common/format.go b/util/common/format.go
index 1ea10877..c73e3a01 100644
--- a/util/common/format.go
+++ b/util/common/format.go
@@ -4,18 +4,14 @@ import (
 	"fmt"
 )
 
-func FormatTraffic(trafficBytes int64) (size string) {
-	if trafficBytes < 1024 {
-		return fmt.Sprintf("%.2fB", float64(trafficBytes)/float64(1))
-	} else if trafficBytes < (1024 * 1024) {
-		return fmt.Sprintf("%.2fKB", float64(trafficBytes)/float64(1024))
-	} else if trafficBytes < (1024 * 1024 * 1024) {
-		return fmt.Sprintf("%.2fMB", float64(trafficBytes)/float64(1024*1024))
-	} else if trafficBytes < (1024 * 1024 * 1024 * 1024) {
-		return fmt.Sprintf("%.2fGB", float64(trafficBytes)/float64(1024*1024*1024))
-	} else if trafficBytes < (1024 * 1024 * 1024 * 1024 * 1024) {
-		return fmt.Sprintf("%.2fTB", float64(trafficBytes)/float64(1024*1024*1024*1024))
-	} else {
-		return fmt.Sprintf("%.2fEB", float64(trafficBytes)/float64(1024*1024*1024*1024*1024))
+func FormatTraffic(trafficBytes int64) string {
+	units := []string{"B", "KB", "MB", "GB", "TB", "PB"}
+	unitIndex := 0
+	size := float64(trafficBytes)
+
+	for size >= 1024 && unitIndex < len(units)-1 {
+		size /= 1024
+		unitIndex++
 	}
+	return fmt.Sprintf("%.2f%s", size, units[unitIndex])
 }