Merge pull request #934 from alanbernstein/cluster-usage

Collect size-on-disk usage data from all nodes

api.go

@@ -25,8 +25,6 @@ import (
 	"io/ioutil"
 	"math"
 	"net/url"
-	"os"
-	"path"
 	"sort"
 	"strconv"
 	"strings"
@@ -793,71 +791,64 @@ 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)
-	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
-		}
-		if api.holder.Txf().IsTxDatabasePath(file.Name()) {
-			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
+// NodeUsage represents all usage measurements for one node.
+type NodeUsage struct {
+	Disk DiskUsage `json:"bytesOnDisk"`
 }
 
-func diskUsage(fname string) (int64, error) {
-	var size int64
+// DiskUsage represents the storage space used on disk by one node.
+type DiskUsage struct {
+	Capacity uint64           `json:"capacity,omitempty"`
+	TotalUse int64            `json:"totalInUse"`
+	Indexes  map[string]int64 `json:"indexes"`
+}
 
-	dir, err := os.Open(fname)
+// Usage gets the disk usage per index, in a map[nodeID]NodeUsage
+func (api *API) Usage(ctx context.Context, remote bool) (map[string]NodeUsage, error) {
+	span, _ := tracing.StartSpanFromContext(ctx, "API.Usage")
+	defer span.Finish()
+
+	nodeUsages := make(map[string]NodeUsage)
+
+	indexSizes, err := api.holder.Txf().IndexSizes()
 	if err != nil {
-		return 0, errors.Wrap(err, "opening data subdirectory")
+		return nil, errors.Wrap(err, "getting index usage")
 	}
-	defer dir.Close()
-
-	files, err := dir.Readdir(-1)
-	if err != nil {
-		return 0, errors.Wrap(err, "reading data subdirectory")
+	var totalSize int64
+	for _, s := range indexSizes {
+		totalSize += s
 	}
 
-	for _, file := range files {
-		if file.IsDir() {
-			sz, err := diskUsage(path.Join(fname, file.Name()))
-			if err != nil {
-				return 0, err
+	capacity, err := api.server.systemInfo.DiskCapacity(api.holder.path)
+	if err != nil {
+		api.server.logger.Printf("couldn't read disk capacity: %s", err)
+	}
+
+	// Insert into result.
+	nodeUsage := NodeUsage{
+		Disk: DiskUsage{
+			Capacity: capacity,
+			TotalUse: totalSize,
+			Indexes:  indexSizes,
+		},
+	}
+	nodeUsages[api.server.nodeID] = nodeUsage
+
+	// Collect size on disk from remote nodes
+	if !remote {
+		nodes := api.cluster.Nodes()
+		for _, node := range nodes {
+			if node.ID == api.server.nodeID {
+				continue
 			}
-			size += sz
-		} else {
-			size += file.Size()
+			nodeUsage, err := api.server.defaultClient.GetNodeUsage(ctx, &node.URI)
+			if err != nil {
+				return nil, errors.Wrapf(err, "collecting disk usage from %s", node.URI)
+			}
+			nodeUsages[node.ID] = nodeUsage[node.ID]
 		}
 	}
-
-	return size, nil
+	return nodeUsages, nil
 }
 
 // RecalculateCaches forces all TopN caches to be updated.

client.go

@@ -81,6 +81,8 @@ type InternalClient interface {
 	FinishTransaction(ctx context.Context, id string) (*Transaction, error)
 	Transactions(ctx context.Context) (map[string]*Transaction, error)
 	GetTransaction(ctx context.Context, id string) (*Transaction, error)
+
+	GetNodeUsage(ctx context.Context, uri *URI) (map[string]NodeUsage, error)
 }
 
 //===============
@@ -227,3 +229,7 @@ func (n nopInternalClient) Transactions(ctx context.Context) (map[string]*Transa
 func (n nopInternalClient) GetTransaction(ctx context.Context, id string) (*Transaction, error) {
 	return nil, nil
 }
+
+func (n nopInternalClient) GetNodeUsage(ctx context.Context, uri *URI) (map[string]NodeUsage, error) {
+	return nil, nil
+}

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
+}

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200709123515-8e896a7af361
 	github.com/dustin/go-humanize v1.0.0
 	github.com/glycerine/idem v0.0.0-20190127113923-7a8083893311
-	github.com/glycerine/lmdb-go v1.9.32
+	github.com/glycerine/lmdb-go v1.9.34
 	github.com/go-ole/go-ole v1.2.4 // indirect
 	github.com/gogo/protobuf v1.2.1
 	github.com/golang/protobuf v1.3.3

go.sum

@@ -64,8 +64,8 @@ github.com/glycerine/goconvey v0.0.0-20190410193231-58a59202ab31 h1:gclg6gY70GLy
 github.com/glycerine/goconvey v0.0.0-20190410193231-58a59202ab31/go.mod h1:Ogl1Tioa0aV7gstGFO7KhffUsb9M4ydbEbbxpcEDc24=
 github.com/glycerine/idem v0.0.0-20190127113923-7a8083893311 h1:AAXH0ZvYIHHqU06ASy0H2tYAkAGrQlZvEy2QZrrtt4E=
 github.com/glycerine/idem v0.0.0-20190127113923-7a8083893311/go.mod h1:B72P/ZM99sNiCmaQJflpmMAF5LsDzStpLdWzn0+Vr2Y=
-github.com/glycerine/lmdb-go v1.9.32 h1:thLnzCykFcmn2rACYnwpR4ovYauLNKaAuk+xj7YMbS0=
-github.com/glycerine/lmdb-go v1.9.32/go.mod h1:DrPeeTGooMg6B7cjNSP14perptTJzzdBy5YoosthrRs=
+github.com/glycerine/lmdb-go v1.9.34 h1:0lymJjpdelYnIMcNzsKROfIaApt99zhaHtjDJTHjGkE=
+github.com/glycerine/lmdb-go v1.9.34/go.mod h1:DrPeeTGooMg6B7cjNSP14perptTJzzdBy5YoosthrRs=
 github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-kit/kit v0.9.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=

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{}

http/client.go

@@ -1247,6 +1247,37 @@ func (c *InternalClient) TranslateIDsNode(ctx context.Context, uri *pilosa.URI,
 	return tkresp.Keys, nil
 }
 
+// GetNodeUsage retrieves the size-on-disk information for the specified node.
+func (c *InternalClient) GetNodeUsage(ctx context.Context, uri *pilosa.URI) (map[string]pilosa.NodeUsage, error) {
+	u := uri.Path("/ui/usage?remote=true")
+	req, err := http.NewRequest("GET", u, nil)
+	if err != nil {
+		return nil, errors.Wrap(err, "creating request")
+	}
+
+	req.Header.Set("Accept", "application/json")
+	req.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
+
+	// Execute request against the host.
+	resp, err := c.executeRequest(req.WithContext(ctx))
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	// Read body and unmarshal response.
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, errors.Wrap(err, "reading")
+	}
+
+	nodeUsages := make(map[string]pilosa.NodeUsage) // map of size 1
+	if err := json.Unmarshal(body, &nodeUsages); err != nil {
+		return nil, fmt.Errorf("unmarshal response: %s", err)
+	}
+	return nodeUsages, nil
+}
+
 func (c *InternalClient) Transactions(ctx context.Context) (map[string]*pilosa.Transaction, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.Transactions")
 	defer span.Finish()

http/handler.go

@@ -665,34 +665,25 @@ func (h *Handler) handleGetUsage(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable)
 		return
 	}
-	usageIndexes, usageTotal, err := h.api.Usage()
+
+	q := r.URL.Query()
+	remoteStr := q.Get("remote")
+	var remote bool
+	if remoteStr == "true" {
+		remote = true
+	}
+
+	nodeUsages, err := h.api.Usage(r.Context(), remote)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 	}
 
-	disk := diskUsage{
-		Total:   usageTotal,
-		Indexes: usageIndexes,
-	}
-
-	usage := getUsageResponse{
-		Disk: disk,
-	}
-
 	w.Header().Set("Content-Type", "application/json")
-	if err := json.NewEncoder(w).Encode(usage); err != nil {
+	if err := json.NewEncoder(w).Encode(nodeUsages); err != nil {
 		h.logger.Printf("write status response error: %s", err)
 	}
 }
 
-type getUsageResponse struct {
-	Disk diskUsage `json:"bytesOnDisk"`
-}
-type diskUsage struct {
-	Total   int64            `json:"total"`
-	Indexes map[string]int64 `json:"indexes"`
-}
-
 // handleGetStatus handles GET /status requests.
 func (h *Handler) handleGetStatus(w http.ResponseWriter, r *http.Request) {
 	if !validHeaderAcceptJSON(r.Header) {

server/handler_test.go

@@ -385,13 +385,18 @@ func TestHandler_Endpoints(t *testing.T) {
 		w := httptest.NewRecorder()
 		h.ServeHTTP(w, test.MustNewHTTPRequest("GET", "/ui/usage", nil))
 		if w.Code != gohttp.StatusOK {
+			fmt.Printf("%+v\n", w.Body)
 			t.Fatalf("unexpected status code: %d", w.Code)
 		}
-		ret := mustJSONDecode(t, w.Body)
-		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)
+		nodeUsages := make(map[string]pilosa.NodeUsage)
+		if err := json.Unmarshal(w.Body.Bytes(), &nodeUsages); err != nil {
+			t.Fatalf("unmarshal")
+		}
+
+		for _, nodeUsage := range nodeUsages {
+			if len(nodeUsage.Disk.Indexes) != 2 {
+				t.Fatalf("wrong length index size list: %#v", nodeUsage.Disk.Indexes)
+			}
 		}
 	})
 

txfactory.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -564,6 +565,67 @@ func (f *TxFactory) DumpAll() {
 	f.dbPerShard.DumpAll()
 }
 
+func (f *TxFactory) IndexSizes() (index2bytes map[string]int64, err error) {
+	// Open storage directory.
+	index2bytes = make(map[string]int64)
+	dirName, err := expandDirName(f.holder.path)
+	if err != nil {
+		return index2bytes, errors.Wrap(err, "expanding data directory")
+	}
+
+	idxs := f.holder.Indexes()
+
+	for _, idx := range idxs {
+		index := idx.name
+		fullName := path.Join(dirName, index)
+		roaringAndMeta, err := directoryUsage(fullName)
+		if err != nil {
+			return index2bytes, errors.Wrap(err, "getting disk usage for roaring and meta")
+		}
+		fullName = index + ".index.txstores@@@"
+		rbfOrLmdb, err := directoryUsage(fullName)
+		if err != nil {
+			return index2bytes, errors.Wrap(err, "getting disk usage for backend")
+		}
+		index2bytes[index] = roaringAndMeta + rbfOrLmdb
+	}
+
+	return index2bytes, nil
+}
+
+func directoryUsage(fname string) (int64, error) {
+	if !DirExists(fname) {
+		return 0, nil
+	}
+
+	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 := directoryUsage(path.Join(fname, file.Name()))
+			if err != nil {
+				return 0, err
+			}
+			size += sz
+		} else {
+			size += file.Size()
+		}
+	}
+
+	return size, nil
+}
+
 func (f *TxFactory) CloseIndex(idx *Index) error {
 	// under roaring and all the new databases, this is a no-op.
 	//idx.Dump("CloseIndex")