feat: close meta store (#114)

- fix: unmarshal typo
- feat: close meta store
- test: add ut for meta store

Author: ifplusor

internal/store/meta/async.go

@@ -35,6 +35,7 @@ type AsyncStore struct {
 	pending *skiplist.SkipList
 
 	commitc chan struct{}
+	closec  chan struct{}
 	donec   chan struct{}
 }
 
@@ -49,6 +50,7 @@ func newAsyncStore(wal *walog.WAL, committed *skiplist.SkipList, version, snapsh
 		},
 		pending: skiplist.New(skiplist.Bytes),
 		commitc: make(chan struct{}, 1),
+		closec:  make(chan struct{}),
 		donec:   make(chan struct{}),
 	}
 
@@ -57,10 +59,16 @@ func newAsyncStore(wal *walog.WAL, committed *skiplist.SkipList, version, snapsh
 	return s
 }
 
-func (s *AsyncStore) Stop() {
-	// TODO(james.yin): stop WAL
-	close(s.commitc)
+func (s *AsyncStore) Close() {
+	s.mu.Lock()
+	close(s.closec)
+	s.mu.Unlock()
+
 	<-s.donec
+
+	// Close WAL.
+	s.wal.Close()
+	s.wal.Wait()
 }
 
 func (s *AsyncStore) Load(key []byte) (interface{}, bool) {
@@ -76,17 +84,38 @@ func (s *AsyncStore) Load(key []byte) (interface{}, bool) {
 }
 
 func (s *AsyncStore) Store(key []byte, value interface{}) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.pending.Set(key, value)
-	s.tryCommit()
+	_ = s.set(KVRange(key, value))
+}
+
+func (s *AsyncStore) BatchStore(kvs Ranger) {
+	_ = s.set(kvs)
 }
 
 func (s *AsyncStore) Delete(key []byte) {
+	_ = s.set(KVRange(key, deletedMark))
+}
+
+func (s *AsyncStore) set(kvs Ranger) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.pending.Set(key, deletedMark)
+
+	select {
+	case <-s.closec:
+		return ErrClosed
+	default:
+	}
+
+	err := kvs.Range(func(key []byte, value interface{}) error {
+		s.pending.Set(key, value)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
 	s.tryCommit()
+
+	return nil
 }
 
 func (s *AsyncStore) tryCommit() {
@@ -113,10 +142,9 @@ func (s *AsyncStore) runCommit() {
 
 	for {
 		select {
-		case _, ok := <-s.commitc:
-			if !ok {
-				return
-			}
+		case <-s.closec:
+			return
+		case <-s.commitc:
 		case <-ticker.C:
 		}
 		s.commit()
@@ -163,7 +191,7 @@ func RecoverAsyncStore(walDir string) (*AsyncStore, error) {
 	version := snapshot
 	wal, err := walog.RecoverWithVisitor(walDir, snapshot, func(data []byte, offset int64) error {
 		m := skiplist.New(skiplist.Bytes)
-		err2 := defaultCodec.Unmarshl(data, func(key []byte, value interface{}) error {
+		err2 := defaultCodec.Unmarshal(data, func(key []byte, value interface{}) error {
 			m.Set(key, value)
 			return nil
 		})

internal/store/meta/async_test.go

@@ -14,8 +14,87 @@
 
 package meta
 
-import "testing"
+import (
+	// standard libraries.
+	"os"
+	"testing"
 
-func TestAsyncStore_Load(t *testing.T) {
+	// third-party libraries.
+	. "github.com/smartystreets/goconvey/convey"
+)
 
+func TestAsyncStore(t *testing.T) {
+	Convey("AsyncStore", t, func() {
+		walDir, err := os.MkdirTemp("", "async-*")
+		So(err, ShouldBeNil)
+
+		Convey("new empty AsyncStore by recovery", func() {
+			ss, err := RecoverAsyncStore(walDir)
+
+			So(err, ShouldBeNil)
+			So(ss, ShouldNotBeNil)
+
+			ss.Close()
+		})
+
+		Convey("setup AsyncStore", func() {
+			ss, err := RecoverAsyncStore(walDir)
+			So(err, ShouldBeNil)
+			ss.Store(key0, "value0")
+			ss.Store(key1, "value1")
+			ss.Close()
+
+			Convey("recover AsyncStore", func() {
+				ss, err = RecoverAsyncStore(walDir)
+				So(err, ShouldBeNil)
+
+				value0, ok0 := ss.Load(key0)
+				So(ok0, ShouldBeTrue)
+				So(value0, ShouldResemble, "value0")
+
+				value1, ok1 := ss.Load(key1)
+				So(ok1, ShouldBeTrue)
+				So(value1, ShouldResemble, "value1")
+
+				_, ok2 := ss.Load(key2)
+				So(ok2, ShouldBeFalse)
+
+				Convey("modify AsyncStore", func() {
+					ss.Delete(key1)
+					_, ok1 = ss.Load(key1)
+					So(ok1, ShouldBeFalse)
+
+					ss.Store(key2, "value2")
+					value2, ok2 := ss.Load(key2)
+					So(ok2, ShouldBeTrue)
+					So(value2, ShouldResemble, "value2")
+
+					ss.Close()
+
+					Convey("recover AsyncStore again", func() {
+						ss, err = RecoverAsyncStore(walDir)
+						So(err, ShouldBeNil)
+
+						value0, ok0 := ss.Load(key0)
+						So(ok0, ShouldBeTrue)
+						So(value0, ShouldResemble, "value0")
+
+						_, ok1 := ss.Load(key1)
+						So(ok1, ShouldBeFalse)
+
+						value2, ok2 := ss.Load(key2)
+						So(ok2, ShouldBeTrue)
+						So(value2, ShouldResemble, "value2")
+
+						ss.Close()
+					})
+				})
+			})
+		})
+
+		Reset(func() {
+			err := os.RemoveAll(walDir)
+			So(err, ShouldBeNil)
+		})
+	})
 }

internal/store/meta/codec.go

@@ -34,7 +34,7 @@ type Marshaler interface {
 }
 
 type Unmarshaler interface {
-	Unmarshl(data []byte, cb RangeCallback) error
+	Unmarshal(data []byte, cb RangeCallback) error
 }
 
 type deletedMarkType struct{}
@@ -176,7 +176,7 @@ func toKind(k reflect.Kind) Kind {
 	return Invalid
 }
 
-func (codec) Unmarshl(data []byte, cb RangeCallback) error {
+func (codec) Unmarshal(data []byte, cb RangeCallback) error {
 	var last []byte
 	for so := 0; so < len(data); {
 		// Decode key.

internal/store/meta/snapshot.go

@@ -97,7 +97,7 @@ func recoverLatestSnopshot(dir string, unmarshaler Unmarshaler) (*skiplist.SkipL
 	}
 
 	m := skiplist.New(skiplist.Bytes)
-	err = unmarshaler.Unmarshl(data, func(key []byte, value interface{}) error {
+	err = unmarshaler.Unmarshal(data, func(key []byte, value interface{}) error {
 		m.Set(key, value)
 		return nil
 	})

internal/store/meta/store.go

@@ -16,6 +16,7 @@ package meta
 
 import (
 	// standard libraries.
+	"errors"
 	"sync"
 
 	// third-party libraries.
@@ -25,6 +26,8 @@ import (
 	walog "github.com/linkall-labs/vanus/internal/store/wal"
 )
 
+var ErrClosed = errors.New("MetaStore: closed")
+
 type store struct {
 	mu sync.RWMutex
 

internal/store/meta/sync.go

@@ -16,6 +16,7 @@ package meta
 
 import (
 	// standard libraries.
+	"errors"
 	"time"
 
 	// third-party libraries.
@@ -33,6 +34,7 @@ type SyncStore struct {
 	store
 
 	snapshotc chan struct{}
+	donec     chan struct{}
 }
 
 func newSyncStore(wal *walog.WAL, committed *skiplist.SkipList, version, snapshot int64) *SyncStore {
@@ -45,15 +47,23 @@ func newSyncStore(wal *walog.WAL, committed *skiplist.SkipList, version, snapsho
 			marshaler: defaultCodec,
 		},
 		snapshotc: make(chan struct{}, 1),
+		donec:     make(chan struct{}),
 	}
 
 	go s.runSnapshot()
 
 	return s
 }
 
-func (s *SyncStore) Stop() {
-	// TODO(james.yin): stop WAL
+func (s *SyncStore) Close() {
+	// Close WAL.
+	s.wal.Close()
+	s.wal.Wait()
+
+	// NOTE: Can not close the snapshotc before close the WAL,
+	// because write to snapshotc in callback of WAL append.
+	close(s.snapshotc)
+	<-s.donec
 }
 
 func (s *SyncStore) Load(key []byte) (interface{}, bool) {
@@ -87,27 +97,25 @@ func (s *SyncStore) set(kvs Ranger) error {
 	}
 
 	ch := make(chan error, 1)
+	// Use callbacks for ordering guarantees.
 	s.wal.AppendOne(entry, walog.WithCallback(func(re walog.Result) {
 		if re.Err != nil {
 			ch <- re.Err
 			return
 		}
 
 		// Update state.
-		func() {
-			s.mu.Lock()
-			defer s.mu.Unlock()
-			_ = kvs.Range(func(key []byte, value interface{}) error {
-				if value == deletedMark {
-					s.committed.Remove(key)
-				} else {
-					s.committed.Set(key, value)
-				}
-				return nil
-			})
-
-			s.version = re.Offset()
-		}()
+		s.mu.Lock()
+		_ = kvs.Range(func(key []byte, value interface{}) error {
+			if value == deletedMark {
+				s.committed.Remove(key)
+			} else {
+				s.committed.Set(key, value)
+			}
+			return nil
+		})
+		s.version = re.Offset()
+		s.mu.Unlock()
 
 		close(ch)
 
@@ -116,12 +124,22 @@ func (s *SyncStore) set(kvs Ranger) error {
 		default:
 		}
 	}))
-	return <-ch
+	err = <-ch
+
+	// Convert ErrClosed.
+	if err != nil && errors.Is(err, walog.ErrClosed) {
+		return ErrClosed
+	}
+
+	return err
 }
 
 func (s *SyncStore) runSnapshot() {
 	ticker := time.NewTicker(runSnapshotInterval)
-	defer ticker.Stop()
+	defer func() {
+		ticker.Stop()
+		close(s.donec)
+	}()
 
 	for {
 		select {
@@ -143,7 +161,7 @@ func RecoverSyncStore(walDir string) (*SyncStore, error) {
 
 	version := snapshot
 	wal, err := walog.RecoverWithVisitor(walDir, snapshot, func(data []byte, offset int64) error {
-		err2 := defaultCodec.Unmarshl(data, func(key []byte, value interface{}) error {
+		err2 := defaultCodec.Unmarshal(data, func(key []byte, value interface{}) error {
 			set(committed, key, value)
 			return nil
 		})