add MarkMany to MarkSet interface

This commit is contained in:
vyzo 2022-01-30 11:27:24 +02:00
parent f9fd47e7d0
commit d140909562
3 changed files with 62 additions and 0 deletions

View File

@ -14,6 +14,7 @@ var errMarkSetClosed = errors.New("markset closed")
type MarkSet interface { type MarkSet interface {
ObjectVisitor ObjectVisitor
Mark(cid.Cid) error Mark(cid.Cid) error
MarkMany([]cid.Cid) error
Has(cid.Cid) (bool, error) Has(cid.Cid) (bool, error)
Close() error Close() error

View File

@ -143,6 +143,23 @@ func (s *BadgerMarkSet) Mark(c cid.Cid) error {
return nil return nil
} }
func (s *BadgerMarkSet) MarkMany(batch []cid.Cid) error {
s.mx.Lock()
if s.pend == nil {
s.mx.Unlock()
return errMarkSetClosed
}
write, seqno := s.putMany(batch)
s.mx.Unlock()
if write {
return s.write(seqno)
}
return nil
}
func (s *BadgerMarkSet) Has(c cid.Cid) (bool, error) { func (s *BadgerMarkSet) Has(c cid.Cid) (bool, error) {
s.mx.RLock() s.mx.RLock()
defer s.mx.RUnlock() defer s.mx.RUnlock()
@ -260,6 +277,24 @@ func (s *BadgerMarkSet) put(key string) (write bool, seqno int) {
return true, seqno return true, seqno
} }
func (s *BadgerMarkSet) putMany(batch []cid.Cid) (write bool, seqno int) {
for _, c := range batch {
key := string(c.Hash())
s.pend[key] = struct{}{}
}
if !s.persist && len(s.pend) < badgerMarkSetBatchSize {
return false, 0
}
seqno = s.seqno
s.seqno++
s.writing[seqno] = s.pend
s.pend = make(map[string]struct{})
return true, seqno
}
func (s *BadgerMarkSet) write(seqno int) (err error) { func (s *BadgerMarkSet) write(seqno int) (err error) {
s.mx.Lock() s.mx.Lock()
if s.pend == nil { if s.pend == nil {

View File

@ -169,6 +169,32 @@ func (s *MapMarkSet) Mark(c cid.Cid) error {
return nil return nil
} }
func (s *MapMarkSet) MarkMany(batch []cid.Cid) error {
s.mx.Lock()
defer s.mx.Unlock()
if s.set == nil {
return errMarkSetClosed
}
for _, c := range batch {
hash := c.Hash()
s.set[string(hash)] = struct{}{}
if s.persist {
if err := s.writeKey(hash, false); err != nil {
return err
}
}
}
if s.persist {
return s.buf.Flush()
}
return nil
}
func (s *MapMarkSet) Has(cid cid.Cid) (bool, error) { func (s *MapMarkSet) Has(cid cid.Cid) (bool, error) {
s.mx.RLock() s.mx.RLock()
defer s.mx.RUnlock() defer s.mx.RUnlock()