perf: Various optimizations to the bulk loader (#6412)

With this change, bulk loader runs fine over very large datasets involving billions of edges, without running out of memory on a 64 GB RAM machine. Tested with loading a client dataset and entire Stack Overflow dataset.

Changes:
- Modifies codec.Encoder to use memory allocated via z.Calloc. This makes UidPack use manual memory, and adds corresponding FreePack calls over the codebase.
- pb.MapEntry is replaced with a simpler struct, which is allocated via z.Buffer.
- countIndexEntry is similarly replaced with a simpler struct, allocated via z.Buffer.
- Keys are partitioned every partition_mb size.
- Map size is set by default to 2GB. This after compression becomes around 300-400MB.
- Instead of allocated Slices, sorting is done via sorting offsets, and then accessing entries in sorted offset order (randomly instead of serially over) directly in z.Buffer.
- Prints detailed stats from jemalloc every minute.
- Prints memory usage from both Go and Jemalloc.
- Lists generated via toList are not accumulated till writing anymore. Instead, they're pushed to a buffered channel, to better control memory usage.
- If reducer is processing encoding requests over 1GB, then it throttles other requests.

Commits:
* Do not use all CPUs for encoding. Only use as many as specified in flags.
* Allocate UidBlocks via Calloc.
* Print out memory allocation as well
* Remove prevoffset logic
* Allocate mapEntries via Calloc too. Try to release memory allocated for Packs.
* Fix a bug in posting list split code where the main pack keeps the blocks allocated to the parts.
* Try without doing mapEntry allocations manually.
* Report memory stats every minute
* Add a throttle to avoid sending out too much data for encoding at once.
* Compiling version of removal of pb.MapEntry buffer.
* Make bulk loader work with the new MapEntry system.
* Don't stop encoding. Only pause what is being sent out for encoding
* Switch to z.Buffer and use z.CallocNoRef
* Avoid doing encoding again, if we already have a clean posting list. Also avoid creating arrays for postings during split.
* Generate a histogram of buffers.
* Do not deallocate Pack twice.
* Encode 4GB and over requests.
* Generate a histogram of key to count.
* Do not accumulate all the UIDs in memory.
* Free packs on encoders used outside the bulk loader.
* Print out the keys for giant buffer
* Add histogram of entries per key. Avoid duplicate partition keys.
* Use channels for kvs and splits.
* Print memstats in Mapper as well.
* use memhash instead of creating strings
* Count indexing is memory efficient now.
* changed partion to take 4 mb
* Sort the offsets before encoding.
* Fixed up the count indexing stuff. Add a new flag partition_mb to allow users to specify how often to pick partition keys.

Co-authored-by: பாலாஜி <balaji@dgraph.io>
Co-authored-by: Martin Martinez Rivera <martinmr@dgraph.io>
Co-authored-by: Ibrahim Jarif <ibrahim@dgraph.io>
Co-authored-by: harshil <harshil@dgraph.io>

Author: 5 people

algo/uidlist_test.go

@@ -330,6 +330,8 @@ func BenchmarkListIntersectRandom(b *testing.B) {
 		if j < len(dst2.Uids) {
 			b.Errorf("Unexpected error in intersection")
 		}
+
+		codec.FreePack(compressedUids)
 	}
 
 	randomTests(10240, 0.3)
@@ -395,6 +397,8 @@ func BenchmarkListIntersectRatio(b *testing.B) {
 			if j < len(dst2.Uids) {
 				b.Errorf("Unexpected error in intersection")
 			}
+
+			codec.FreePack(compressedUids)
 		}
 	}
 
@@ -464,6 +468,7 @@ func TestIntersectCompressedWithLinJump(t *testing.T) {
 			actual := make([]uint64, 0)
 			IntersectCompressedWithLinJump(&dec, otherNums, &actual)
 			require.Equal(t, commonNums, actual)
+			codec.FreePack(pack)
 		}
 	}
 }
@@ -488,6 +493,7 @@ func TestIntersectCompressedWithBin(t *testing.T) {
 			actual := make([]uint64, 0)
 			IntersectCompressedWithBin(&dec, otherNums, &actual)
 			require.Equal(t, commonNums, actual)
+			codec.FreePack(pack)
 		}
 	}
 }
@@ -513,6 +519,7 @@ func TestIntersectCompressedWithBinMissingSize(t *testing.T) {
 			actual := make([]uint64, 0)
 			IntersectCompressedWithBin(&dec, otherNums, &actual)
 			require.Equal(t, commonNums, actual)
+			codec.FreePack(pack)
 		}
 	}
 }

codec/benchmark/benchmark.go

