From 7a32eadc6487aa4a4c95333a01c6d212d69dff00 Mon Sep 17 00:00:00 2001
From: Alan Bernstein <alan.aaron.bernstein@gmail.com>
Date: Fri, 9 Oct 2020 10:21:37 -0500
Subject: [PATCH] Move disk capacity lookup to gopsutil wrapper package

---
 api.go                 |  5 +----
 diagnostics.go         |  6 ++++++
 gopsutil/systeminfo.go | 11 +++++++++++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/api.go b/api.go
index 180988a7e..4d52ab1d9 100644
--- a/api.go
+++ b/api.go
@@ -38,7 +38,6 @@ import (
 	"github.com/pilosa/pilosa/v2/stats"
 	"github.com/pilosa/pilosa/v2/tracing"
 	"github.com/pkg/errors"
-	"github.com/shirou/gopsutil/disk"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -847,11 +846,9 @@ func (api *API) Usage(ctx context.Context, remote bool) (map[string]NodeUsage, e
 		totalSize += indexSizes[file.Name()]
 	}
 
-	usageStats, err := disk.Usage("/")
-	capacity := usageStats.Total
+	capacity, err := api.server.systemInfo.DiskCapacity(api.server.dataDir)
 
 	if err != nil {
-		capacity = uint64(0)
 		api.server.logger.Printf("failed to get disk capacity: %s", err)
 	}
 
diff --git a/diagnostics.go b/diagnostics.go
index 9c380c1b8..3899eb14d 100644
--- a/diagnostics.go
+++ b/diagnostics.go
@@ -275,6 +275,7 @@ type SystemInfo interface {
 	CPUCores() (physical int, logical int, err error)
 	CPUMHz() (int, error)
 	CPUArch() string
+	DiskCapacity(string) (uint64, error)
 }
 
 // newNopSystemInfo creates a no-op implementation of SystemInfo.
@@ -345,3 +346,8 @@ func (n *nopSystemInfo) CPUMHz() (int, error) {
 func (n *nopSystemInfo) CPUCores() (physical, logical int, err error) {
 	return 0, 0, nil
 }
+
+// DiskCapacity returns the disk capacity
+func (n *nopSystemInfo) DiskCapacity(path string) (uint64, error) {
+	return 0, nil
+}
diff --git a/gopsutil/systeminfo.go b/gopsutil/systeminfo.go
index b251c39a4..312dd3f36 100644
--- a/gopsutil/systeminfo.go
+++ b/gopsutil/systeminfo.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/pilosa/pilosa/v2"
 	"github.com/shirou/gopsutil/cpu"
+	"github.com/shirou/gopsutil/disk"
 	"github.com/shirou/gopsutil/host"
 	"github.com/shirou/gopsutil/mem"
 )
@@ -243,6 +244,16 @@ func (s *systemInfo) CPUCores() (physical, logical int, err error) {
 	return s.cpuPhysicalCores, s.cpuLogicalCores, nil
 }
 
+// DiskCapacity returns the disk capacity.
+func (s *systemInfo) DiskCapacity(path string) (uint64, error) {
+	diskInfo, err := disk.Usage(path)
+
+	if err != nil {
+		return 0, err
+	}
+	return diskInfo.Total, nil
+}
+
 // NewSystemInfo is a constructor for the gopsutil implementation of SystemInfo.
 func NewSystemInfo() *systemInfo {
 	return &systemInfo{}