fix(replay) - Update head for LSM entires also (#1456)

#1372 tried to fix the `replay from
start` issue but it partially fixed the issue. The head was not being updated
in case all the entries are inserted only in the LSM tree.
This commit fixes it.

backup_test.go

@@ -112,6 +112,7 @@ func TestBackupRestore1(t *testing.T) {
 		return nil
 	})
 	require.NoError(t, err)
+	require.Equal(t, db.orc.nextTs(), uint64(3))
 }
 
 func TestBackupRestore2(t *testing.T) {
@@ -163,6 +164,9 @@ func TestBackupRestore2(t *testing.T) {
 	err = db2.Load(&backup, 16)
 	require.NoError(t, err)
 
+	// Check nextTs is correctly set.
+	require.Equal(t, db1.orc.nextTs(), db2.orc.nextTs())
+
 	for i := byte(1); i < N; i++ {
 		err = db2.View(func(tx *Txn) error {
 			k := append(key1, i)
@@ -210,6 +214,9 @@ func TestBackupRestore2(t *testing.T) {
 	err = db3.Load(&backup, 16)
 	require.NoError(t, err)
 
+	// Check nextTs is correctly set.
+	require.Equal(t, db2.orc.nextTs(), db3.orc.nextTs())
+
 	for i := byte(1); i < N; i++ {
 		err = db3.View(func(tx *Txn) error {
 			k := append(key1, i)
@@ -325,6 +332,7 @@ func TestBackupRestore3(t *testing.T) {
 	N := 1000
 	entries := createEntries(N)
 
+	var db1NextTs uint64
 	// backup
 	{
 		db1, err := Open(DefaultOptions(filepath.Join(tmpdir, "backup1")))
@@ -335,6 +343,8 @@ func TestBackupRestore3(t *testing.T) {
 
 		_, err = db1.Backup(&bb, 0)
 		require.NoError(t, err)
+
+		db1NextTs = db1.orc.nextTs()
 		require.NoError(t, db1.Close())
 	}
 	require.True(t, len(entries) == N)
@@ -345,7 +355,9 @@ func TestBackupRestore3(t *testing.T) {
 	require.NoError(t, err)
 
 	defer db2.Close()
+	require.NotEqual(t, db1NextTs, db2.orc.nextTs())
 	require.NoError(t, db2.Load(&bb, 16))
+	require.Equal(t, db1NextTs, db2.orc.nextTs())
 
 	// verify
 	err = db2.View(func(txn *Txn) error {
@@ -383,6 +395,7 @@ func TestBackupLoadIncremental(t *testing.T) {
 	updates := make(map[int]byte)
 	var bb bytes.Buffer
 
+	var db1NextTs uint64
 	// backup
 	{
 		db1, err := Open(DefaultOptions(filepath.Join(tmpdir, "backup2")))
@@ -439,6 +452,9 @@ func TestBackupLoadIncremental(t *testing.T) {
 		require.NoError(t, err)
 		_, err = db1.Backup(&bb, since)
 		require.NoError(t, err)
+
+		db1NextTs = db1.orc.nextTs()
+
 		require.NoError(t, db1.Close())
 	}
 	require.True(t, len(entries) == N)
@@ -450,7 +466,9 @@ func TestBackupLoadIncremental(t *testing.T) {
 
 	defer db2.Close()
 
+	require.NotEqual(t, db1NextTs, db2.orc.nextTs())
 	require.NoError(t, db2.Load(&bb, 16))
+	require.Equal(t, db1NextTs, db2.orc.nextTs())
 
 	// verify
 	actual := make(map[int]byte)
@@ -517,6 +535,8 @@ func TestBackupBitClear(t *testing.T) {
 	_, err = db.Backup(bak, 0)
 	require.NoError(t, err)
 	require.NoError(t, bak.Close())
+
+	oldValue := db.orc.nextTs()
 	require.NoError(t, db.Close())
 
 	opt = getTestOptions(dir)
@@ -530,6 +550,8 @@ func TestBackupBitClear(t *testing.T) {
 	defer bak.Close()
 
 	require.NoError(t, db.Load(bak, 16))
+	// Ensure nextTs is still the same.
+	require.Equal(t, oldValue, db.orc.nextTs())
 
 	require.NoError(t, db.View(func(txn *Txn) error {
 		e, err := txn.Get(key)

db.go

@@ -136,11 +136,11 @@ func (db *DB) replayFunction() func(Entry, valuePointer) error {
 		} else {
 			nv = vp.Encode()
 			meta = meta | bitValuePointer
-			// Update vhead. If the crash happens while replay was in progess
-			// and the head is not updated, we will end up replaying all the
-			// files again.
-			db.updateHead([]valuePointer{vp})
 		}
+		// Update vhead. If the crash happens while replay was in progess
+		// and the head is not updated, we will end up replaying all the
+		// files starting from file zero, again.
+		db.updateHead([]valuePointer{vp})
 
 		v := y.ValueStruct{
 			Value:     nv,
@@ -994,24 +994,41 @@ type flushTask struct {
 	dropPrefixes [][]byte
 }
 
-// handleFlushTask must be run serially.
-func (db *DB) handleFlushTask(ft flushTask) error {
-	// There can be a scenario, when empty memtable is flushed. For example, memtable is empty and
-	// after writing request to value log, rotation count exceeds db.LogRotatesToFlush.
-	if ft.mt.Empty() {
+func (db *DB) pushHead(ft flushTask) error {
+	// We don't need to store head pointer in the in-memory mode since we will
+	// never be replay anything.
+	if db.opt.InMemory {
 		return nil
 	}
+	// Ensure we never push a zero valued head pointer.
+	if ft.vptr.IsZero() {
+		return errors.New("Head should not be zero")
+	}
 
 	// Store badger head even if vptr is zero, need it for readTs
 	db.opt.Debugf("Storing value log head: %+v\n", ft.vptr)
-	db.opt.Debugf("Storing offset: %+v\n", ft.vptr)
 	val := ft.vptr.Encode()
 
 	// Pick the max commit ts, so in case of crash, our read ts would be higher than all the
 	// commits.
 	headTs := y.KeyWithTs(head, db.orc.nextTs())
 	ft.mt.Put(headTs, y.ValueStruct{Value: val})
 
+	return nil
+}
+
+// handleFlushTask must be run serially.
+func (db *DB) handleFlushTask(ft flushTask) error {
+	// There can be a scenario, when empty memtable is flushed. For example, memtable is empty and
+	// after writing request to value log, rotation count exceeds db.LogRotatesToFlush.
+	if ft.mt.Empty() {
+		return nil
+	}
+
+	if err := db.pushHead(ft); err != nil {
+		return err
+	}
+
 	dk, err := db.registry.latestDataKey()
 	if err != nil {
 		return y.Wrapf(err, "failed to get datakey in db.handleFlushTask")

db_test.go

@@ -2001,8 +2001,9 @@ func TestNoCrash(t *testing.T) {
 	}
 
 	db.Lock()
-	// make head to point to first file
-	db.vhead = valuePointer{0, 0, 0}
+	// make head to point to second file. We cannot make it point to the first
+	// vlog file because we cannot push a zero head pointer.
+	db.vhead = valuePointer{1, 0, 0}
 	db.Unlock()
 	db.Close()