feat: add --cache_mb and --cache_percentage flag (#6271)

Badger now has 2 separate caches blockCache and indexCache 
(see hypermodeinc/badger#1476)
This PR adds --cache_mb and --cache_percentage flags for alpha and zero.
The total cache is split among various caches used by zero and alpha based on percentages defined.
Cache size is in MBs and format of caches is as follows
For alpha:
PostingListCache,PstoreBlockCache,PstoreIndexCache,WstoreBlockCache,WstoreIndexCache
For zero:
blockCache,indexCache

dgraph/cmd/alpha/run.go

@@ -203,6 +203,12 @@ 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.")
+
+	// Cache flags
+	flag.Int64("cache_mb", 0, "Total size of cache (in MB) to be used in alpha.")
+	flag.String("cache_percentage", "0,65,25,0,10",
+		`Cache percentages summing up to 100 for various caches (FORMAT:
+		PostingListCache,PstoreBlockCache,PstoreIndexCache,WstoreBlockCache,WstoreIndexCache).`)
 }
 
 func setupCustomTokenizers() {
@@ -577,10 +583,26 @@ 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, err := x.GetCachePercentages(cachePercentage, 5)
+	x.Check(err)
+	postingListCacheSize := (cachePercent[0] * (totalCache << 20)) / 100
+	pstoreBlockCacheSize := (cachePercent[1] * (totalCache << 20)) / 100
+	pstoreIndexCacheSize := (cachePercent[2] * (totalCache << 20)) / 100
+	wstoreBlockCacheSize := (cachePercent[3] * (totalCache << 20)) / 100
+	wstoreIndexCacheSize := (cachePercent[4] * (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,
+		WBlockCacheSize:        wstoreBlockCacheSize,
+		WIndexCacheSize:        wstoreIndexCacheSize,
 
 		MutationsMode:  worker.AllowMutations,
 		AuthToken:      Alpha.Conf.GetString("auth_token"),
@@ -701,7 +723,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,7 +18,6 @@ package zero
 
 import (
 	"context"
-	//	"errors"
 	"fmt"
 	"log"
 	"net"
@@ -55,6 +54,9 @@ type options struct {
 	w                 string
 	rebalanceInterval time.Duration
 	LudicrousMode     bool
+
+	totalCache      int64
+	cachePercentage string
 }
 
 var opts options
@@ -101,6 +103,11 @@ 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.")
+
+	// 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 summing up to 100 for various caches (FORMAT: blockCache,indexCache).")
 }
 
 func setupListener(addr string, port int, kind string) (listener net.Listener, err error) {
@@ -183,6 +190,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)
 
@@ -232,10 +241,22 @@ func run() {
 	httpListener, err := setupListener(addr, x.PortZeroHTTP+opts.portOffset, "http")
 	x.Check(err)
 
+	x.AssertTruef(opts.totalCache >= 0, "ERROR: Cache size must be non-negative")
+
+	cachePercent, err := x.GetCachePercentages(opts.cachePercentage, 2)
+	x.Check(err)
+	blockCacheSz := (cachePercent[0] * (opts.totalCache << 20)) / 100
+	indexCacheSz := (cachePercent[1] * (opts.totalCache << 20)) / 100
+
 	// 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)
+	kvOpt := badger.LSMOnlyOptions(opts.w).
+		WithSyncWrites(false).
+		WithTruncate(true).
+		WithValueLogFileSize(64 << 20).
+		WithBlockCacheSize(blockCacheSz).
+		WithIndexCacheSize(indexCacheSz).
+		WithLoadBloomsOnOpen(false)
 
 	kvOpt.ZSTDCompressionLevel = 3
 

go.mod

@@ -17,7 +17,7 @@ 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.2007.2-0.20200826122734-bc243f38bfe1
 	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

go.sum

@@ -6,6 +6,7 @@ contrib.go.opencensus.io/exporter/prometheus v0.1.0 h1:SByaIoWwNgMdPSgl5sMqM2KDE
 contrib.go.opencensus.io/exporter/prometheus v0.1.0/go.mod h1:cGFniUXGZlKRjzOyuZJ6mgB+PgBcCIa79kEKR8YCW+A=
 github.com/99designs/gqlgen v0.12.2 h1:aOdpsiCycFtCnAv8CAI1exnKrIDHMqtMzQoXeTziY4o=
 github.com/99designs/gqlgen v0.12.2/go.mod h1:7zdGo6ry9u1YBp/qlb2uxSU5Mt2jQKLcBETQiKk+Bxo=
+github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9 h1:HD8gA2tkByhMAwYaFAX9w2l7vxvBQ5NMoxDrkhqhtn4=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
@@ -83,8 +84,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/badger v1.6.0 h1:DshxFxZWXUcO0xX476VJC07Xsr6ZCBVRHKZ93Oh7Evo=
 github.com/dgraph-io/badger v1.6.0/go.mod h1:zwt7syl517jmP8s94KqSxTlM6IMsdhYy6psNgSztDR4=
-github.com/dgraph-io/badger/v2 v2.2007.1 h1:t36VcBCpo4SsmAD5M8wVv1ieVzcALyGfaJ92z4ccULM=
-github.com/dgraph-io/badger/v2 v2.2007.1/go.mod h1:26P/7fbL4kUZVEVKLAKXkBXKOydDmM2p1e+NhhnBCAE=
+github.com/dgraph-io/badger/v2 v2.2007.2-0.20200826122734-bc243f38bfe1 h1:vPPlQYByX3+Z3NOaB06i7VCb0bNOznVQ9eEnqLxbrH0=
+github.com/dgraph-io/badger/v2 v2.2007.2-0.20200826122734-bc243f38bfe1/go.mod h1:26P/7fbL4kUZVEVKLAKXkBXKOydDmM2p1e+NhhnBCAE=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6 h1:toHzMCdCUgYsjM0cW9+wafnKFXfp1HizIJUyzihN+vk=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6/go.mod h1:rHa+h3kI4M8ASOirxyIyNeXBfHFgeskVUum2OrDMN3U=
 github.com/dgraph-io/graphql-transport-ws v0.0.0-20200715131837-c0460019ead2 h1:NSl3XXyON9bgmBJSAvr5FPrgILAovtoTs7FwdtaZZq0=

posting/list_test.go

@@ -1346,7 +1346,8 @@ func TestMain(m *testing.M) {
 
 	ps, err = badger.OpenManaged(badger.DefaultOptions(dir))
 	x.Check(err)
-	Init(ps)
+	// Not using posting list cache
+	Init(ps, 0)
 	schema.Init(ps)
 
 	r := m.Run()

posting/lists.go

@@ -136,21 +136,21 @@ var (
 )
 
 // Init initializes the posting lists package, the in memory and dirty list hash.
-func Init(ps *badger.DB) {
+func Init(ps *badger.DB, cacheSize int64) {
 	pstore = ps
 	closer = y.NewCloser(1)
 	go updateMemoryMetrics(closer)
 
 	// Initialize cache.
-	// TODO(Ibrahim): Add flag to switch cache on and off. For now cache is disabled.
-	if true {
+	if cacheSize == 0 {
 		return
 	}
-	// TODO(Ibrahim): Replace hard-coded value with value from flag.
+
 	var err error
 	lCache, err = ristretto.NewCache(&ristretto.Config{
-		NumCounters: 200e6,
-		MaxCost:     int64(1000 * 1024 * 1024),
+		// Use 5% of cache memory for storing counters.
+		NumCounters: int64(float64(cacheSize) * 0.05 * 2),
+		MaxCost:     int64(float64(cacheSize) * 0.95),
 		BufferItems: 64,
 		Metrics:     true,
 		Cost: func(val interface{}) int64 {

worker/config.go

@@ -61,6 +61,15 @@ type Options struct {
 	// AllottedMemory is the estimated size taken by the LRU cache.
 	AllottedMemory float64
 
+	// PBlockCacheSize is the size of block cache for pstore
+	PBlockCacheSize int64
+	// PIndexCacheSize is the size of index cache for pstore
+	PIndexCacheSize int64
+	// WBlockCacheSize is the size of block cache for wstore
+	WBlockCacheSize int64
+	// WIndexCacheSize is the size of index cache for wstore
+	WIndexCacheSize int64
+
 	// HmacSecret stores the secret used to sign JSON Web Tokens (JWT).
 	HmacSecret x.SensitiveByteSlice
 	// AccessJwtTtl is the TTL for the access JWT.

worker/server_state.go

@@ -142,7 +142,8 @@ func (s *ServerState) initStorage() {
 		opt := badger.LSMOnlyOptions(Config.WALDir)
 		opt = setBadgerOptions(opt, true)
 		opt.ValueLogMaxEntries = 10000 // Allow for easy space reclamation.
-		opt.MaxCacheSize = 10 << 20    // 10 mb of cache size for WAL.
+		opt.BlockCacheSize = Config.WBlockCacheSize
+		opt.IndexCacheSize = Config.WIndexCacheSize
 
 		// Print the options w/o exposing key.
 		// TODO: Build a stringify interface in Badger options, which is used to print nicely here.
@@ -162,10 +163,8 @@ func (s *ServerState) initStorage() {
 		opt := badger.DefaultOptions(Config.PostingDir).
 			WithValueThreshold(1 << 10 /* 1KB */).
 			WithNumVersionsToKeep(math.MaxInt32).
-			WithMaxCacheSize(1 << 30).
-			WithKeepBlockIndicesInCache(true).
-			WithKeepBlocksInCache(true).
-			WithMaxBfCacheSize(500 << 20) // 500 MB of bloom filter cache.
+			WithBlockCacheSize(Config.PBlockCacheSize).
+			WithIndexCacheSize(Config.PIndexCacheSize)
 		opt = setBadgerOptions(opt, false)
 
 		// Print the options w/o exposing key.

worker/worker_test.go

@@ -386,7 +386,8 @@ func TestMain(m *testing.M) {
 	ps, err := badger.OpenManaged(opt)
 	x.Check(err)
 	pstore = ps
-	posting.Init(ps)
+	// Not using posting list cache
+	posting.Init(ps, 0)
 	Init(ps)
 
 	os.Exit(m.Run())

x/x.go

@@ -1099,3 +1099,35 @@ func DeepCopyJsonArray(a []interface{}) []interface{} {
 	}
 	return aCopy
 }
+
+// GetCachePercentages returns the slice of cache percentages given the "," (comma) separated
+// cache percentages(integers) string and expected number of caches.
+func GetCachePercentages(cpString string, numExpected int) ([]int64, error) {
+	cp := strings.Split(cpString, ",")
+	// Sanity checks
+	if len(cp) != numExpected {
+		return nil, errors.Errorf("ERROR: expected %d cache percentages, got %d",
+			numExpected, len(cp))
+	}
+
+	var cachePercent []int64
+	percentSum := 0
+	for _, percent := range cp {
+		x, err := strconv.Atoi(percent)
+		if err != nil {
+			return nil, errors.Errorf("ERROR: unable to parse cache percentage(%s)", percent)
+		}
+		if x < 0 {
+			return nil, errors.Errorf("ERROR: cache percentage(%s) cannot be negative", percent)
+		}
+		cachePercent = append(cachePercent, int64(x))
+		percentSum += x
+	}
+
+	if percentSum != 100 {
+		return nil, errors.Errorf("ERROR: cache percentages (%s) does not sum up to 100",
+			strings.Join(cp, "+"))
+	}
+
+	return cachePercent, nil
+}