Store total key-value size in table footer (#1137)

This PR stores the total key-value size in a table in the table
footer. The key-value size can be accessed via db.Tables(..) call.
It returns the list of all tables along with the total size of key-values.

db.go

@@ -899,11 +899,16 @@ func buildL0Table(ft flushTask, bopts table.Options) []byte {
 	defer iter.Close()
 	b := table.NewTableBuilder(bopts)
 	defer b.Close()
+	var vp valuePointer
 	for iter.SeekToFirst(); iter.Valid(); iter.Next() {
 		if len(ft.dropPrefix) > 0 && bytes.HasPrefix(iter.Key(), ft.dropPrefix) {
 			continue
 		}
-		b.Add(iter.Key(), iter.Value())
+		vs := iter.Value()
+		if vs.Meta&bitValuePointer > 0 {
+			vp.Decode(vs.Value)
+		}
+		b.Add(iter.Key(), iter.Value(), vp.Len)
 	}
 	return b.Finish()
 }

db2_test.go

@@ -538,7 +538,7 @@ func createTableWithRange(t *testing.T, db *DB, start, end int) *table.Table {
 		binary.BigEndian.PutUint64(key[:], uint64(i))
 		key = y.KeyWithTs(key, uint64(0))
 		val := y.ValueStruct{Value: []byte(fmt.Sprintf("%d", i))}
-		b.Add(key, val)
+		b.Add(key, val, 0)
 	}
 
 	fileID := db.lc.reserveFileID()

levels.go

@@ -508,6 +508,7 @@ func (s *levelsController) compactBuildTables(
 	resultCh := make(chan newTableResult)
 	var numBuilds, numVersions int
 	var lastKey, skipKey []byte
+	var vp valuePointer
 	for it.Valid() {
 		timeStart := time.Now()
 		dk, err := s.kv.registry.latestDataKey()
@@ -584,7 +585,10 @@ func (s *levelsController) compactBuildTables(
 				}
 			}
 			numKeys++
-			builder.Add(it.Key(), it.Value())
+			if vs.Meta&bitValuePointer > 0 {
+				vp.Decode(vs.Value)
+			}
+			builder.Add(it.Key(), vs, vp.Len)
 		}
 		// It was true that it.Valid() at least once in the loop above, which means we
 		// called Add() at least once, and builder is not Empty().
@@ -1009,11 +1013,12 @@ func (s *levelsController) appendIterators(
 
 // TableInfo represents the information about a table.
 type TableInfo struct {
-	ID       uint64
-	Level    int
-	Left     []byte
-	Right    []byte
-	KeyCount uint64 // Number of keys in the table
+	ID          uint64
+	Level       int
+	Left        []byte
+	Right       []byte
+	KeyCount    uint64 // Number of keys in the table
+	EstimatedSz uint64
 }
 
 func (s *levelsController) getTableInfo(withKeysCount bool) (result []TableInfo) {
@@ -1030,11 +1035,12 @@ func (s *levelsController) getTableInfo(withKeysCount bool) (result []TableInfo)
 			}
 
 			info := TableInfo{
-				ID:       t.ID(),
-				Level:    l.level,
-				Left:     t.Smallest(),
-				Right:    t.Biggest(),
-				KeyCount: count,
+				ID:          t.ID(),
+				Level:       l.level,
+				Left:        t.Smallest(),
+				Right:       t.Biggest(),
+				KeyCount:    count,
+				EstimatedSz: t.EstimatedSize(),
 			}
 			result = append(result, info)
 		}

manifest_test.go

@@ -153,7 +153,7 @@ func buildTable(t *testing.T, keyValues [][]string, bopts table.Options) *os.Fil
 			Value:    []byte(kv[1]),
 			Meta:     'A',
 			UserMeta: 0,
-		})
+		}, 0)
 	}
 	_, err = f.Write(b.Finish())
 	require.NoError(t, err, "unable to write to file.")

pb/pb.proto

