unify marksets and visitors

blockstore/splitstore/markset.go

@@ -24,6 +24,7 @@ type MarkSet interface {
 type MarkSetVisitor interface {
 	ObjectVisitor
 	Close() error
+	SetConcurrent()
 }
 
 type MarkSetEnv interface {

blockstore/splitstore/markset_badger.go

@@ -33,15 +33,7 @@ type BadgerMarkSet struct {
 }
 
 var _ MarkSet = (*BadgerMarkSet)(nil)
-
+var _ MarkSetVisitor = (*BadgerMarkSet)(nil)
-type BadgerMarkSetVisitor struct {
-	pend map[string]struct{}
-
-	db   *badger.DB
-	path string
-}
-
-var _ MarkSetVisitor = (*BadgerMarkSetVisitor)(nil)
 
 var badgerMarkSetBatchSize = 16384
 
@@ -55,7 +47,7 @@ func NewBadgerMarkSetEnv(path string) (MarkSetEnv, error) {
 	return &BadgerMarkSetEnv{path: msPath}, nil
 }
 
-func (e *BadgerMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error) {
+func (e *BadgerMarkSetEnv) create(name string, sizeHint int64) (*BadgerMarkSet, error) {
 	path := filepath.Join(e.path, name)
 
 	db, err := openTransientBadgerDB(path)
@@ -74,21 +66,12 @@ func (e *BadgerMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error)
 	return ms, nil
 }
 
+func (e *BadgerMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error) {
+	return e.create(name, sizeHint)
+}
+
 func (e *BadgerMarkSetEnv) CreateVisitor(name string, sizeHint int64) (MarkSetVisitor, error) {
-	path := filepath.Join(e.path, name)
+	return e.create(name, sizeHint)
-
-	db, err := openTransientBadgerDB(path)
-	if err != nil {
-		return nil, xerrors.Errorf("error creating badger db: %w", err)
-	}
-
-	v := &BadgerMarkSetVisitor{
-		pend: make(map[string]struct{}),
-		db:   db,
-		path: path,
-	}
-
-	return v, nil
 }
 
 func (e *BadgerMarkSetEnv) Close() error {
@@ -110,42 +93,7 @@ func (s *BadgerMarkSet) Mark(c cid.Cid) error {
 		return nil
 	}
 
-	pend := s.pend
+	return s.putPending()
-	seqno := s.seqno
-	s.seqno++
-	s.writing[seqno] = pend
-	s.pend = make(map[string]struct{})
-	s.writers++
-	s.mx.Unlock()
-
-	defer func() {
-		s.mx.Lock()
-		defer s.mx.Unlock()
-
-		delete(s.writing, seqno)
-		s.writers--
-		if s.writers == 0 {
-			s.cond.Broadcast()
-		}
-	}()
-
-	empty := []byte{} // not nil
-
-	batch := s.db.NewWriteBatch()
-	defer batch.Cancel()
-
-	for k := range pend {
-		if err := batch.Set([]byte(k), empty); err != nil {
-			return err
-		}
-	}
-
-	err := batch.Flush()
-	if err != nil {
-		return xerrors.Errorf("error flushing batch to badger markset: %w", err)
-	}
-
-	return nil
 }
 
 func (s *BadgerMarkSet) Has(c cid.Cid) (bool, error) {
@@ -187,6 +135,98 @@ func (s *BadgerMarkSet) Has(c cid.Cid) (bool, error) {
 	}
 }
 
+func (s *BadgerMarkSet) Visit(c cid.Cid) (bool, error) {
+	s.mx.Lock()
+
+	if s.pend == nil {
+		s.mx.Unlock()
+		return false, errMarkSetClosed
+	}
+
+	key := c.Hash()
+	pendKey := string(key)
+	if _, ok := s.pend[pendKey]; ok {
+		s.mx.Unlock()
+		return false, nil
+	}
+
+	for _, wr := range s.writing {
+		if _, ok := wr[pendKey]; ok {
+			s.mx.Unlock()
+			return false, nil
+		}
+	}
+
+	err := s.db.View(func(txn *badger.Txn) error {
+		_, err := txn.Get(key)
+		return err
+	})
+
+	switch err {
+	case nil:
+		s.mx.Unlock()
+		return false, nil
+
+	case badger.ErrKeyNotFound:
+		s.pend[pendKey] = struct{}{}
+
+		if len(s.pend) < badgerMarkSetBatchSize {
+			s.mx.Unlock()
+			return true, nil
+		}
+
+		if err := s.putPending(); err != nil {
+			return false, err
+		}
+
+		return true, nil
+
+	default:
+		s.mx.Unlock()
+		return false, xerrors.Errorf("error checking badger markset: %w", err)
+	}
+}
+
+// writer holds the lock
+func (s *BadgerMarkSet) putPending() error {
+	pend := s.pend
+	seqno := s.seqno
+	s.seqno++
+	s.writing[seqno] = pend
+	s.pend = make(map[string]struct{})
+	s.writers++
+	s.mx.Unlock()
+
+	defer func() {
+		s.mx.Lock()
+		defer s.mx.Unlock()
+
+		delete(s.writing, seqno)
+		s.writers--
+		if s.writers == 0 {
+			s.cond.Broadcast()
+		}
+	}()
+
+	empty := []byte{} // not nil
+
+	batch := s.db.NewWriteBatch()
+	defer batch.Cancel()
+
+	for k := range pend {
+		if err := batch.Set([]byte(k), empty); err != nil {
+			return xerrors.Errorf("error setting batch: %w", err)
+		}
+	}
+
+	err := batch.Flush()
+	if err != nil {
+		return xerrors.Errorf("error flushing batch to badger markset: %w", err)
+	}
+
+	return nil
+}
+
 func (s *BadgerMarkSet) Close() error {
 	s.mx.Lock()
 	defer s.mx.Unlock()
@@ -208,72 +248,6 @@ func (s *BadgerMarkSet) Close() error {
 
 func (s *BadgerMarkSet) SetConcurrent() {}
 
-func (v *BadgerMarkSetVisitor) Visit(c cid.Cid) (bool, error) {
-	if v.pend == nil {
-		return false, errMarkSetClosed
-	}
-
-	key := c.Hash()
-	pendKey := string(key)
-	_, ok := v.pend[pendKey]
-	if ok {
-		return false, nil
-	}
-
-	err := v.db.View(func(txn *badger.Txn) error {
-		_, err := txn.Get(key)
-		return err
-	})
-
-	switch err {
-	case nil:
-		return false, nil
-
-	case badger.ErrKeyNotFound:
-		v.pend[pendKey] = struct{}{}
-
-		if len(v.pend) < badgerMarkSetBatchSize {
-			return true, nil
-		}
-
-		pend := v.pend
-		v.pend = make(map[string]struct{})
-
-		empty := []byte{} // not nil
-
-		batch := v.db.NewWriteBatch()
-		defer batch.Cancel()
-
-		for k := range pend {
-			if err := batch.Set([]byte(k), empty); err != nil {
-				return false, err
-			}
-		}
-
-		err := batch.Flush()
-		if err != nil {
-			return false, xerrors.Errorf("error flushing batch to badger markset: %w", err)
-		}
-
-		return true, nil
-
-	default:
-		return false, xerrors.Errorf("error checking badger markset: %w", err)
-	}
-}
-
-func (v *BadgerMarkSetVisitor) Close() error {
-	if v.pend == nil {
-		return nil
-	}
-
-	v.pend = nil
-	db := v.db
-	v.db = nil
-
-	return closeTransientBadgerDB(db, v.path)
-}
-
 func openTransientBadgerDB(path string) (*badger.DB, error) {
 	// clean up first
 	err := os.RemoveAll(path)

blockstore/splitstore/markset_map.go

@@ -18,27 +18,24 @@ type MapMarkSet struct {
 }
 
 var _ MarkSet = (*MapMarkSet)(nil)
-
+var _ MarkSetVisitor = (*MapMarkSet)(nil)
-type MapMarkSetVisitor struct {
-	set map[string]struct{}
-}
-
-var _ MarkSetVisitor = (*MapMarkSetVisitor)(nil)
 
 func NewMapMarkSetEnv() (*MapMarkSetEnv, error) {
 	return &MapMarkSetEnv{}, nil
 }
 
-func (e *MapMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error) {
+func (e *MapMarkSetEnv) create(name string, sizeHint int64) (*MapMarkSet, error) {
 	return &MapMarkSet{
 		set: make(map[string]struct{}, sizeHint),
 	}, nil
 }
 
+func (e *MapMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error) {
+	return e.create(name, sizeHint)
+}
+
 func (e *MapMarkSetEnv) CreateVisitor(name string, sizeHint int64) (MarkSetVisitor, error) {
-	return &MapMarkSetVisitor{
+	return e.create(name, sizeHint)
-		set: make(map[string]struct{}, sizeHint),
-	}, nil
 }
 
 func (e *MapMarkSetEnv) Close() error {
@@ -73,6 +70,25 @@ func (s *MapMarkSet) Has(cid cid.Cid) (bool, error) {
 	return ok, nil
 }
 
+func (s *MapMarkSet) Visit(c cid.Cid) (bool, error) {
+	if s.ts {
+		s.mx.Lock()
+		defer s.mx.Unlock()
+	}
+
+	if s.set == nil {
+		return false, errMarkSetClosed
+	}
+
+	key := string(c.Hash())
+	if _, ok := s.set[key]; ok {
+		return false, nil
+	}
+
+	s.set[key] = struct{}{}
+	return true, nil
+}
+
 func (s *MapMarkSet) Close() error {
 	if s.ts {
 		s.mx.Lock()
@@ -85,22 +101,3 @@ func (s *MapMarkSet) Close() error {
 func (s *MapMarkSet) SetConcurrent() {
 	s.ts = true
 }
-
-func (v *MapMarkSetVisitor) Visit(c cid.Cid) (bool, error) {
-	if v.set == nil {
-		return false, errMarkSetClosed
-	}
-
-	key := string(c.Hash())
-	if _, ok := v.set[key]; ok {
-		return false, nil
-	}
-
-	v.set[key] = struct{}{}
-	return true, nil
-}
-
-func (v *MapMarkSetVisitor) Close() error {
-	v.set = nil
-	return nil
-}