bugfix:dead lock on store.mu when store.Compact in store.Restore happens

mvcc/kvstore.go

@@ -139,6 +139,8 @@ func NewStore(lg *zap.Logger, b backend.Backend, le lease.Lessor, ig ConsistentI
 	tx.Unlock()
 	s.b.ForceCommit()
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	if err := s.restore(); err != nil {
 		// TODO: return the error instead of panic here?
 		panic("failed to recover store from backend")
@@ -225,25 +227,20 @@ func (s *store) HashByRev(rev int64) (hash uint32, currentRev int64, compactRev
 	return hash, currentRev, compactRev, err
 }
 
-func (s *store) Compact(rev int64) (<-chan struct{}, error) {
+func (s *store) updateCompactRev(rev int64) (<-chan struct{}, error) {
-	s.mu.Lock()
 	s.revMu.Lock()
 	if rev <= s.compactMainRev {
 		ch := make(chan struct{})
 		f := func(ctx context.Context) { s.compactBarrier(ctx, ch) }
 		s.fifoSched.Schedule(f)
-		s.mu.Unlock()
 		s.revMu.Unlock()
 		return ch, ErrCompacted
 	}
 	if rev > s.currentRev {
-		s.mu.Unlock()
 		s.revMu.Unlock()
 		return nil, ErrFutureRev
 	}
 
-	start := time.Now()
-
 	s.compactMainRev = rev
 
 	rbytes := newRevBytes()
@@ -256,8 +253,13 @@ func (s *store) Compact(rev int64) (<-chan struct{}, error) {
 	// ensure that desired compaction is persisted
 	s.b.ForceCommit()
 
-	s.mu.Unlock()
 	s.revMu.Unlock()
+
+	return nil, nil
+}
+
+func (s *store) compact(rev int64) (<-chan struct{}, error) {
+	start := time.Now()
 	keep := s.kvindex.Compact(rev)
 	ch := make(chan struct{})
 	var j = func(ctx context.Context) {
@@ -278,6 +280,29 @@ func (s *store) Compact(rev int64) (<-chan struct{}, error) {
 	return ch, nil
 }
 
+func (s *store) compactLockfree(rev int64) (<-chan struct{}, error) {
+	ch, err := s.updateCompactRev(rev)
+	if nil != err {
+		return ch, err
+	}
+
+	return s.compact(rev)
+}
+
+func (s *store) Compact(rev int64) (<-chan struct{}, error) {
+	s.mu.Lock()
+
+	ch, err := s.updateCompactRev(rev)
+
+	if err != nil {
+		s.mu.Unlock()
+		return ch, err
+	}
+	s.mu.Unlock()
+
+	return s.compact(rev)
+}
+
 // DefaultIgnores is a map of keys to ignore in hash checking.
 var DefaultIgnores map[backend.IgnoreKey]struct{}
 
@@ -415,7 +440,7 @@ func (s *store) restore() error {
 	tx.Unlock()
 
 	if scheduledCompact != 0 {
-		s.Compact(scheduledCompact)
+		s.compactLockfree(scheduledCompact)
 
 		if s.lg != nil {
 			s.lg.Info(

mvcc/kvstore_test.go

@@ -466,6 +466,8 @@ func TestRestoreDelete(t *testing.T) {
 }
 
 func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
+	tests := []string{"recreate", "restore"}
+	for _, test := range tests {
 		b, tmpPath := backend.NewDefaultTmpBackend()
 		s0 := NewStore(zap.NewExample(), b, &lease.FakeLessor{}, nil)
 		defer os.Remove(tmpPath)
@@ -484,12 +486,19 @@ func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
 
 		s0.Close()
 
-	s1 := NewStore(zap.NewExample(), b, &lease.FakeLessor{}, nil)
+		var s *store
+		switch test {
+		case "recreate":
+			s = NewStore(zap.NewExample(), b, &lease.FakeLessor{}, nil)
+		case "restore":
+			s0.Restore(b)
+			s = s0
+		}
 
 		// wait for scheduled compaction to be finished
 		time.Sleep(100 * time.Millisecond)
 
-	if _, err := s1.Range([]byte("foo"), nil, RangeOptions{Rev: 1}); err != ErrCompacted {
+		if _, err := s.Range([]byte("foo"), nil, RangeOptions{Rev: 1}); err != ErrCompacted {
 			t.Errorf("range on compacted rev error = %v, want %v", err, ErrCompacted)
 		}
 		// check the key in backend is deleted
@@ -499,7 +508,7 @@ func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
 		// The disk compaction is done asynchronously and requires more time on slow disk.
 		// try 5 times for CI with slow IO.
 		for i := 0; i < 5; i++ {
-		tx = s1.b.BatchTx()
+			tx = s.b.BatchTx()
 			tx.Lock()
 			ks, _ := tx.UnsafeRange(keyBucketName, revbytes, nil, 0)
 			tx.Unlock()
@@ -512,6 +521,7 @@ func TestRestoreContinueUnfinishedCompaction(t *testing.T) {
 
 		t.Errorf("key for rev %+v still exists, want deleted", bytesToRev(revbytes))
 	}
+}
 
 type hashKVResult struct {
 	hash       uint32