From 604df7088a399a54941dbf941130cedd8c2076f4 Mon Sep 17 00:00:00 2001
From: Marius van der Wijden <m.vanderwijden@live.de>
Date: Fri, 15 Sep 2023 11:40:31 +0200
Subject: [PATCH] common/bloom: implement expiring bloom filters

---
 common/bloom/bloom.go      | 81 ++++++++++++++++++++++++++++++++++++++
 common/bloom/bloom_test.go | 50 +++++++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 common/bloom/bloom.go
 create mode 100644 common/bloom/bloom_test.go

diff --git a/common/bloom/bloom.go b/common/bloom/bloom.go
new file mode 100644
index 0000000000..affa717f1c
--- /dev/null
+++ b/common/bloom/bloom.go
@@ -0,0 +1,81 @@
+package bloom
+
+import (
+	"hash"
+	"sync"
+	"time"
+
+	bloomfilter "github.com/holiman/bloomfilter/v2"
+)
+
+type ExpiringBloom struct {
+	currentBloom int
+	blooms       []*bloomfilter.Filter
+	filterM      uint64
+	filterK      uint64
+
+	timer   *time.Ticker
+	mu      sync.RWMutex // Mutex only locks the currentBloom variable
+	closeCh chan struct{}
+}
+
+func NewExpiringBloom(n, m, k uint64, timeout time.Duration) *ExpiringBloom {
+	blooms := make([]*bloomfilter.Filter, 0, n)
+	for i := 0; i < int(n); i++ {
+		filter, err := bloomfilter.New(m, k)
+		if err != nil {
+			panic(err)
+		}
+		blooms = append(blooms, filter)
+	}
+	filter := ExpiringBloom{
+		currentBloom: 0,
+		blooms:       blooms,
+		filterM:      m,
+		filterK:      k,
+		timer:        time.NewTicker(timeout),
+		closeCh:      make(chan struct{}),
+	}
+	go filter.loop()
+	return &filter
+}
+
+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
+			}
+			e.mu.Unlock()
+		case <-e.closeCh:
+			break
+		}
+	}
+}
+
+func (e *ExpiringBloom) Stop() {
+	close(e.closeCh)
+}
+
+func (e *ExpiringBloom) Put(key hash.Hash64) {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	e.blooms[e.currentBloom].Add(key)
+}
+
+func (e *ExpiringBloom) Contain(key hash.Hash64) bool {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	return e.blooms[e.currentBloom].Contains(key)
+}
diff --git a/common/bloom/bloom_test.go b/common/bloom/bloom_test.go
new file mode 100644
index 0000000000..9536c809f5
--- /dev/null
+++ b/common/bloom/bloom_test.go
@@ -0,0 +1,50 @@
+package bloom
+
+import (
+	"testing"
+	"time"
+)
+
+type hashable struct {
+	b []byte
+}
+
+func (h hashable) BlockSize() int {
+	return len(h.b)
+}
+
+func (h hashable) Hash() []byte {
+	return h.b
+}
+
+func (h hashable) Sum([]byte) []byte {
+	return h.b
+}
+
+func (h hashable) Sum64() uint64 {
+	return 1
+}
+
+func (h hashable) Write([]byte) (int, error) {
+	return 0, nil
+}
+
+func (h hashable) Reset() {}
+
+func (h hashable) Size() int {
+	return len(h.b)
+}
+
+func TestBloom(t *testing.T) {
+	bloom := NewExpiringBloom(2, 10, 10, 10*time.Millisecond)
+
+	testKey := hashable{[]byte{0x01}}
+	bloom.Put(testKey)
+	if !bloom.Contain(testKey) {
+		t.Fail()
+	}
+	time.Sleep(10 * time.Millisecond)
+	if bloom.Contain(testKey) {
+		t.Fail()
+	}
+}