[Breaking flag] fix(Backups): new badger SuperFlag, NumGoroutines option solves OOM crashes

dgraph/cmd/alpha/run.go

@@ -108,11 +108,12 @@ they form a Raft group and provide synchronous replication.
 	flag.StringP("postings", "p", "p", "Directory to store posting lists.")
 	flag.String("tmp", "t", "Directory to store temporary buffers.")
 
-	// Options around how to set up Badger.
-	flag.String("badger.compression", "snappy",
-		"[none, zstd:level, snappy] Specifies the compression algorithm and the compression"+
-			"level (if applicable) for the postings directory. none would disable compression,"+
-			" while zstd:1 would set zstd compression at level 1.")
+	flag.String("badger", worker.BadgerDefaults,
+		`Various badger options.
+	goroutines=N provides the number of goroutines to use in badger.Stream.
+	compression=[none, zstd:level, snappy] specifies the compression algorithm and the compression 
+		level (if applicable) for the postings directory. "none" would disable compression, while 
+		"zstd:1" would set zstd compression at level 1.`)
 	enc.RegisterFlags(flag)
 
 	// Snapshot and Transactions.
@@ -605,7 +606,8 @@ func run() {
 	pstoreIndexCacheSize := (cachePercent[2] * (totalCache << 20)) / 100
 	walCache := (cachePercent[3] * (totalCache << 20)) / 100
 
-	ctype, clevel := x.ParseCompression(Alpha.Conf.GetString("badger.compression"))
+	badger := x.NewSuperFlag(Alpha.Conf.GetString("badger")).MergeAndCheckDefault(worker.BadgerDefaults)
+	ctype, clevel := x.ParseCompression(badger.GetString("compression"))
 
 	conf := audit.GetAuditConf(Alpha.Conf.GetString("audit"))
 	opts := worker.Options{
@@ -672,6 +674,7 @@ func run() {
 		NumPendingProposals:  Alpha.Conf.GetInt("pending_proposals"),
 		ZeroAddr:             strings.Split(Alpha.Conf.GetString("zero"), ","),
 		Raft:                 raft,
+		Badger:               badger,
 		WhiteListedIPRanges:  ips,
 		MaxRetries:           Alpha.Conf.GetInt("max_retries"),
 		StrictMutations:      opts.MutationsMode == worker.StrictMutations,

dgraph/cmd/bulk/run.go

@@ -115,21 +115,22 @@ func init() {
 		"Ignore UIDs in load files and assign new ones.")
 
 	// Options around how to set up Badger.
-	flag.String("badger.compression", "snappy",
-		"[none, zstd:level, snappy] Specifies the compression algorithm and the compression"+
-			"level (if applicable) for the postings directory. none would disable compression,"+
-			" while zstd:1 would set zstd compression at level 1.")
-	flag.Int64("badger.cache_mb", 64, "Total size of cache (in MB) per shard in reducer.")
-	flag.String("badger.cache_percentage", "70,30",
-		"Cache percentages summing up to 100 for various caches"+
-			" (FORMAT: BlockCacheSize, IndexCacheSize).")
+	flag.String("badger", worker.BadgerDefaults,
+		`Various badger options.
+	compression=[none, zstd:level, snappy] specifies the compression algorithm and the compression
+		level (if applicable) for the postings directory. "none" would disable compression, while
+		"zstd:1" would set zstd compression at level 1.
+	cache_mb=N total size of cache (in MB) per shard in the reducer.
+	cache_percentage=N cache percentages summing up to 100 for various caches.
+		(FORMAT: BlockCacheSize, IndexCacheSize)`)
 	x.RegisterClientTLSFlags(flag)
 	// Encryption and Vault options
 	enc.RegisterFlags(flag)
 }
 
 func run() {
-	ctype, clevel := x.ParseCompression(Bulk.Conf.GetString("badger.compression"))
+	badger := x.NewSuperFlag(Bulk.Conf.GetString("badger")).MergeAndCheckDefault(worker.BadgerDefaults)
+	ctype, clevel := x.ParseCompression(badger.GetString("compression"))
 	opt := options{
 		DataFiles:        Bulk.Conf.GetString("files"),
 		DataFormat:       Bulk.Conf.GetString("format"),
@@ -166,9 +167,9 @@ func run() {
 		os.Exit(0)
 	}
 
-	totalCache := int64(Bulk.Conf.GetInt("badger.cache_mb"))
+	totalCache := int64(badger.GetUint64("cache_mb"))
 	x.AssertTruef(totalCache >= 0, "ERROR: Cache size must be non-negative")
-	cachePercent, err := x.GetCachePercentages(Bulk.Conf.GetString("badger.cache_percentage"), 2)
+	cachePercent, err := x.GetCachePercentages(badger.GetString("cache_percentage"), 2)
 	x.Check(err)
 	totalCache <<= 20 // Convert to MB.
 	opt.BlockCacheSize = (cachePercent[0] * totalCache) / 100

worker/backup_processor.go

@@ -37,11 +37,6 @@ import (
 	"github.com/dgraph-io/dgraph/x"
 )
 
-const (
-	// backupNumGo is the number of go routines used by the backup stream writer.
-	backupNumGo = 16
-)
-
 // BackupProcessor handles the different stages of the backup process.
 type BackupProcessor struct {
 	// DB is the Badger pstore managed by this node.
@@ -65,7 +60,7 @@ func NewBackupProcessor(db *badger.DB, req *pb.BackupRequest) *BackupProcessor {
 	bp := &BackupProcessor{
 		DB:      db,
 		Request: req,
-		threads: make([]*threadLocal, backupNumGo),
+		threads: make([]*threadLocal, x.WorkerConfig.Badger.GetUint64("goroutines")),
 	}
 	for i := range bp.threads {
 		bp.threads[i] = &threadLocal{
@@ -128,7 +123,6 @@ func (pr *BackupProcessor) WriteBackup(ctx context.Context) (*pb.BackupResponse,
 
 	stream := pr.DB.NewStreamAt(pr.Request.ReadTs)
 	stream.LogPrefix = "Dgraph.Backup"
-	stream.NumGo = backupNumGo
 
 	stream.KeyToList = func(key []byte, itr *badger.Iterator) (*bpb.KVList, error) {
 		tl := pr.threads[itr.ThreadId]

worker/server_state.go

@@ -30,6 +30,10 @@ import (
 	"github.com/golang/glog"
 )
 
+const (
+	BadgerDefaults = "goroutines=8; compression=snappy; cache_mb=64; cache_percentage=70,30"
+)
+
 // ServerState holds the state of the Dgraph server.
 type ServerState struct {
 	FinishCh chan struct{} // channel to wait for all pending reqs to finish.
@@ -109,6 +113,7 @@ func (s *ServerState) initStorage() {
 		opt := badger.DefaultOptions(Config.PostingDir).
 			WithValueThreshold(1 << 10 /* 1KB */).
 			WithNumVersionsToKeep(math.MaxInt32).
+			WithNumGoroutines(int(x.WorkerConfig.Badger.GetUint64("goroutines"))).
 			WithBlockCacheSize(Config.PBlockCacheSize).
 			WithIndexCacheSize(Config.PIndexCacheSize)
 		opt = setBadgerOptions(opt)

x/config.go

@@ -79,6 +79,8 @@ type WorkerOptions struct {
 	TLSServerConfig *tls.Config
 	// Raft stores options related to Raft.
 	Raft *SuperFlag
+	// Badger stores options related to Badger.
+	Badger *SuperFlag
 	// WhiteListedIPRanges is a list of IP ranges from which requests will be allowed.
 	WhiteListedIPRanges []IPRange
 	// MaxRetries is the maximum number of times to retry a commit before giving up.