From 474261f12f1a513b4aa765006c470b9445b9e450 Mon Sep 17 00:00:00 2001
From: Alan Bernstein <alan.aaron.bernstein@gmail.com>
Date: Wed, 30 Sep 2020 07:14:35 -0500
Subject: [PATCH] Add single-node 'bytesOnDisk' object to /status response

---
 api.go                 | 66 ++++++++++++++++++++++++++++++++++++++++++
 http/handler.go        | 27 +++++++++++++----
 server/handler_test.go |  5 ++++
 3 files changed, 92 insertions(+), 6 deletions(-)

diff --git a/api.go b/api.go
index b6e8187c5..a307e323d 100644
--- a/api.go
+++ b/api.go
@@ -25,6 +25,8 @@ import (
 	"io/ioutil"
 	"math"
 	"net/url"
+	"os"
+	"path"
 	"sort"
 	"strconv"
 	"strings"
@@ -780,6 +782,70 @@ func (api *API) Node() *Node {
 	return &node
 }
 
+// Usage gets the disk usage per index
+func (api *API) Usage() (map[string]int64, int64, error) {
+	indexSizes := make(map[string]int64, 0)
+	var totalSize int64
+
+	dirName, err := expandDirName(api.server.dataDir)
+	if err != nil {
+		return indexSizes, totalSize, errors.Wrap(err, "expanding data directory")
+	}
+	dir, err := os.Open(dirName)
+	if err != nil {
+		return indexSizes, totalSize, errors.Wrap(err, "opening data directory")
+	}
+	defer dir.Close()
+
+	files, err := dir.Readdir(-1)
+	if err != nil {
+		return indexSizes, totalSize, errors.Wrap(err, "reading data directory")
+	}
+
+	for _, file := range files {
+		if !file.IsDir() {
+			continue
+		}
+		fullName := path.Join(dirName, file.Name())
+		indexSizes[file.Name()], err = diskUsage(fullName)
+		if err != nil {
+			break
+		}
+		totalSize += indexSizes[file.Name()]
+	}
+
+	return indexSizes, totalSize, nil
+}
+
+func diskUsage(fname string) (int64, error) {
+	var size int64
+
+	dir, err := os.Open(fname)
+	if err != nil {
+		return 0, errors.Wrap(err, "opening data subdirectory")
+	}
+	defer dir.Close()
+
+	files, err := dir.Readdir(-1)
+	if err != nil {
+		return 0, errors.Wrap(err, "reading data subdirectory")
+	}
+
+	for _, file := range files {
+		if file.IsDir() {
+			sz, err := diskUsage(path.Join(fname, file.Name()))
+			if err != nil {
+				return 0, err
+			}
+			size += sz
+		} else {
+			size += file.Size()
+		}
+	}
+
+	return size, nil
+}
+
 // RecalculateCaches forces all TopN caches to be updated.
 // This is done internally within a TopN query, but a user may want to do it ahead of time?
 func (api *API) RecalculateCaches(ctx context.Context) error {
diff --git a/http/handler.go b/http/handler.go
index c1020c941..606aaffa4 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -663,10 +663,19 @@ func (h *Handler) handleGetStatus(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable)
 		return
 	}
+	usageIndexes, usageTotal, err := h.api.Usage()
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+	usage := diskUsage{
+		Total:   usageTotal,
+		Indexes: usageIndexes,
+	}
 	status := getStatusResponse{
-		State:   h.api.State(),
-		Nodes:   h.api.Hosts(r.Context()),
-		LocalID: h.api.Node().ID,
+		State:       h.api.State(),
+		Nodes:       h.api.Hosts(r.Context()),
+		LocalID:     h.api.Node().ID,
+		BytesOnDisk: usage,
 	}
 	w.Header().Set("Content-Type", "application/json")
 	if err := json.NewEncoder(w).Encode(status); err != nil {
@@ -722,9 +731,15 @@ type getSchemaResponse struct {
 }
 
 type getStatusResponse struct {
-	State   string         `json:"state"`
-	Nodes   []*pilosa.Node `json:"nodes"`
-	LocalID string         `json:"localID"`
+	State       string         `json:"state"`
+	Nodes       []*pilosa.Node `json:"nodes"`
+	LocalID     string         `json:"localID"`
+	BytesOnDisk diskUsage      `json:"bytesOnDisk"`
+}
+
+type diskUsage struct {
+	Total   int64            `json:"total"`
+	Indexes map[string]int64 `json:"indexes"`
 }
 
 func hash(s string) string {
diff --git a/server/handler_test.go b/server/handler_test.go
index 4a1323452..c093d873d 100644
--- a/server/handler_test.go
+++ b/server/handler_test.go
@@ -379,6 +379,11 @@ func TestHandler_Endpoints(t *testing.T) {
 		if len(ret["nodes"].([]interface{})) != 1 {
 			t.Fatalf("wrong length nodes list: %#v", ret)
 		}
+		usage := ret["bytesOnDisk"].(map[string]interface{})
+		indexes := usage["indexes"].(map[string]interface{})
+		if len(indexes) != 2 {
+			t.Fatalf("wrong length index size list: %#v", indexes)
+		}
 	})
 
 	t.Run("Metrics", func(t *testing.T) {