allow modifying unit prefixes

coder · mtojek · Jun 20, 2023 · May 29, 2023 · Jun 7, 2023 · Jun 7, 2023
commit 665bf7fea59b2d1de10a737d10e6be19fa57165e
diff --git a/cli/clistat/cgroup.go b/cli/clistat/cgroup.go
@@ -45,8 +45,13 @@ const (
 )
 
 // ContainerCPU returns the CPU usage of the container cgroup.
+// This is calculated as difference of two samples of the
+// CPU usage of the container cgroup.
+// The total is read from the relevant path in /sys/fs/cgroup.
+// If there is no limit set, the total is assumed to be the
+// number of host cores multiplied by the CFS period.
 // If the system is not containerized, this always returns nil.
-func (s *Statter) ContainerCPU() (*Result, error) {
+func (s *Statter) ContainerCPU(m Prefix) (*Result, error) {
 	// Firstly, check if we are containerized.
 	if ok, err := IsContainerized(s.fs); err != nil || !ok {
 		return nil, nil //nolint: nilnil
@@ -61,6 +66,11 @@ func (s *Statter) ContainerCPU() (*Result, error) {
 	if err != nil {
 		return nil, xerrors.Errorf("get cgroup CPU usage: %w", err)
 	}
+
+	// The measurements in /sys/fs/cgroup are counters.
+	// We need to wait for a bit to get a difference.
+	// Note that someone could reset the counter in the meantime.
+	// We can't do anything about that.
 	s.wait(s.sampleInterval)
 
 	used2, err := s.cGroupCPUUsed()
@@ -69,9 +79,10 @@ func (s *Statter) ContainerCPU() (*Result, error) {
 	}
 
 	r := &Result{
-		Unit:  "cores",
-		Used:  (used2 - used1),
-		Total: ptr.To(total),
+		Unit:   "cores",
+		Prefix: m,
+		Used:   (used2 - used1),
+		Total:  ptr.To(total),
 	}
 	return r, nil
 }
@@ -169,20 +180,20 @@ func (s *Statter) cGroupV1CPUUsed() (float64, error) {
 
 // ContainerMemory returns the memory usage of the container cgroup.
 // If the system is not containerized, this always returns nil.
-func (s *Statter) ContainerMemory() (*Result, error) {
+func (s *Statter) ContainerMemory(m Prefix) (*Result, error) {
 	if ok, err := IsContainerized(s.fs); err != nil || !ok {
 		return nil, nil //nolint:nilnil
 	}
 
 	if s.isCGroupV2() {
-		return s.cGroupV2Memory()
+		return s.cGroupV2Memory(m)
 	}
 
 	// Fall back to CGroupv1
-	return s.cGroupV1Memory()
+	return s.cGroupV1Memory(m)
 }
 
-func (s *Statter) cGroupV2Memory() (*Result, error) {
+func (s *Statter) cGroupV2Memory(m Prefix) (*Result, error) {
 	maxUsageBytes, err := readInt64(s.fs, cgroupV2MemoryMaxBytes)
 	if err != nil {
 		return nil, xerrors.Errorf("read memory total: %w", err)
@@ -199,13 +210,14 @@ func (s *Statter) cGroupV2Memory() (*Result, error) {
 	}
 
 	return &Result{
-		Total: ptr.To(float64(maxUsageBytes) / 1024 / 1024 / 1024),
-		Used:  float64(currUsageBytes-inactiveFileBytes) / 1024 / 1024 / 1024,
-		Unit:  "GB",
+		Total:  ptr.To(float64(maxUsageBytes)),
+		Used:   float64(currUsageBytes - inactiveFileBytes),
+		Unit:   "B",
+		Prefix: m,
 	}, nil
 }
 
-func (s *Statter) cGroupV1Memory() (*Result, error) {
+func (s *Statter) cGroupV1Memory(m Prefix) (*Result, error) {
 	maxUsageBytes, err := readInt64(s.fs, cgroupV1MemoryMaxUsageBytes)
 	if err != nil {
 		return nil, xerrors.Errorf("read memory total: %w", err)
@@ -224,9 +236,10 @@ func (s *Statter) cGroupV1Memory() (*Result, error) {
 
 	// Total memory used is usage - total_inactive_file
 	return &Result{
-		Total: ptr.To(float64(maxUsageBytes) / 1024 / 1024 / 1024),
-		Used:  float64(usageBytes-totalInactiveFileBytes) / 1024 / 1024 / 1024,
-		Unit:  "GB",
+		Total:  ptr.To(float64(maxUsageBytes)),
+		Used:   float64(usageBytes - totalInactiveFileBytes),
+		Unit:   "B",
+		Prefix: m,
 	}, nil
 }
 

diff --git a/cli/clistat/disk.go b/cli/clistat/disk.go
@@ -10,7 +10,7 @@ import (
 
 // Disk returns the disk usage of the given path.
 // If path is empty, it returns the usage of the root directory.
-func (*Statter) Disk(path string) (*Result, error) {
+func (*Statter) Disk(path string, m Prefix) (*Result, error) {
 	if path == "" {
 		path = "/"
 	}
@@ -19,8 +19,9 @@ func (*Statter) Disk(path string) (*Result, error) {
 		return nil, err
 	}
 	var r Result
-	r.Total = ptr.To(float64(stat.Blocks*uint64(stat.Bsize)) / 1024 / 1024 / 1024)
-	r.Used = float64(stat.Blocks-stat.Bfree) * float64(stat.Bsize) / 1024 / 1024 / 1024
-	r.Unit = "GB"
+	r.Total = ptr.To(float64(stat.Blocks * uint64(stat.Bsize)))
+	r.Used = float64(stat.Blocks-stat.Bfree) * float64(stat.Bsize)
+	r.Unit = "B"
+	r.Prefix = m
 	return &r, nil
 }
diff --git a/cli/clistat/disk_windows.go b/cli/clistat/disk_windows.go
@@ -7,7 +7,7 @@ import (
 
 // Disk returns the disk usage of the given path.
 // If path is empty, it defaults to C:\
-func (*Statter) Disk(path string) (*Result, error) {
+func (*Statter) Disk(path string, m Prefix) (*Result, error) {
 	if path == "" {
 		path = `C:\`
 	}
@@ -28,8 +28,9 @@ func (*Statter) Disk(path string) (*Result, error) {
 	}
 
 	var r Result
-	r.Total = ptr.To(float64(totalBytes) / 1024 / 1024 / 1024)
-	r.Used = float64(totalBytes-freeBytes) / 1024 / 1024 / 1024
-	r.Unit = "GB"
+	r.Total = ptr.To(float64(totalBytes))
+	r.Used = float64(totalBytes - freeBytes)
+	r.Unit = "B"
+	r.Prefix = m
 	return &r, nil
 }
diff --git a/cli/clistat/stat.go b/cli/clistat/stat.go
@@ -14,6 +14,77 @@ import (
 	sysinfotypes "github.com/elastic/go-sysinfo/types"
 )
 
+// Prefix is an SI prefix for a unit.
+type Prefix string
+
+// Float64 returns the prefix as a float64.
+func (m *Prefix) Float64() (float64, error) {
+	switch *m {
+	case PrefixDeciShort, PrefixDeci:
+		return 0.1, nil
+	case PrefixCentiShort, PrefixCenti:
+		return 0.01, nil
+	case PrefixMilliShort, PrefixMilli:
+		return 0.001, nil
+	case PrefixMicroShort, PrefixMicro:
+		return 0.000_001, nil
+	case PrefixNanoShort, PrefixNano:
+		return 0.000_000_001, nil
+	case PrefixKiloShort, PrefixKilo:
+		return 1_000.0, nil
+	case PrefixMegaShort, PrefixMega:
+		return 1_000_000.0, nil
+	case PrefixGigaShort, PrefixGiga:
+		return 1_000_000_000.0, nil
+	case PrefixTeraShort, PrefixTera:
+		return 1_000_000_000_000.0, nil
+	case PrefixKibiShort, PrefixKibi:
+		return 1024.0, nil
+	case PrefixMebiShort, PrefixMebi:
+		return 1_048_576.0, nil
+	case PrefixGibiShort, PrefixGibi:
+		return 1_073_741_824.0, nil
+	case PrefixTebiShort, PrefixTebi:
+		return 1_099_511_627_776.0, nil
+	default:
+		return 0, xerrors.Errorf("unknown prefix: %s", *m)
+	}
+}
+
+const (
+	PrefixDeci  Prefix = "deci"
+	PrefixCenti Prefix = "centi"
+	PrefixMilli Prefix = "milli"
+	PrefixMicro Prefix = "micro"
+	PrefixNano  Prefix = "nano"
+
+	PrefixDeciShort  Prefix = "d"
+	PrefixCentiShort Prefix = "c"
+	PrefixMilliShort Prefix = "m"
+	PrefixMicroShort Prefix = "u"
+	PrefixNanoShort  Prefix = "n"
+
+	PrefixKilo Prefix = "kilo"
+	PrefixMega Prefix = "mega"
+	PrefixGiga Prefix = "giga"
+	PrefixTera Prefix = "tera"
+
+	PrefixKiloShort Prefix = "K"
+	PrefixMegaShort Prefix = "M"
+	PrefixGigaShort Prefix = "G"
+	PrefixTeraShort Prefix = "T"
+
+	PrefixKibi = "kibi"
+	PrefixMebi = "mebi"
+	PrefixGibi = "gibi"
+	PrefixTebi = "tebi"
+
+	PrefixKibiShort Prefix = "Ki"
+	PrefixMebiShort Prefix = "Mi"
+	PrefixGibiShort Prefix = "Gi"
+	PrefixTebiShort Prefix = "Ti"
+)
+
 // Result is a generic result type for a statistic.
 // Total is the total amount of the resource available.
 // It is nil if the resource is not a finite quantity.
@@ -23,6 +94,8 @@ type Result struct {
 	Total *float64 `json:"total"`
 	Unit  string   `json:"unit"`
 	Used  float64  `json:"used"`
+	// Prefix controls the string representation of the result.
+	Prefix Prefix `json:"-"`
 }
 
 // String returns a human-readable representation of the result.
@@ -31,13 +104,20 @@ func (r *Result) String() string {
 		return "-"
 	}
 	var sb strings.Builder
-	_, _ = sb.WriteString(strconv.FormatFloat(r.Used, 'f', 1, 64))
+	scale, err := r.Prefix.Float64()
+	prefix := string(r.Prefix)
+	if err != nil {
+		prefix = ""
+		scale = 1.0
+	}
+	_, _ = sb.WriteString(strconv.FormatFloat(r.Used/scale, 'f', 1, 64))
 	if r.Total != (*float64)(nil) {
 		_, _ = sb.WriteString("/")
-		_, _ = sb.WriteString(strconv.FormatFloat(*r.Total, 'f', 1, 64))
+		_, _ = sb.WriteString(strconv.FormatFloat(*r.Total/scale, 'f', 1, 64))
 	}
 	if r.Unit != "" {
 		_, _ = sb.WriteString(" ")
+		_, _ = sb.WriteString(prefix)
 		_, _ = sb.WriteString(r.Unit)
 	}
 	return sb.String()
@@ -96,7 +176,7 @@ func New(opts ...Option) (*Statter, error) {
 // This is calculated by taking the difference between the total and idle HostCPU time
 // and scaling it by the number of cores.
 // Units are in "cores".
-func (s *Statter) HostCPU() (*Result, error) {
+func (s *Statter) HostCPU(m Prefix) (*Result, error) {
 	r := &Result{
 		Unit:  "cores",
 		Total: ptr.To(float64(s.nproc)),
@@ -115,19 +195,21 @@ func (s *Statter) HostCPU() (*Result, error) {
 	used := total - idle
 	scaleFactor := float64(s.nproc) / total.Seconds()
 	r.Used = used.Seconds() * scaleFactor
+	r.Prefix = m
 	return r, nil
 }
 
 // HostMemory returns the memory usage of the host, in gigabytes.
-func (s *Statter) HostMemory() (*Result, error) {
+func (s *Statter) HostMemory(m Prefix) (*Result, error) {
 	r := &Result{
-		Unit: "GB",
+		Unit:   "B",
+		Prefix: m,
 	}
 	hm, err := s.hi.Memory()
 	if err != nil {
 		return nil, xerrors.Errorf("get memory info: %w", err)
 	}
-	r.Total = ptr.To(float64(hm.Total) / 1024 / 1024 / 1024)
-	r.Used = float64(hm.Used) / 1024 / 1024 / 1024
+	r.Total = ptr.To(float64(hm.Total))
+	r.Used = float64(hm.Used)
 	return r, nil
 }