Optimize XidtoUID map used by live and bulk loader (#2998)

I've spent the last few days looking at how to optimize the live mutation path in Dgraph server. While trying many things in the server (past commits included), I realized my optimizations in the server are not improving things much, the throughput saturating at 20-30K NQuads/sec.

Turns out, it was the live loader which was causing the saturation. The XID to UID assigner was the bottleneck causing the throughput to stagnate, despite the server being underutilized.

This PR fixes that by optimizing the assigner. In particular, I've removed the slow LRU cache. Added buffer to `newRanges` channel to ensure we always have a range handy when we run out. Made passing badger DB instance optional, so we can avoid doing disk writes if not required. And made other optimizations around how we lock, etc. I also added benchmarks for the assigner, which shows each allocation (tested via parallel benchmark) takes 350 ns/op on my desktop.

With these changes, the live loader throughput jumps to 100K-120K NQuads/sec on my desktop. In particular, pre-assigning UIDs to the RDF/JSON file yields maximum throughput. I can load 140M friend graph RDFs in 25 mins.

Helps with #2975 .

Changes:
* Work on optimizing XidToUid map.
* Add the test and benchmark for xid to uid map
* Working code with decreased memory usage. Includes a new BumpUp API.
* Working live loader, which can optionally just keep all the mapping in memory.
* Adding shards back to XidMap speed up operations by a huge factor. Benchmark shows each allocation is 300ns.
* Make BumpTo much faster by calling Zero directly, instead of looping through the newRanges channel.
* Improve how BumpTo() happens by using a maxSeenUid variable.

Author: manishrjain

dgraph/cmd/bulk/loader.go

@@ -147,6 +147,7 @@ func readSchema(filename string) []*pb.SchemaUpdate {
 func (ld *loader) mapStage() {
 	ld.prog.setPhase(mapPhase)
 
+	// TODO: Consider if we need to always store the XIDs in Badger. Things slow down if we do.
 	xidDir := filepath.Join(ld.opt.TmpDir, "xids")
 	x.Check(os.Mkdir(xidDir, 0755))
 	opt := badger.DefaultOptions
@@ -157,10 +158,7 @@ func (ld *loader) mapStage() {
 	var err error
 	ld.xidDB, err = badger.Open(opt)
 	x.Check(err)
-	ld.xids = xidmap.New(ld.xidDB, ld.zero, xidmap.Options{
-		NumShards: 1 << 10,
-		LRUSize:   1 << 19,
-	})
+	ld.xids = xidmap.New(ld.zero, ld.xidDB)
 
 	var dir, ext string
 	var loaderType int
@@ -226,7 +224,7 @@ func (ld *loader) mapStage() {
 	for i := range ld.mappers {
 		ld.mappers[i] = nil
 	}
-	ld.xids.EvictAll()
+	x.Check(ld.xids.Flush())
 	x.Check(ld.xidDB.Close())
 	ld.xids = nil
 	runtime.GC()

dgraph/cmd/bulk/mapper.go

@@ -217,8 +217,8 @@ func (m *mapper) processNQuad(nq gql.NQuad) {
 }
 
 func (m *mapper) lookupUid(xid string) uint64 {
-	uid, isNew := m.xids.AssignUid(xid)
-	if !isNew || !m.opt.StoreXids {
+	uid := m.xids.AssignUid(xid)
+	if !m.opt.StoreXids {
 		return uid
 	}
 	if strings.HasPrefix(xid, "_:") {

dgraph/cmd/live/batch.go

@@ -63,7 +63,7 @@ type loader struct {
 	dc         *dgo.Dgraph
 	alloc      *xidmap.XidMap
 	ticker     *time.Ticker
-	kv         *badger.DB
+	db         *badger.DB
 	requestsWg sync.WaitGroup
 	// If we retry a request, we add one to retryRequestsWg.
 	retryRequestsWg sync.WaitGroup
@@ -188,15 +188,18 @@ func (l *loader) makeRequests() {
 }
 
 func (l *loader) printCounters() {
-	l.ticker = time.NewTicker(2 * time.Second)
+	period := 5 * time.Second
+	l.ticker = time.NewTicker(period)
 	start := time.Now()
 
+	var last Counter
 	for range l.ticker.C {
 		counter := l.Counter()
-		rate := float64(counter.Nquads) / counter.Elapsed.Seconds()
+		rate := float64(counter.Nquads-last.Nquads) / period.Seconds()
 		elapsed := time.Since(start).Round(time.Second)
-		fmt.Printf("[%6s] Txns: %d N-Quads: %d N-Quads/sec: %5.0f Aborts: %d\n",
+		fmt.Printf("[%6s] Txns: %d N-Quads: %d N-Quads/s [last 5s]: %5.0f Aborts: %d\n",
 			elapsed, counter.TxnsDone, counter.Nquads, rate, counter.Aborts)
+		last = counter
 	}
 }
 

dgraph/cmd/live/run.go

@@ -158,13 +158,12 @@ func (l *loader) uid(val string) string {
 	// to be an existing node in the graph. There is limited protection against
 	// a user selecting an unassigned UID in this way - it may be assigned
 	// later to another node. It is up to the user to avoid this.
-	if strings.HasPrefix(val, "0x") {
-		if _, err := strconv.ParseUint(val[2:], 16, 64); err == nil {
-			return val
-		}
+	if uid, err := strconv.ParseUint(val, 0, 64); err == nil {
+		l.alloc.BumpTo(uid)
+		return fmt.Sprintf("%#x", uid)
 	}
 
-	uid, _ := l.alloc.AssignUid(val)
+	uid := l.alloc.AssignUid(val)
 	return fmt.Sprintf("%#x", uint64(uid))
 }
 
@@ -241,36 +240,32 @@ func (l *loader) processLoadFile(ctx context.Context, rd *bufio.Reader, ck chunk
 }
 
 func setup(opts batchMutationOptions, dc *dgo.Dgraph) *loader {
-	x.Check(os.MkdirAll(opt.clientDir, 0700))
-	o := badger.DefaultOptions
-	o.SyncWrites = true // So that checkpoints are persisted immediately.
-	o.TableLoadingMode = bopt.MemoryMap
-	o.Dir = opt.clientDir
-	o.ValueDir = opt.clientDir
+	var db *badger.DB
+	if len(opt.clientDir) > 0 {
+		x.Check(os.MkdirAll(opt.clientDir, 0700))
+		o := badger.DefaultOptions
+		o.Dir = opt.clientDir
+		o.ValueDir = opt.clientDir
+		o.TableLoadingMode = bopt.MemoryMap
+		o.SyncWrites = false
 
-	kv, err := badger.Open(o)
-	x.Checkf(err, "Error while creating badger KV posting store")
+		var err error
+		db, err = badger.Open(o)
+		x.Checkf(err, "Error while creating badger KV posting store")
+	}
 
 	// compression with zero server actually makes things worse
 	connzero, err := x.SetupConnection(opt.zero, &tlsConf, false)
 	x.Checkf(err, "Unable to connect to zero, Is it running at %s?", opt.zero)
 
-	alloc := xidmap.New(
-		kv,
-		connzero,
-		xidmap.Options{
-			NumShards: 100,
-			LRUSize:   1e5,
-		},
-	)
-
+	alloc := xidmap.New(connzero, db)
 	l := &loader{
 		opts:     opts,
 		dc:       dc,
 		start:    time.Now(),
 		reqs:     make(chan api.Mutation, opts.Pending*2),
 		alloc:    alloc,
-		kv:       kv,
+		db:       db,
 		zeroconn: connzero,
 	}
 
@@ -322,17 +317,8 @@ func run() error {
 	}
 	dgraphClient := dgo.NewDgraphClient(clients...)
 
-	if len(opt.clientDir) == 0 {
-		var err error
-		opt.clientDir, err = ioutil.TempDir("", "x")
-		x.Checkf(err, "Error while trying to create temporary client directory.")
-		fmt.Printf("Creating temp client directory at %s\n", opt.clientDir)
-		defer os.RemoveAll(opt.clientDir)
-	}
 	l := setup(bmOpts, dgraphClient)
 	defer l.zeroconn.Close()
-	defer l.kv.Close()
-	defer l.alloc.EvictAll()
 
 	if len(opt.schemaFile) > 0 {
 		if err := processSchemaFile(ctx, opt.schemaFile, dgraphClient); err != nil {
@@ -397,5 +383,9 @@ func run() error {
 	fmt.Printf("Time spent                   : %v\n", c.Elapsed)
 	fmt.Printf("N-Quads processed per second : %d\n", rate)
 
+	if l.db != nil {
+		l.alloc.Flush()
+		l.db.Close()
+	}
 	return nil
 }

worker/mutation.go

@@ -290,7 +290,7 @@ func ValidateAndConvert(edge *pb.DirectedEdge, su *pb.SchemaUpdate) error {
 		return x.Errorf("Input for predicate %s of type uid is scalar", edge.Attr)
 
 	case schemaType.IsScalar() && !storageType.IsScalar():
-		return x.Errorf("Input for predicate %s of type scalar is uid", edge.Attr)
+		return x.Errorf("Input for predicate %s of type scalar is uid. Edge: %v", edge.Attr, edge)
 
 	// The suggested storage type matches the schema, OK!
 	case storageType == schemaType && schemaType != types.DefaultID:

x/x.go

@@ -433,7 +433,7 @@ func SetupConnection(host string, tlsConf *TLSHelperConfig, useGz bool) (*grpc.C
 		grpc.WithBlock(),
 		grpc.WithTimeout(10*time.Second))
 
-	if tlsConf.CertRequired {
+	if tlsConf != nil && tlsConf.CertRequired {
 		tlsConf.ConfigType = TLSClientConfig
 		tlsCfg, _, err := GenerateTLSConfig(*tlsConf)
 		if err != nil {