feat: add --cache_mb and --cache_percentage flag

dgraph/cmd/alpha/run.go

@@ -203,6 +203,11 @@ they form a Raft group and provide synchronous replication.
 	flag.Bool("ludicrous_mode", false, "Run alpha in ludicrous mode")
 	flag.Bool("graphql_extensions", true, "Set to false if extensions not required in GraphQL response body")
 	flag.Duration("graphql_poll_interval", time.Second, "polling interval for graphql subscription.")
+
+	// Add cache flags
+	flag.Int64("cache_mb", 0, "Total size of cache (in MB) to be used in alpha")
+	flag.String("cache_percentage", "30:30:20:20", `Cache percentages for various caches
+		(PstoreBlockCache:PstoreIndexCache:PstoreIndexCache:PostingListCache)`)
 }
 
 func setupCustomTokenizers() {
@@ -577,10 +582,24 @@ func run() {
 	}
 	bindall = Alpha.Conf.GetBool("bindall")
 
+	totalCache := int64(Alpha.Conf.GetInt("cache_mb"))
+	x.AssertTruef(totalCache >= 0, "ERROR: Cache size must be non-negative")
+
+	cachePercentage := Alpha.Conf.GetString("cache_percentage")
+	cachePercent := x.GetCachePercentages(cachePercentage,
+		"PstoreBlockCache:PstoreIndexCache:PstoreIndexCache:PostingListCache")
+	pstoreBlockCacheSize := (cachePercent[0] * (totalCache << 20)) / 100
+	pstoreIndexCacheSize := (cachePercent[1] * (totalCache << 20)) / 100
+	wstoreIndexCacheSize := (cachePercent[2] * (totalCache << 20)) / 100
+	postingListCacheSize := (cachePercent[3] * (totalCache << 20)) / 100
+
 	opts := worker.Options{
 		BadgerCompressionLevel: Alpha.Conf.GetInt("badger.compression_level"),
 		PostingDir:             Alpha.Conf.GetString("postings"),
 		WALDir:                 Alpha.Conf.GetString("wal"),
+		PBlockCacheSize:        pstoreBlockCacheSize,
+		PIndexCacheSize:        pstoreIndexCacheSize,
+		WIndexCacheSize:        wstoreIndexCacheSize,
 
 		MutationsMode:  worker.AllowMutations,
 		AuthToken:      Alpha.Conf.GetString("auth_token"),
@@ -701,7 +720,7 @@ func run() {
 	// Posting will initialize index which requires schema. Hence, initialize
 	// schema before calling posting.Init().
 	schema.Init(worker.State.Pstore)
-	posting.Init(worker.State.Pstore)
+	posting.Init(worker.State.Pstore, postingListCacheSize)
 	defer posting.Cleanup()
 	worker.Init(worker.State.Pstore)
 

dgraph/cmd/bulk/reduce.go

@@ -84,10 +84,10 @@ func (r *reducer) run() error {
 			splitWriter := tmpDb.NewManagedWriteBatch()
 
 			ci := &countIndexer{
-				reducer: r,
-				writer: writer,
+				reducer:     r,
+				writer:      writer,
 				splitWriter: splitWriter,
-				tmpDb: tmpDb,
+				tmpDb:       tmpDb,
 			}
 			sort.Slice(partitionKeys, func(i, j int) bool {
 				return bytes.Compare(partitionKeys[i], partitionKeys[j]) < 0
@@ -130,7 +130,7 @@ func (r *reducer) createBadgerInternal(dir string, compression bool) *badger.DB
 
 	opt := badger.DefaultOptions(dir).WithSyncWrites(false).
 		WithTableLoadingMode(bo.MemoryMap).WithValueThreshold(1 << 10 /* 1 KB */).
-		WithLogger(nil).WithMaxCacheSize(1 << 20).
+		WithLogger(nil).WithBlockCacheSize(1 << 20).
 		WithEncryptionKey(r.opt.EncryptionKey)
 
 	// Overwrite badger options based on the options provided by the user.

dgraph/cmd/debug/run.go

@@ -784,7 +784,8 @@ func run() {
 		db, err = badger.OpenManaged(bopts)
 	}
 	x.Check(err)
-	posting.Init(db)
+	// Not using posting list cache
+	posting.Init(db, 0)
 	defer db.Close()
 
 	if isWal {

dgraph/cmd/zero/run.go

@@ -18,6 +18,7 @@ package zero
 
 import (
 	"context"
+
 	//	"errors"
 	"fmt"
 	"log"
@@ -55,6 +56,9 @@ type options struct {
 	w                 string
 	rebalanceInterval time.Duration
 	LudicrousMode     bool
+
+	totalCache      int64
+	cachePercentage string
 }
 
 var opts options
@@ -101,6 +105,10 @@ instances to achieve high-availability.
 		" exporter does not support annotation logs and would discard them.")
 	flag.Bool("ludicrous_mode", false, "Run zero in ludicrous mode")
 	flag.String("enterprise_license", "", "Path to the enterprise license file.")
+
+	// Add cache flags
+	flag.Int64("cache_mb", 0, "Total size of cache (in MB) to be used in zero")
+	flag.String("cache_percentage", "100:0", "Cache percentages for various caches (blockcache:indexCache)")
 }
 
 func setupListener(addr string, port int, kind string) (listener net.Listener, err error) {
@@ -183,6 +191,8 @@ func run() {
 		w:                 Zero.Conf.GetString("wal"),
 		rebalanceInterval: Zero.Conf.GetDuration("rebalance_interval"),
 		LudicrousMode:     Zero.Conf.GetBool("ludicrous_mode"),
+		totalCache:        int64(Zero.Conf.GetInt("cache_mb")),
+		cachePercentage:   Zero.Conf.GetString("cache_percentage"),
 	}
 	glog.Infof("Setting Config to: %+v", opts)
 
@@ -215,6 +225,13 @@ func run() {
 			opts.rebalanceInterval)
 	}
 
+	x.AssertTruef(opts.totalCache >= 0, "ERROR: Cache size must be non-negative")
+
+	glog.Info(opts.cachePercentage)
+	cachePercent := x.GetCachePercentages(opts.cachePercentage, "blockCache:indexCache")
+	blockCache := (cachePercent[0] * (opts.totalCache << 20)) / 100
+	indexCache := (cachePercent[1] * (opts.totalCache << 20)) / 100
+
 	grpc.EnableTracing = false
 	otrace.ApplyConfig(otrace.Config{
 		DefaultSampler: otrace.ProbabilitySampler(Zero.Conf.GetFloat64("trace"))})
@@ -235,7 +252,7 @@ func run() {
 	// Open raft write-ahead log and initialize raft node.
 	x.Checkf(os.MkdirAll(opts.w, 0700), "Error while creating WAL dir.")
 	kvOpt := badger.LSMOnlyOptions(opts.w).WithSyncWrites(false).WithTruncate(true).
-		WithValueLogFileSize(64 << 20).WithMaxCacheSize(10 << 20).WithLoadBloomsOnOpen(false)
+		WithValueLogFileSize(64 << 20).WithBlockCacheSize(blockCache).WithIndexCacheSize(indexCache).WithLoadBloomsOnOpen(false)
 
 	kvOpt.ZSTDCompressionLevel = 3
 

go.mod

@@ -17,10 +17,11 @@ require (
 	github.com/blevesearch/segment v0.0.0-20160915185041-762005e7a34f // indirect
 	github.com/blevesearch/snowballstem v0.0.0-20180110192139-26b06a2c243d // indirect
 	github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd
-	github.com/dgraph-io/badger/v2 v2.2007.1
+	github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200825103817-c3b15824dc14
+	github.com/dgraph-io/dgo v1.0.0
 	github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6
 	github.com/dgraph-io/graphql-transport-ws v0.0.0-20200715131837-c0460019ead2
-	github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de
+	github.com/dgraph-io/ristretto v0.0.4-0.20200820164438-623d8ef1614b
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/dgrijalva/jwt-go/v4 v4.0.0-preview1
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13
@@ -31,23 +32,23 @@ require (
 	github.com/gogo/protobuf v1.3.1
 	github.com/golang/geo v0.0.0-20170810003146-31fb0106dc4a
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
-	github.com/golang/protobuf v1.3.5
+	github.com/golang/protobuf v1.4.2
 	github.com/golang/snappy v0.0.1
 	github.com/google/codesearch v1.0.0
-	github.com/google/go-cmp v0.4.0
+	github.com/google/go-cmp v0.5.0
 	github.com/google/uuid v1.0.0
 	github.com/gorilla/websocket v1.4.2
 	github.com/graph-gophers/graphql-go v0.0.0-20200309224638-dae41bde9ef9
 	github.com/graph-gophers/graphql-transport-ws v0.0.0-20190611222414-40c048432299 // indirect
 	github.com/hashicorp/vault/api v1.0.4
+	github.com/kr/pretty v0.2.0 // indirect
 	github.com/minio/minio-go/v6 v6.0.55
 	github.com/mitchellh/panicwrap v1.0.0
 	github.com/paulmach/go.geojson v0.0.0-20170327170536-40612a87147b
 	github.com/philhofer/fwd v1.0.0 // indirect
 	github.com/pkg/errors v0.9.1
 	github.com/pkg/profile v1.2.1
 	github.com/prometheus/client_golang v0.9.3-0.20190127221311-3c4408c8b829
-	github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90 // indirect
 	github.com/prometheus/common v0.4.1 // indirect
 	github.com/prometheus/procfs v0.0.0-20190517135640-51af30a78b0e // indirect
 	github.com/spf13/cast v1.3.0
@@ -61,13 +62,13 @@ require (
 	github.com/willf/bitset v0.0.0-20181014161241-71fa2377963f // indirect
 	go.etcd.io/etcd v0.0.0-20190228193606-a943ad0ee4c9
 	go.opencensus.io v0.21.0
-	golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550
-	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e
+	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
+	golang.org/x/net v0.0.0-20200822124328-c89045814202
 	golang.org/x/sync v0.0.0-20190423024810-112230192c58
-	golang.org/x/sys v0.0.0-20200515095857-1151b9dac4a9
+	golang.org/x/sys v0.0.0-20200824131525-c12d262b63d8
 	golang.org/x/text v0.3.2
-	google.golang.org/genproto v0.0.0-20190516172635-bb713bdc0e52 // indirect
-	google.golang.org/grpc v1.23.0
+	google.golang.org/grpc v1.27.0
+	google.golang.org/protobuf v1.25.0 // indirect
 	gopkg.in/DataDog/dd-trace-go.v1 v1.13.1 // indirect
 	gopkg.in/ini.v1 v1.48.0 // indirect
 	gopkg.in/yaml.v2 v2.2.4