lotus/blockstore/splitstore/markset_bloom.go
2021-07-08 10:20:29 +03:00

106 lines
1.8 KiB
Go

package splitstore
import (
"crypto/rand"
"crypto/sha256"
"sync"
"golang.org/x/xerrors"
bbloom "github.com/ipfs/bbloom"
cid "github.com/ipfs/go-cid"
)
const (
BloomFilterMinSize = 10_000_000
BloomFilterProbability = 0.01
)
type BloomMarkSetEnv struct {
ts bool
}
var _ MarkSetEnv = (*BloomMarkSetEnv)(nil)
type BloomMarkSet struct {
salt []byte
mx sync.RWMutex
bf *bbloom.Bloom
ts bool
}
var _ MarkSet = (*BloomMarkSet)(nil)
func NewBloomMarkSetEnv(ts bool) (*BloomMarkSetEnv, error) {
return &BloomMarkSetEnv{ts: ts}, nil
}
func (e *BloomMarkSetEnv) Create(name string, sizeHint int64) (MarkSet, error) {
size := int64(BloomFilterMinSize)
for size < sizeHint {
size += BloomFilterMinSize
}
salt := make([]byte, 4)
_, err := rand.Read(salt)
if err != nil {
return nil, xerrors.Errorf("error reading salt: %w", err)
}
bf, err := bbloom.New(float64(size), BloomFilterProbability)
if err != nil {
return nil, xerrors.Errorf("error creating bloom filter: %w", err)
}
return &BloomMarkSet{salt: salt, bf: bf, ts: e.ts}, nil
}
func (e *BloomMarkSetEnv) Close() error {
return nil
}
func (s *BloomMarkSet) saltedKey(cid cid.Cid) []byte {
hash := cid.Hash()
key := make([]byte, len(s.salt)+len(hash))
n := copy(key, s.salt)
copy(key[n:], hash)
rehash := sha256.Sum256(key)
return rehash[:]
}
func (s *BloomMarkSet) Mark(cid cid.Cid) error {
if s.ts {
s.mx.Lock()
defer s.mx.Unlock()
}
if s.bf == nil {
return errMarkSetClosed
}
s.bf.Add(s.saltedKey(cid))
return nil
}
func (s *BloomMarkSet) Has(cid cid.Cid) (bool, error) {
if s.ts {
s.mx.RLock()
defer s.mx.RUnlock()
}
if s.bf == nil {
return false, errMarkSetClosed
}
return s.bf.Has(s.saltedKey(cid)), nil
}
func (s *BloomMarkSet) Close() error {
if s.ts {
s.mx.Lock()
defer s.mx.Unlock()
}
s.bf = nil
return nil
}