Adding getDbSize functionality

clients/java/dkv-client/src/main/java/org/dkv/client/DKVClient.java

@@ -280,6 +280,12 @@ public interface DKVClient extends Closeable {
      */
     Iterator<DKVEntry> iterate(byte[] startKey, byte[] keyPref);
 
+    /**
+     * This methods returns the approximate count of keys in the DKV database
+     * @return retunrn the count of keys
+     */
+    long getDbSize();
+
     @Override
     void close();
 }

clients/java/dkv-client/src/main/java/org/dkv/client/ShardedDKVClient.java

@@ -260,6 +260,12 @@ public Iterator<DKVEntry> iterate(byte[] startKey, byte[] keyPref) {
         return dkvClient.iterate(startKey, keyPref);
     }
 
+    @Override
+    public long getDbSize() {
+        //TODO: implement this behaviour
+        return 0;
+    }
+
     @Override
     public void close() {
         pool.close();

clients/java/dkv-client/src/main/java/org/dkv/client/SimpleDKVClient.java

@@ -3,6 +3,7 @@
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.jmx.JmxReporter;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Empty;
 import dkv.serverpb.Api;
 import dkv.serverpb.DKVGrpc;
 import io.grpc.ManagedChannel;
@@ -272,6 +273,17 @@ public Iterator<DKVEntry> iterate(byte[] startKey, byte[] keyPref) {
         return iterate(copyFrom(startKey), copyFrom(keyPref));
     }
 
+    @Override
+    public long getDbSize() {
+        Empty emptyReq = Empty.newBuilder().build();
+        Api.KeySpaceSizeResponse res = blockingStub.getKeySpaceSize(emptyReq);
+        Api.Status status = res.getStatus();
+        if (status.getCode() != 0) {
+            throw new DKVException(status, "getDbSize", new Object[]{});
+        }
+        return res.getDbSize();
+    }
+
     @Override
     public void close() {
         reporter.stop();

cmd/dkvctl/main.go

@@ -32,6 +32,7 @@ var cmds = []*cmd{
 	{"removeNode", "<nexusUrl>", "Remove a master node from DKV cluster", (*cmd).removeNode, "", false},
 	{"listNodes", "", "Lists the various DKV nodes that are part of the Nexus cluster", (*cmd).listNodes, "", true},
 	{"getClusterInfo", "<dcId> <database> <vBucket>", "Gets the latest cluster info", (*cmd).getStatus, "", true},
+	{"getDbSize", "", "Fetched the approximate count of keys in the db", (*cmd).getDbSize, "", true},
 }
 
 func (c *cmd) usage() {
@@ -220,6 +221,15 @@ func (c *cmd) getStatus(client *ctl.DKVClient, args ...string) {
 	}
 }
 
+func (c *cmd) getDbSize(client *ctl.DKVClient, args ...string) {
+	dbSize, err := client.GetDbSize()
+	if err != nil {
+		fmt.Println("Unable to perform getDbSize. Error: %v\n", err)
+	}
+	fmt.Printf("%d\n", dbSize.DbSize)
+
+}
+
 var dkvAddr, dkvAuthority string
 
 func init() {

internal/master/service.go

@@ -174,6 +174,20 @@ func (ss *standaloneService) CompareAndSet(ctx context.Context, casReq *serverpb
 	return res, err
 }
 
+func (ss *standaloneService) GetKeySpaceSize(ctx context.Context, e *emptypb.Empty) (*serverpb.KeySpaceSizeResponse, error) {
+	ss.rwl.RLock()
+	defer ss.rwl.RUnlock()
+
+	dbSizeResults, err := ss.store.GetKeySpaceSize()
+	res := &serverpb.KeySpaceSizeResponse{Status: newEmptyStatus(), DbSize: dbSizeResults}
+	if err != nil {
+		ss.opts.Logger.Error("Unable to fetch DBKeySpaceSize", zap.Error(err))
+		res.Status = newErrorStatus(err)
+	}
+	res.DbSize = dbSizeResults
+	return res, err
+}
+
 func (ss *standaloneService) GetChanges(ctx context.Context, getChngsReq *serverpb.GetChangesRequest) (*serverpb.GetChangesResponse, error) {
 	ss.rwl.RLock()
 	defer ss.rwl.RUnlock()
@@ -417,6 +431,10 @@ func (ds *distributedService) CompareAndSet(ctx context.Context, casReq *serverp
 	return res, err
 }
 
+func (ds *distributedService) GetKeySpaceSize(ctx context.Context, req *emptypb.Empty) (*serverpb.KeySpaceSizeResponse, error) {
+	return nil, nil
+}
+
 func (ds *distributedService) Delete(ctx context.Context, delReq *serverpb.DeleteRequest) (*serverpb.DeleteResponse, error) {
 	reqBts, err := proto.Marshal(&raftpb.InternalRaftRequest{Delete: delReq})
 	res := &serverpb.DeleteResponse{Status: newEmptyStatus()}

internal/slave/service.go

@@ -108,6 +108,11 @@ func (ss *slaveService) CompareAndSet(_ context.Context, _ *serverpb.CompareAndS
 	return nil, errors.New("DKV slave service does not support keyspace mutations")
 }
 
+func (ss *slaveService) GetKeySpaceSize(ctx context.Context, empty *emptypb.Empty) (*serverpb.KeySpaceSizeResponse, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
 func (ss *slaveService) Get(ctx context.Context, getReq *serverpb.GetRequest) (*serverpb.GetResponse, error) {
 	readResults, err := ss.store.Get(getReq.Key)
 	res := &serverpb.GetResponse{Status: newEmptyStatus()}

internal/storage/badger/store.go

@@ -285,6 +285,10 @@ func (bdb *badgerDB) CompareAndSet(key, expect, update []byte) (bool, error) {
 	return err == nil, err
 }
 
+func (bdb *badgerDB) GetKeySpaceSize() (int64, error) {
+	return 0, nil
+}
+
 const (
 	badgerSSTPrefix = "badger-snapshot-"
 )

internal/storage/rocksdb/store.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -380,6 +381,17 @@ func (rdb *rocksDB) CompareAndSet(key, expect, update []byte) (bool, error) {
 	return err == nil, err
 }
 
+func (rdb *rocksDB) GetKeySpaceSize() (int64, error) {
+	defer rdb.opts.statsCli.Timing("rocksdb.property.latency.dbSize", time.Now())
+	keySize := rdb.db.GetProperty("rocksdb.estimate-num-keys")
+	dbKeySize, err := strconv.ParseInt(keySize, 10, 64)
+	if err != nil {
+		rdb.opts.statsCli.Incr("rocksdb.property.latency.dbSize", 1)
+		return 0, err
+	}
+	return dbKeySize, nil
+}
+
 const (
 	sstPrefix               = "rocksdb-sstfile-"
 	sstDefaultCF            = "/default.cf"

internal/storage/store.go

@@ -42,6 +42,8 @@ type KVStore interface {
 	// If the expected value is `nil`, then the key is created and
 	// initialized with the given value, atomically.
 	CompareAndSet(key, expect, update []byte) (bool, error)
+	//GetKeySpaceSize returns the estimated number of keys the db
+	GetKeySpaceSize() (int64, error)
 }
 
 // A Backupable represents the capability of the underlying store

internal/sync/repl_test.go

@@ -205,6 +205,10 @@ func (ms *memStore) CompareAndSet(key, expect, update []byte) (bool, error) {
 	return true, nil
 }
 
+func (ms *memStore) GetKeySpaceSize() (int64, error) {
+	return 0, nil
+}
+
 func (ms *memStore) Close() error {
 	ms.store = nil
 	return nil

pkg/ctl/client.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"github.com/flipkart-incubator/dkv/internal/hlc"
 	"github.com/flipkart-incubator/nexus/models"
+	"google.golang.org/protobuf/types/known/emptypb"
 	"io"
 	"time"
 
@@ -262,6 +263,13 @@ func (dkvClnt *DKVClient) Iterate(keyPrefix, startKey []byte) (<-chan *KVPair, e
 	return ch, nil
 }
 
+// GetDbSize returns the approximate count of the number of the keys in the db
+func (dkvClnt *DKVClient) GetDbSize() (*serverpb.KeySpaceSizeResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), Timeout)
+	defer cancel()
+	return dkvClnt.dkvCli.GetKeySpaceSize(ctx, &emptypb.Empty{})
+}
+
 // Close closes the underlying GRPC client connection to DKV service
 func (dkvClnt *DKVClient) Close() error {
 	if dkvClnt.cliConn != nil {