Add options to enable cache in benchmark read (dgraph-io#1461)

Adds options in benchmark read to enable cache for blocks, block 
indices. It also adds a functionality to do a full DB scan using 
iterators. 

Fixes the argument of cache size in benchmark write to accept the size 
in MB instead of Bytes

badger/cmd/read_bench.go

@@ -46,10 +46,16 @@ var (
 	entriesRead uint64    // will store entries read till now
 	startTime   time.Time // start time of read benchmarking
 
+	sizeMaxCache      int64
+	keepBlockIdxCache bool
+	keepBlocksCache   bool
+	sizeMaxBfCache    int64
+
 	sampleSize  int
 	loadingMode string
 	keysOnly    bool
 	readOnly    bool
+	fullScan    bool
 )
 
 func init() {
@@ -67,6 +73,35 @@ func init() {
 	readBenchCmd.Flags().StringVar(
 		&loadingMode, "loading-mode", "mmap", "Mode for accessing SSTables and value log files. "+
 			"Valid loading modes are fileio and mmap.")
+	readBenchCmd.Flags().BoolVar(
+		&fullScan, "full-scan", false, "If true, full db will be scanned using iterators.")
+	readBenchCmd.Flags().Int64VarP(&sizeMaxCache, "max-cache", "C", 0, "Max size of cache in MB")
+	readBenchCmd.Flags().BoolVarP(&keepBlockIdxCache, "keep-bidx", "b", false,
+		"Keep block indices in cache")
+	readBenchCmd.Flags().BoolVarP(&keepBlocksCache, "keep-blocks", "B", false,
+		"Keep blocks in cache")
+	readBenchCmd.Flags().Int64VarP(&sizeMaxBfCache, "max-bf-cache", "c", 0,
+		"Maximum Bloom Filter Cache Size in MB")
+}
+
+// Scan the whole database using the iterators
+func fullScanDB(db *badger.DB) {
+	txn := db.NewTransaction(false)
+	defer txn.Discard()
+
+	startTime = time.Now()
+	// Print the stats
+	c := y.NewCloser(0)
+	c.AddRunning(1)
+	go printStats(c)
+
+	it := txn.NewIterator(badger.DefaultIteratorOptions)
+	defer it.Close()
+	for it.Rewind(); it.Valid(); it.Next() {
+		i := it.Item()
+		atomic.AddUint64(&entriesRead, 1)
+		atomic.AddUint64(&sizeRead, uint64(i.EstimatedSize()))
+	}
 }
 
 func readBench(cmd *cobra.Command, args []string) error {
@@ -82,16 +117,28 @@ func readBench(cmd *cobra.Command, args []string) error {
 		WithValueDir(vlogDir).
 		WithReadOnly(readOnly).
 		WithTableLoadingMode(mode).
-		WithValueLogLoadingMode(mode)
-
+		WithValueLogLoadingMode(mode).
+		WithMaxCacheSize(sizeMaxCache << 20).
+		WithKeepBlockIndicesInCache(keepBlockIdxCache).
+		WithKeepBlocksInCache(keepBlocksCache).
+		WithMaxBfCacheSize(sizeMaxBfCache << 20)
 	fmt.Printf("Opening badger with options = %+v\n", opt)
 	db, err := badger.Open(opt)
 	if err != nil {
 		return y.Wrapf(err, "unable to open DB")
 	}
 	defer db.Close()
-
 	now := time.Now()
+
+	fmt.Println("*********************************************************")
+	fmt.Println("Starting to benchmark Reads")
+	fmt.Println("*********************************************************")
+
+	// if fullScan is true then do a complete scan of the db and return
+	if fullScan {
+		fullScanDB(db)
+		return nil
+	}
 	keys, err := getSampleKeys(db)
 	if err != nil {
 		return y.Wrapf(err, "error while sampling keys")
@@ -104,10 +151,6 @@ func readBench(cmd *cobra.Command, args []string) error {
 		fmt.Println("DB is empty, hence returning")
 		return nil
 	}
-
-	fmt.Println("*********************************************************")
-	fmt.Println("Starting to benchmark Reads")
-	fmt.Println("*********************************************************")
 	c := y.NewCloser(0)
 	startTime = time.Now()
 	for i := 0; i < numGoroutines; i++ {

badger/cmd/write_bench.go

@@ -102,13 +102,13 @@ func init() {
 	writeBenchCmd.Flags().BoolVarP(&showLogs, "logs", "l", false, "Show Badger logs.")
 	writeBenchCmd.Flags().IntVarP(&valueThreshold, "value-th", "t", 1<<10, "Value threshold")
 	writeBenchCmd.Flags().IntVarP(&numVersions, "num-version", "n", 1, "Number of versions to keep")
-	writeBenchCmd.Flags().Int64VarP(&maxCacheSize, "max-cache", "C", 0, "Max size of cache")
+	writeBenchCmd.Flags().Int64VarP(&maxCacheSize, "max-cache", "C", 0, "Max size of cache in MB")
 	writeBenchCmd.Flags().BoolVarP(&keepBlockIdxInCache, "keep-bidx", "b", false,
 		"Keep block indices in cache")
 	writeBenchCmd.Flags().BoolVarP(&keepBlocksInCache, "keep-blocks", "B", false,
 		"Keep blocks in cache")
 	writeBenchCmd.Flags().Int64VarP(&maxBfCacheSize, "max-bf-cache", "c", 0,
-		"Maximum Bloom Filter Cache Size")
+		"Maximum Bloom Filter Cache Size in MB")
 	writeBenchCmd.Flags().Uint32Var(&vlogMaxEntries, "vlog-maxe", 1000000, "Value log Max Entries")
 	writeBenchCmd.Flags().StringVarP(&encryptionKey, "encryption-key", "e", "",
 		"If it is true, badger will encrypt all the data stored on the disk.")
@@ -249,10 +249,10 @@ func writeBench(cmd *cobra.Command, args []string) error {
 		WithCompactL0OnClose(force).
 		WithValueThreshold(valueThreshold).
 		WithNumVersionsToKeep(numVersions).
-		WithMaxCacheSize(maxCacheSize).
+		WithMaxCacheSize(maxCacheSize << 20).
 		WithKeepBlockIndicesInCache(keepBlockIdxInCache).
 		WithKeepBlocksInCache(keepBlocksInCache).
-		WithMaxBfCacheSize(maxBfCacheSize).
+		WithMaxBfCacheSize(maxBfCacheSize << 20).
 		WithValueLogMaxEntries(vlogMaxEntries).
 		WithTableLoadingMode(mode).
 		WithEncryptionKey([]byte(encryptionKey)).