@@ -121,7 +121,8 @@ func benchmarkPack(trials int, chunks *chunks) int {
 	for i := range times {
 		start := time.Now()
 		for _, c := range chunks.data {
-			codec.Encode(c, 256)
+			pack := codec.Encode(c, 256)
+			codec.FreePack(pack)
 			// bp128.DeltaPack(c)
 		}
 		times[i] = int(time.Since(start).Nanoseconds())

codec/codec.go

@@ -20,9 +20,11 @@ import (
 	"bytes"
 	"math"
 	"sort"
+	"unsafe"
 
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/x"
+	"github.com/dgraph-io/ristretto/z"
 	"github.com/dgryski/go-groupvarint"
 )
 
@@ -46,15 +48,55 @@ type Encoder struct {
 	uids      []uint64
 }
 
+var blockSize = int(unsafe.Sizeof(pb.UidBlock{}))
+
+// AllocateBlock would allocate a block via Calloc. This block must be released via a call to
+// ReleaseBlock.
+func AllocateBlock() *pb.UidBlock {
+	// Allocate blocks manually.
+	// TODO: Avoid calling z.Calloc repeatedly. Instead use the allocator. The tricky question here
+	// is how would FreePack call be able to trace the allocator.
+	b := z.CallocNoRef(blockSize)
+	if len(b) == 0 {
+		return &pb.UidBlock{}
+	}
+	return (*pb.UidBlock)(unsafe.Pointer(&b[0]))
+}
+
+// FreeBlock releases a previously manually allocated UidBlock.
+func FreeBlock(ub *pb.UidBlock) {
+	buf := (*[z.MaxArrayLen]byte)(unsafe.Pointer(ub))[:blockSize:blockSize]
+	z.Free(buf)
+}
+
+func FreePack(pack *pb.UidPack) {
+	if pack == nil {
+		return
+	}
+	for _, b := range pack.Blocks {
+		z.Free(b.Deltas)
+		FreeBlock(b)
+	}
+}
+
 func (e *Encoder) packBlock() {
 	if len(e.uids) == 0 {
 		return
 	}
-	block := &pb.UidBlock{Base: e.uids[0], NumUids: uint32(len(e.uids))}
+
+	// Allocate blocks manually.
+	block := AllocateBlock()
+	block.Base = e.uids[0]
+	block.NumUids = uint32(len(e.uids))
+
+	// block := &pb.UidBlock{Base: e.uids[0], NumUids: uint32(len(e.uids))}
 	last := e.uids[0]
 	e.uids = e.uids[1:]
 
-	var out bytes.Buffer
+	var out z.Buffer
+	// We are not releasing the allocated memory here. Instead, block.Deltas would need to be
+	// released at the end.
+
 	buf := make([]byte, 17)
 	tmpUids := make([]uint32, 4)
 	for {
@@ -79,6 +121,7 @@ func (e *Encoder) packBlock() {
 		e.uids = e.uids[4:]
 	}
 
+	// TODO: Instead of one z.Buffer for every delta, allocate one big one.
 	block.Deltas = out.Bytes()
 	e.pack.Blocks = append(e.pack.Blocks, block)
 }
@@ -102,7 +145,7 @@ func (e *Encoder) Add(uid uint64) {
 	}
 }
 
-// Done returns the final output of the encoder.
+// Done returns the final output of the encoder. This UidPack MUST BE FREED via a call to FreePack.
 func (e *Encoder) Done() *pb.UidPack {
 	e.packBlock()
 	return e.pack

codec/codec_test.go

@@ -47,7 +47,8 @@ func TestUidPack(t *testing.T) {
 	rand.Seed(time.Now().UnixNano())
 
 	// Some edge case tests.
-	Encode([]uint64{}, 128)
+	pack := Encode([]uint64{}, 128)
+	FreePack(pack)
 	require.Equal(t, 0, ApproxLen(&pb.UidPack{}))
 	require.Equal(t, 0, len(Decode(&pb.UidPack{}, 0)))
 
@@ -63,6 +64,7 @@ func TestUidPack(t *testing.T) {
 		require.Equal(t, len(expected), ExactLen(pack))
 		actual := Decode(pack, 0)
 		require.Equal(t, expected, actual)
+		FreePack(pack)
 	}
 }
 
@@ -73,6 +75,7 @@ func TestSeek(t *testing.T) {
 		enc.Add(uint64(i))
 	}
 	pack := enc.Done()
+	defer FreePack(pack)
 	dec := Decoder{Pack: pack}
 
 	tests := []struct {
@@ -121,6 +124,7 @@ func TestLinearSeek(t *testing.T) {
 		enc.Add(uint64(i))
 	}
 	pack := enc.Done()
+	defer FreePack(pack)
 	dec := Decoder{Pack: pack}
 
 	for i := 0; i < 2*N; i += 10 {
@@ -150,6 +154,7 @@ func TestDecoder(t *testing.T) {
 		expected = append(expected, uint64(i))
 	}
 	pack := enc.Done()
+	defer FreePack(pack)
 
 	dec := Decoder{Pack: pack}
 	for i := 3; i < N; i += 3 {
@@ -215,6 +220,7 @@ func benchmarkUidPackEncode(b *testing.B, blockSize int) {
 	for i := 0; i < b.N; i++ {
 		pack := Encode(uids, blockSize)
 		out, err := pack.Marshal()
+		FreePack(pack)
 		if err != nil {
 			b.Fatalf("Error marshaling uid pack: %s", err.Error())
 		}
@@ -249,6 +255,7 @@ func benchmarkUidPackDecode(b *testing.B, blockSize int) {
 
 	pack := Encode(uids, blockSize)
 	data, err := pack.Marshal()
+	defer FreePack(pack)
 	x.Check(err)
 	b.Logf("Output size: %s. Compression: %.2f",
 		humanize.Bytes(uint64(len(data))),
@@ -301,6 +308,7 @@ func newUidPack(data []uint64) *pb.UidPack {
 
 func TestCopyUidPack(t *testing.T) {
 	pack := newUidPack([]uint64{1, 2, 3, 4, 5})
+	defer FreePack(pack)
 	copy := CopyUidPack(pack)
 	require.Equal(t, Decode(pack, 0), Decode(copy, 0))
 }