@@ -43,7 +43,7 @@ message ManifestChangeSet {
 }
 
 enum EncryptionAlgo {
-                aes = 0;  
+  aes = 0;
 }
 
 message ManifestChange {
@@ -68,6 +68,7 @@ message BlockOffset {
 message TableIndex {
   repeated BlockOffset offsets = 1;
   bytes bloom_filter = 2;
+  uint64 estimated_size = 3;
 }
 
 message Checksum {
@@ -84,4 +85,4 @@ message DataKey {
   bytes  data       = 2;
   bytes  iv         = 3;
   int64  created_at = 4;
-}
+}

stream_writer.go

@@ -362,7 +362,11 @@ func (w *sortedWriter) Add(key []byte, vs y.ValueStruct) error {
 	}
 
 	w.lastKey = y.SafeCopy(w.lastKey, key)
-	w.builder.Add(key, vs)
+	var vp valuePointer
+	if vs.Meta&bitValuePointer > 0 {
+		vp.Decode(vs.Value)
+	}
+	w.builder.Add(key, vs, vp.Len)
 	return nil
 }
 

table/builder.go

@@ -69,11 +69,9 @@ type Builder struct {
 	baseKey      []byte   // Base key for the current block.
 	baseOffset   uint32   // Offset for the current block.
 	entryOffsets []uint32 // Offsets of entries present in current block.
-
-	tableIndex *pb.TableIndex
-	keyHashes  []uint64
-
-	opt *Options
+	tableIndex   *pb.TableIndex
+	keyHashes    []uint64 // Used for building the bloomfilter.
+	opt          *Options
 }
 
 // NewTableBuilder makes a new TableBuilder.
@@ -103,7 +101,7 @@ func (b *Builder) keyDiff(newKey []byte) []byte {
 	return newKey[i:]
 }
 
-func (b *Builder) addHelper(key []byte, v y.ValueStruct) {
+func (b *Builder) addHelper(key []byte, v y.ValueStruct, vpLen uint64) {
 	b.keyHashes = append(b.keyHashes, farm.Fingerprint64(y.ParseKey(key)))
 
 	// diffKey stores the difference of key with baseKey.
@@ -131,6 +129,10 @@ func (b *Builder) addHelper(key []byte, v y.ValueStruct) {
 	b.buf.Write(diffKey) // We only need to store the key difference.
 
 	v.EncodeTo(b.buf)
+	// Size of KV on SST.
+	sstSz := uint64(uint32(headerSize) + uint32(len(diffKey)) + v.EncodedSize())
+	// Total estimated size = size on SST + size on vlog (length of value pointer).
+	b.tableIndex.EstimatedSize += (sstSz + vpLen)
 }
 
 /*
@@ -210,7 +212,7 @@ func (b *Builder) shouldFinishBlock(key []byte, value y.ValueStruct) bool {
 }
 
 // Add adds a key-value pair to the block.
-func (b *Builder) Add(key []byte, value y.ValueStruct) {
+func (b *Builder) Add(key []byte, value y.ValueStruct, valueLen uint32) {
 	if b.shouldFinishBlock(key, value) {
 		b.finishBlock()
 		// Start a new block. Initialize the block.
@@ -219,7 +221,7 @@ func (b *Builder) Add(key []byte, value y.ValueStruct) {
 		b.baseOffset = uint32(b.buf.Len())
 		b.entryOffsets = b.entryOffsets[:0]
 	}
-	b.addHelper(key, value)
+	b.addHelper(key, value, uint64(valueLen))
 }
 
 // TODO: vvv this was the comment on ReachedCapacity.

table/builder_test.go

@@ -74,7 +74,7 @@ func TestTableIndex(t *testing.T) {
 					blockCount++
 					blockFirstKeys = append(blockFirstKeys, k)
 				}
-				builder.Add(k, vs)
+				builder.Add(k, vs, 0)
 			}
 			_, err = f.Write(builder.Finish())
 			require.NoError(t, err, "unable to write to file")
@@ -129,7 +129,7 @@ func BenchmarkBuilder(b *testing.B) {
 		builder := NewTableBuilder(opts)
 
 		for i := 0; i < keysCount; i++ {
-			builder.Add(key(i), vs)
+			builder.Add(key(i), vs, 0)
 		}
 
 		_ = builder.Finish()

table/table.go

@@ -99,6 +99,8 @@ type Table struct {
 
 	bf       *z.Bloom
 	Checksum []byte
+	// Stores the total size of key-values stored in this table (including the size on vlog).
+	estimatedSize uint64
 
 	IsInmemory bool // Set to true if the table is on level 0 and opened in memory.
 	opt        *Options
@@ -351,6 +353,7 @@ func (t *Table) readIndex() error {
 	err := proto.Unmarshal(data, &index)
 	y.Check(err)
 
+	t.estimatedSize = index.EstimatedSize
 	t.bf = z.JSONUnmarshal(index.BloomFilter)
 	t.blockIndex = index.Offsets
 	return nil
@@ -439,6 +442,10 @@ func (t *Table) blockCacheKey(idx int) uint64 {
 	return (t.ID() << 32) | uint64(idx)
 }
 
+// EstimatedSize returns the total size of key-values stored in this table (including the
+// disk space occupied on the value log).
+func (t *Table) EstimatedSize() uint64 { return t.estimatedSize }
+
 // Size is its file size in bytes
 func (t *Table) Size() int64 { return int64(t.tableSize) }