chore(stream): log the estimated size of data to be streamed (#1632)

When you're scanning a lot, but sending little data, the stream framework seems to work really slowly. But, that's deceiving, because it is scanning really fast.
We calculate the rough uncompressed size of the data the stream framework will send. We also periodically log the amount of data scanned by produceKVs.

Author: NamanJain8

db.go

@@ -1300,6 +1300,19 @@ func (db *DB) Levels() []LevelInfo {
 	return db.lc.getLevelInfo()
 }
 
+// EstimateSize can be used to get rough estimate of data size for a given prefix.
+func (db *DB) EstimateSize(prefix []byte) (uint64, uint64) {
+	var onDiskSize, uncompressedSize uint64
+	tables := db.Tables()
+	for _, ti := range tables {
+		if bytes.HasPrefix(ti.Left, prefix) && bytes.HasPrefix(ti.Right, prefix) {
+			onDiskSize += uint64(ti.OnDiskSize)
+			uncompressedSize += uint64(ti.UncompressedSize)
+		}
+	}
+	return onDiskSize, uncompressedSize
+}
+
 // KeySplits can be used to get rough key ranges to divide up iteration over
 // the DB.
 func (db *DB) KeySplits(prefix []byte) []string {

iterator.go

@@ -429,7 +429,8 @@ type Iterator struct {
 
 	lastKey []byte // Used to skip over multiple versions of the same key.
 
-	closed bool
+	closed  bool
+	scanned int // Used to estimate the size of data scanned by iterator.
 
 	// ThreadId is an optional value that can be set to identify which goroutine created
 	// the iterator. It can be used, for example, to uniquely identify each of the
@@ -559,11 +560,11 @@ func (it *Iterator) Close() {
 func (it *Iterator) Next() {
 	// Reuse current item
 	it.item.wg.Wait() // Just cleaner to wait before pushing to avoid doing ref counting.
+	it.scanned += len(it.item.key) + len(it.item.val) + len(it.item.vptr) + 2
 	it.waste.push(it.item)
 
 	// Set next item to current
 	it.item = it.data.pop()
-
 	for it.iitr.Valid() {
 		if it.parseItem() {
 			// parseItem calls one extra next.

stream.go

@@ -87,6 +87,7 @@ type Stream struct {
 	kvChan       chan *z.Buffer
 	nextStreamId uint32
 	doneMarkers  bool
+	scanned      uint64 // used to estimate the ETA for data scan.
 }
 
 // SendDoneMarkers when true would send out done markers on the stream. False by default.
@@ -202,11 +203,14 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
 
 		// This unique stream id is used to identify all the keys from this iteration.
 		streamId := atomic.AddUint32(&st.nextStreamId, 1)
+		var scanned int
 
 		sendIt := func() error {
 			select {
 			case st.kvChan <- outList:
 				outList = z.NewBuffer(2 * batchSize)
+				atomic.AddUint64(&st.scanned, uint64(itr.scanned-scanned))
+				scanned = itr.scanned
 			case <-ctx.Done():
 				return ctx.Err()
 			}
@@ -227,6 +231,7 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
 			if len(kr.right) > 0 && bytes.Compare(item.Key(), kr.right) >= 0 {
 				break
 			}
+
 			// Check if we should pick this key.
 			if st.ChooseKey != nil && !st.ChooseKey(item) {
 				continue
@@ -281,6 +286,10 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
 }
 
 func (st *Stream) streamKVs(ctx context.Context) error {
+	onDiskSize, uncompressedSize := st.db.EstimateSize(st.Prefix)
+	st.db.opt.Infof("%s Streaming about %s of uncompressed data (%s on disk)\n",
+		st.LogPrefix, humanize.IBytes(uncompressedSize), humanize.IBytes(onDiskSize))
+
 	var bytesSent uint64
 	t := time.NewTicker(time.Second)
 	defer t.Stop()
@@ -340,9 +349,10 @@ outer:
 				continue
 			}
 			speed := bytesSent / durSec
-
-			st.db.opt.Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec, jemalloc: %s\n",
-				st.LogPrefix, y.FixedDuration(dur), humanize.IBytes(bytesSent),
+			scanned := atomic.LoadUint64(&st.scanned)
+			st.db.opt.Infof("%s Time elapsed: %s, scanned: ~%s/%s, bytes sent: %s, speed: %s/sec,"+
+				"jemalloc: %s\n", st.LogPrefix, y.FixedDuration(dur), humanize.IBytes(scanned),
+				humanize.IBytes(uncompressedSize), humanize.IBytes(bytesSent),
 				humanize.IBytes(speed), humanize.IBytes(uint64(z.NumAllocBytes())))
 
 		case kvs, ok := <-st.kvChan: