From 5224612df3694fa914a7218a548ad5b65292810d Mon Sep 17 00:00:00 2001
From: reesporte <reesedporter@gmail.com>
Date: Mon, 14 Mar 2022 12:35:45 -0500
Subject: [PATCH] add /internal/disk-usage endpoint for testing

will be used to ensure delete work doesn't result in an ever-inflating usage of
memory
---
 http_handler.go               | 20 ++++++++++++++++++++
 http_handler_internal_test.go | 21 +++++++++++++++++++++
 util.go                       | 31 +++++++++++++++++++++++++++++++
 util_test.go                  | 13 ++++++++++++-
 4 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/http_handler.go b/http_handler.go
index 3fc47cad4..9dc5319ae 100644
--- a/http_handler.go
+++ b/http_handler.go
@@ -468,6 +468,7 @@ func newRouter(handler *Handler) http.Handler {
 
 	// other ones
 	router.HandleFunc("/internal/mem-usage", handler.chkAuthZ(handler.handleGetMemUsage, authz.Read)).Methods("GET").Name("GetUsage")
+	router.HandleFunc("/internal/disk-usage", handler.chkAuthZ(handler.handleGetDiskUsage, authz.Read)).Methods("GET").Name("GetUsage")
 	router.HandleFunc("/internal/fragment/block/data", handler.chkAuthN(handler.handleGetFragmentBlockData)).Methods("GET").Name("GetFragmentBlockData")
 	router.HandleFunc("/internal/fragment/blocks", handler.chkAuthN(handler.handleGetFragmentBlocks)).Methods("GET").Name("GetFragmentBlocks")
 	router.HandleFunc("/internal/fragment/data", handler.chkAuthN(handler.handleGetFragmentData)).Methods("GET").Name("GetFragmentData")
@@ -1013,6 +1014,25 @@ func (h *Handler) handleGetMemUsage(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// handleGetDiskUsage handles GET /internal/disk-usage requests.
+func (h *Handler) handleGetDiskUsage(w http.ResponseWriter, r *http.Request) {
+	if !validHeaderAcceptJSON(r.Header) {
+		http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable)
+		return
+	}
+
+	use, err := GetDiskUsage(h.api.server.dataDir)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(use); err != nil {
+		h.logger.Errorf("write disk usage response error: %s", err)
+	}
+}
+
 // handleGetShardDistribution handles GET /ui/shard-distribution requests.
 func (h *Handler) handleGetShardDistribution(w http.ResponseWriter, r *http.Request) {
 	dist := h.api.ShardDistribution(r.Context())
diff --git a/http_handler_internal_test.go b/http_handler_internal_test.go
index 30bf3a782..882eebfe6 100644
--- a/http_handler_internal_test.go
+++ b/http_handler_internal_test.go
@@ -789,3 +789,24 @@ func TestHandleGetMemUsage(t *testing.T) {
 		t.Fatalf("expected %v, got %v", http.StatusOK, resp.StatusCode)
 	}
 }
+
+func TestHandleGetDiskUsage(t *testing.T) {
+	h := Handler{
+		logger:      logger.NewStandardLogger(os.Stdout),
+		queryLogger: logger.NewStandardLogger(os.Stdout),
+		api: &API{
+			server: &Server{
+				dataDir: t.TempDir(),
+			},
+		},
+	}
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "/whatever", nil)
+
+	h.handleGetDiskUsage(w, r)
+
+	resp := w.Result()
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("expected %v, got %v", http.StatusOK, resp.StatusCode)
+	}
+}
diff --git a/util.go b/util.go
index ce4323ec3..7b20a2224 100644
--- a/util.go
+++ b/util.go
@@ -5,7 +5,11 @@ package pilosa
 
 import (
 	"fmt"
+	"os"
+	"os/user"
+	"path/filepath"
 	"reflect"
+	"strings"
 	"time"
 
 	"github.com/shirou/gopsutil/v3/mem"
@@ -71,3 +75,30 @@ func GetMemoryUsage() (MemoryUsage, error) {
 	}
 	return MemoryUsage{Capacity: usage.Total, TotalUse: usage.Used}, nil
 }
+
+type DiskUsage struct {
+	Usage int64 `json:"usage"`
+}
+
+// GetDiskUsage gets the disk usage of the path
+func GetDiskUsage(path string) (DiskUsage, error) {
+	usr, _ := user.Current()
+	dir := usr.HomeDir
+	if path == "~" {
+		path = dir
+	} else if strings.HasPrefix(path, "~/") {
+		path = filepath.Join(dir, path[2:])
+	}
+
+	var size int64
+	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if !info.IsDir() {
+			size += info.Size()
+		}
+		return err
+	})
+	return DiskUsage{size}, err
+}
diff --git a/util_test.go b/util_test.go
index 9a1bfe7f9..9aaa78c57 100644
--- a/util_test.go
+++ b/util_test.go
@@ -2,7 +2,6 @@
 package pilosa
 
 // util_test.go has unit tests for utility functions from util.go
-//
 
 import (
 	"testing"
@@ -96,3 +95,15 @@ func TestGetMemoryUsage(t *testing.T) {
 		t.Fatalf("unexpected error getting memory usage: %v", err)
 	}
 }
+
+func TestGetDiskUsage(t *testing.T) {
+	tdir := t.TempDir()
+	if _, err := GetDiskUsage(tdir); err != nil {
+		t.Fatalf("unexpected error getting disk usage: %v", err)
+	}
+
+	if _, err := GetDiskUsage(""); err == nil {
+		t.Fatal("expected error getting disk usage but got nil")
+	}
+
+}