Fix deadlock when flushing discard stats.

value.go

@@ -1386,47 +1386,50 @@ func (vlog *valueLog) runGC(discardRatio float64, head valuePointer) error {
 
 func (vlog *valueLog) updateDiscardStats(stats map[uint32]int64) error {
 	vlog.lfDiscardStats.Lock()
-	defer vlog.lfDiscardStats.Unlock()
 
 	for fid, sz := range stats {
 		vlog.lfDiscardStats.m[fid] += sz
 		vlog.lfDiscardStats.updatesSinceFlush++
 	}
 	if vlog.lfDiscardStats.updatesSinceFlush > discardStatsFlushThreshold {
-		if err := vlog.flushDiscardStats(); err != nil {
-			return err
-		}
-		vlog.lfDiscardStats.updatesSinceFlush = 0
+		vlog.lfDiscardStats.Unlock()
+		// flushDiscardStats also acquires lock. So, we need to unlock here.
+		return vlog.flushDiscardStats()
 	}
+	vlog.lfDiscardStats.Unlock()
 	return nil
 }
 
 // flushDiscardStats inserts discard stats into badger. Returns error on failure.
 func (vlog *valueLog) flushDiscardStats() error {
 	vlog.lfDiscardStats.Lock()
+	defer vlog.lfDiscardStats.Unlock()
+
 	if len(vlog.lfDiscardStats.m) == 0 {
-		vlog.lfDiscardStats.Unlock()
 		return nil
 	}
-	vlog.lfDiscardStats.Unlock()
-
 	entries := []*Entry{{
 		Key:   y.KeyWithTs(lfDiscardStatsKey, 1),
 		Value: vlog.encodedDiscardStats(),
 	}}
 	req, err := vlog.db.sendToWriteCh(entries)
-	if err != nil {
+	if err == ErrBlockedWrites {
+		// We'll block write while closing db.
+		// When L0 compaction in close may push discard stats.
+		// So ignoring it.
+		// https://github.com/dgraph-io/badger/issues/970
+		return nil
+	} else if err != nil {
 		return errors.Wrapf(err, "failed to push discard stats to write channel")
 	}
+	vlog.lfDiscardStats.updatesSinceFlush = 0
 	return req.Wait()
 }
 
 // encodedDiscardStats returns []byte representation of lfDiscardStats
 // This will be called while storing stats in BadgerDB
+// caller should acquire lock before encoding the stats.
 func (vlog *valueLog) encodedDiscardStats() []byte {
-	vlog.lfDiscardStats.Lock()
-	defer vlog.lfDiscardStats.Unlock()
-
 	encodedStats, _ := json.Marshal(vlog.lfDiscardStats.m)
 	return encodedStats
 }

value_test.go

@@ -1085,3 +1085,38 @@ func TestDiscardStatsMove(t *testing.T) {
 	require.Equal(t, stat, db.vlog.lfDiscardStats.m)
 	require.NoError(t, db.Close())
 }
+
+// This test ensures, flushDiscardStats() doesn't crash.
+func TestBlockedDiscardStats(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err)
+	defer os.Remove(dir)
+	db, err := Open(getTestOptions(dir))
+	require.NoError(t, err)
+	// Set discard stats.
+	db.vlog.lfDiscardStats = &lfDiscardStats{
+		m: map[uint32]int64{0: 0},
+	}
+	db.blockWrite()
+	require.NoError(t, db.vlog.flushDiscardStats())
+	db.unblockWrite()
+	require.NoError(t, db.Close())
+}
+
+// Regression test for https://github.com/dgraph-io/badger/issues/970
+func TestBlockedDiscardStatsOnClose(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err)
+	defer os.Remove(dir)
+
+	db, err := Open(getTestOptions(dir))
+	require.NoError(t, err)
+	// Set discard stats.
+	db.vlog.lfDiscardStats = &lfDiscardStats{
+		m: map[uint32]int64{0: 0},
+	}
+	// This is important. Set updateSinceFlush to discardStatsFlushThresold so
+	// that the next update call flushes the discard stats.
+	db.vlog.lfDiscardStats.updatesSinceFlush = discardStatsFlushThreshold + 1
+	require.NoError(t, db.Close())
+}