diff --git a/common/bloom/bloom.go b/common/bloom/bloom.go index ce55e439b3..853d60d441 100644 --- a/common/bloom/bloom.go +++ b/common/bloom/bloom.go @@ -8,74 +8,79 @@ import ( bloomfilter "github.com/holiman/bloomfilter/v2" ) +const ( + k = 4 +) + type ExpiringBloom struct { currentBloom int union *bloomfilter.Filter blooms []*bloomfilter.Filter - filterM uint64 - filterK uint64 + size uint64 - timer *time.Ticker - mu sync.RWMutex // Mutex only locks the currentBloom variable + timer *time.Ticker + // Mutex lock the currentBloom and union variables + mu sync.RWMutex closeCh chan struct{} } -func NewExpiringBloom(n, m, k uint64, timeout time.Duration) *ExpiringBloom { +func NewExpiringBloom(n, m uint64, timeout time.Duration) (*ExpiringBloom, error) { + union, err := bloomfilter.New(m*8, k) + if err != nil { + return nil, err + } + blooms := make([]*bloomfilter.Filter, 0, n) for i := 0; i < int(n); i++ { - filter, err := bloomfilter.New(m, k) + filter, err := union.NewCompatible() if err != nil { - panic(err) + return nil, err } blooms = append(blooms, filter) } - union, err := bloomfilter.New(m, k) - if err != nil { - panic(err) - } + filter := ExpiringBloom{ currentBloom: 0, blooms: blooms, union: union, - filterM: m, - filterK: k, + size: m * 8, timer: time.NewTicker(timeout), closeCh: make(chan struct{}), } go filter.loop() - return &filter + + return &filter, nil } func (e *ExpiringBloom) loop() { for { select { case <-e.timer.C: - // Reset the filters on every tick - e.mu.Lock() - var err error - e.blooms[e.currentBloom], err = bloomfilter.New(e.filterM, e.filterK) - if err != nil { - panic(err) - } - e.currentBloom++ - if e.currentBloom == len(e.blooms)-1 { - e.currentBloom = 0 - } - // Recreate the union filter - e.union, err = bloomfilter.New(e.filterM, e.filterK) - if err != nil { - panic(err) - } - for _, bloom := range e.blooms { - e.union.UnionInPlace(bloom) - } - e.mu.Unlock() + e.tick() case <-e.closeCh: - break + return } } } +func (e *ExpiringBloom) tick() { + e.mu.Lock() + defer e.mu.Unlock() + + // Advance the current bloom + e.currentBloom++ + if e.currentBloom == len(e.blooms) { + e.currentBloom = 0 + } + // Clear the filter + e.blooms[e.currentBloom], _ = e.blooms[e.currentBloom].NewCompatible() + // Recreate the union filter + e.union, _ = e.union.NewCompatible() + for _, bloom := range e.blooms { + e.union.UnionInPlace(bloom) + } +} + func (e *ExpiringBloom) Stop() { close(e.closeCh) } diff --git a/common/bloom/bloom_test.go b/common/bloom/bloom_test.go index 26d6391e90..a55908990c 100644 --- a/common/bloom/bloom_test.go +++ b/common/bloom/bloom_test.go @@ -36,19 +36,66 @@ func (h hashable) Size() int { } func TestBloom(t *testing.T) { - bloom := NewExpiringBloom(2, 10, 10, 10*time.Millisecond) + bloom, _ := NewExpiringBloom(3, 1024, 10*time.Millisecond) testKey := hashable{[]byte{0x01}} bloom.Put(testKey) if !bloom.Contain(testKey) { t.Fatal() } - time.Sleep(10 * time.Millisecond) + time.Sleep(11 * time.Millisecond) if !bloom.Contain(testKey) { t.Fatal() } - time.Sleep(10 * time.Millisecond) + time.Sleep(11 * time.Millisecond) + if !bloom.Contain(testKey) { + t.Fatal() + } + time.Sleep(11 * time.Millisecond) if bloom.Contain(testKey) { t.Fatal() } } + +func TestBloom2(t *testing.T) { + bloom, _ := NewExpiringBloom(3, 1024, 10*time.Second) + + testKey := hashable{[]byte{0x01}} + // Put key in bloom 0 + bloom.Put(testKey) + if !bloom.Contain(testKey) { + t.Fatal() + } + // Override bloom 1 + bloom.tick() + if !bloom.Contain(testKey) { + t.Fatal() + } + // Override bloom 2 + bloom.tick() + if !bloom.Contain(testKey) { + t.Fatal() + } + // Override bloom 0 + bloom.tick() + if bloom.Contain(testKey) { + t.Fatal() + } +} + +func BenchmarkPut(b *testing.B) { + bloom, _ := NewExpiringBloom(2, 1024, 10*time.Second) + + testKey := hashable{[]byte{0x01}} + for i := 0; i < b.N; i++ { + bloom.Put(testKey) + } +} + +func BenchmarkTick(b *testing.B) { + bloom, _ := NewExpiringBloom(2, 1024, 10*time.Second) + + for i := 0; i < b.N; i++ { + bloom.tick() + } +}