fix(table): Add onDisk size (#1569)

This PR adds onDisk size to the table index

fb/flatbuffer.fbs

@@ -19,10 +19,10 @@ namespace fb;
 table TableIndex {
   offsets:[BlockOffset];
   bloom_filter:[ubyte];
-  estimated_size:uint32;
   max_version:uint64;
-  uncompressed_size:uint32;
   key_count:uint32;
+  uncompressed_size:uint32;
+  on_disk_size:uint32;
 }
 
 table BlockOffset {

levels.go

@@ -1459,7 +1459,7 @@ type TableInfo struct {
 	Left             []byte
 	Right            []byte
 	KeyCount         uint32 // Number of keys in the table
-	EstimatedSz      uint32
+	OnDiskSize       uint32
 	UncompressedSize uint32
 	MaxVersion       uint64
 	IndexSz          int
@@ -1476,7 +1476,7 @@ func (s *levelsController) getTableInfo() (result []TableInfo) {
 				Left:             t.Smallest(),
 				Right:            t.Biggest(),
 				KeyCount:         t.KeyCount(),
-				EstimatedSz:      t.EstimatedSize(),
+				OnDiskSize:       t.OnDiskSize(),
 				IndexSz:          t.IndexSize(),
 				BloomFilterSize:  t.BloomFilterSize(),
 				UncompressedSize: t.UncompressedSize(),

table/builder.go

@@ -85,12 +85,12 @@ 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.
-	offsets       *z.Buffer
-	estimatedSize uint32
-	keyHashes     []uint32 // Used for building the bloomfilter.
-	opt           *Options
-	maxVersion    uint64
+	entryOffsets []uint32 // Offsets of entries present in current block.
+	offsets      *z.Buffer
+	onDiskSize   uint32
+	keyHashes    []uint32 // Used for building the bloomfilter.
+	opt          *Options
+	maxVersion   uint64
 
 	// Used to concurrently compress/encrypt blocks.
 	wg        sync.WaitGroup
@@ -230,10 +230,9 @@ func (b *Builder) addHelper(key []byte, v y.ValueStruct, vpLen uint32) {
 	}
 	b.sz += v.Encode(b.buf[b.sz:])
 
-	// Size of KV on SST.
-	sstSz := uint32(headerSize) + uint32(len(diffKey)) + v.EncodedSize()
-	// Total estimated size = size on SST + size on vlog (length of value pointer).
-	b.estimatedSize += (sstSz + vpLen)
+	// Add the vpLen to the onDisk size. We'll add the size of the block to
+	// onDisk size in Finish() function.
+	b.onDiskSize += vpLen
 }
 
 // grow increases the size of b.buf by atleast 50%.
@@ -438,6 +437,8 @@ func (b *Builder) Finish(allocate bool) []byte {
 		b.sz = dstLen
 	}
 
+	// b.sz is the total size of the compressed table without the index.
+	b.onDiskSize += b.sz
 	var f y.Filter
 	if b.opt.BloomFalsePositive > 0 {
 		bits := y.BloomBitsPerKey(len(b.keyHashes), b.opt.BloomFalsePositive)
@@ -566,13 +567,17 @@ func (b *Builder) buildIndex(bloom []byte, tableSz uint32) []byte {
 	fb.TableIndexStart(builder)
 	fb.TableIndexAddOffsets(builder, boEnd)
 	fb.TableIndexAddBloomFilter(builder, bfoff)
-	fb.TableIndexAddEstimatedSize(builder, b.estimatedSize)
 	fb.TableIndexAddMaxVersion(builder, b.maxVersion)
 	fb.TableIndexAddUncompressedSize(builder, tableSz)
 	fb.TableIndexAddKeyCount(builder, uint32(len(b.keyHashes)))
+	fb.TableIndexAddOnDiskSize(builder, b.onDiskSize)
 	builder.Finish(fb.TableIndexEnd(builder))
 
-	return builder.FinishedBytes()
+	buf := builder.FinishedBytes()
+	index := fb.GetRootAsTableIndex(buf, 0)
+	// Mutate the ondisk size to include the size of the index as well.
+	y.AssertTrue(index.MutateOnDiskSize(index.OnDiskSize() + uint32(len(buf))))
+	return buf
 }
 
 // writeBlockOffsets writes all the blockOffets in b.offsets and returns the

table/table.go

@@ -108,7 +108,7 @@ type Table struct {
 	Checksum  []byte
 	CreatedAt time.Time
 	// Stores the total size of key-values stored in this table (including the size on vlog).
-	estimatedSize  uint32
+	onDiskSize     uint32
 	indexStart     int
 	indexLen       int
 	hasBloomFilter bool
@@ -388,17 +388,7 @@ func (t *Table) initIndex() (*fb.BlockOffset, error) {
 		t._index = index
 	}
 
-	if t.opt.Compression == options.None {
-		t.estimatedSize = index.EstimatedSize()
-	} else {
-		// TODO(ibrahim): This estimatedSize doesn't make any sense. If it is tracking the size of
-		// values including in value log, then index.EstimatedSize() should be used irrespective of
-		// compression or not.
-		//
-		// Due to compression the real size on disk is much
-		// smaller than what we estimate from index.EstimatedSize.
-		t.estimatedSize = uint32(t.tableSize)
-	}
+	t.onDiskSize = index.OnDiskSize()
 	t.hasBloomFilter = len(index.BloomFilterBytes()) > 0
 
 	var bo fb.BlockOffset
@@ -580,7 +570,7 @@ func (t *Table) indexKey() uint64 {
 	return t.id
 }
 
-// UncompressedSize is the size of table index in bytes.
+// UncompressedSize is the size uncompressed data stored in this file.
 func (t *Table) UncompressedSize() uint32 {
 	return t.fetchIndex().UncompressedSize()
 }
@@ -600,9 +590,9 @@ func (t *Table) BloomFilterSize() int {
 	return t.fetchIndex().BloomFilterLength()
 }
 
-// EstimatedSize returns the total size of key-values stored in this table (including the
+// OnDiskSize returns the total size of key-values stored in this table (including the
 // disk space occupied on the value log).
-func (t *Table) EstimatedSize() uint32 { return t.estimatedSize }
+func (t *Table) OnDiskSize() uint32 { return t.onDiskSize }
 
 // Size is its file size in bytes
 func (t *Table) Size() int64 { return int64(t.tableSize) }

table/table_test.go

@@ -865,31 +865,6 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
-func TestOpenKVSize(t *testing.T) {
-	t.Run("compression", func(t *testing.T) {
-		// When compression is on
-		opts := getTestTableOptions()
-		opts.Compression = options.ZSTD
-		table := buildTestTable(t, "foo", 1000, opts)
-		defer table.DecrRef()
-
-		// The estimated size is same as table size in case compression is enabled.
-		require.Equal(t, uint32(table.tableSize), table.EstimatedSize())
-	})
-
-	t.Run("no compressin", func(t *testing.T) {
-		// When compression is off
-		opts := getTestTableOptions()
-		opts.Compression = options.None
-		table := buildTestTable(t, "foo", 1, opts)
-		defer table.DecrRef()
-
-		stat, err := table.Fd.Stat()
-		require.NoError(t, err)
-		require.Less(t, table.EstimatedSize(), uint32(stat.Size()))
-	})
-}
-
 // Run this test with command "go test -race -run TestDoesNotHaveRace"
 func TestDoesNotHaveRace(t *testing.T) {
 	opts := getTestTableOptions()