From 18bddca86f2efe32157404be49ec6392565f33b9 Mon Sep 17 00:00:00 2001
From: reesporte <reesedporter@gmail.com>
Date: Mon, 28 Feb 2022 12:05:30 -0600
Subject: [PATCH] add get internal mem usage endpoint

for use in benchmarking deletes
---
 http_handler.go               | 20 ++++++++++++++++++++
 http_handler_internal_test.go | 16 ++++++++++++++++
 util.go                       | 17 +++++++++++++++++
 util_test.go                  |  6 ++++++
 4 files changed, 59 insertions(+)

diff --git a/http_handler.go b/http_handler.go
index fdbbb0688..13ecd1f95 100644
--- a/http_handler.go
+++ b/http_handler.go
@@ -465,6 +465,7 @@ func newRouter(handler *Handler) http.Handler {
 	router.HandleFunc("/internal/translate/data", handler.chkAuthZ(handler.handlePostTranslateData, authz.Write)).Methods("POST").Name("PostTranslateData")
 
 	// other ones
+	router.HandleFunc("/internal/mem-usage", handler.chkAuthZ(handler.handleGetMemUsage, 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")
@@ -986,6 +987,25 @@ func (h *Handler) handlePostSchema(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusNoContent)
 }
 
+// handleGetMemUsage handles GET /internal/mem-usage requests.
+func (h *Handler) handleGetMemUsage(w http.ResponseWriter, r *http.Request) {
+	if !validHeaderAcceptJSON(r.Header) {
+		http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable)
+		return
+	}
+
+	use, err := GetMemoryUsage()
+	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 mem 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 455a40c64..59ae88e66 100644
--- a/http_handler_internal_test.go
+++ b/http_handler_internal_test.go
@@ -771,3 +771,19 @@ func NewTestAuth(t *testing.T) *authn.Auth {
 	}
 	return a
 }
+
+func TestHandleGetMemUsage(t *testing.T) {
+	h := Handler{
+		logger:      logger.NewStandardLogger(os.Stdout),
+		queryLogger: logger.NewStandardLogger(os.Stdout),
+	}
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "/whatever", nil)
+
+	h.handleGetMemUsage(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 eb9f958ce..ce4323ec3 100644
--- a/util.go
+++ b/util.go
@@ -4,8 +4,11 @@ package pilosa
 // util.go: a place for generic, reusable utilities.
 
 import (
+	"fmt"
 	"reflect"
 	"time"
+
+	"github.com/shirou/gopsutil/v3/mem"
 )
 
 // LeftShifted16MaxContainerKey is 0xffffffffffff0000. It is similar
@@ -54,3 +57,17 @@ func GetLoopProgress(start time.Time, now time.Time, iteration uint, total uint)
 func FormatTimestampNano(value, base int64, timeUnit string) string {
 	return time.Unix(0, (value+base)*TimeUnitNanos(timeUnit)).UTC().Format(time.RFC3339Nano)
 }
+
+type MemoryUsage struct {
+	Capacity uint64 `json:"capacity"`
+	TotalUse uint64 `json:"totalUsed"`
+}
+
+// GetMemoryUsage gets the memory usage
+func GetMemoryUsage() (MemoryUsage, error) {
+	usage, err := mem.VirtualMemory()
+	if usage == nil || err != nil {
+		return MemoryUsage{}, fmt.Errorf("reading virtual memory: %v", err)
+	}
+	return MemoryUsage{Capacity: usage.Total, TotalUse: usage.Used}, nil
+}
diff --git a/util_test.go b/util_test.go
index 870625ad8..9a1bfe7f9 100644
--- a/util_test.go
+++ b/util_test.go
@@ -90,3 +90,9 @@ func TestFormatTimestampNano(t *testing.T) {
 		t.Fatal("Timestamp not formatted properly")
 	}
 }
+
+func TestGetMemoryUsage(t *testing.T) {
+	if _, err := GetMemoryUsage(); err != nil {
+		t.Fatalf("unexpected error getting memory usage: %v", err)
+	}
+}