add Put method

Author: dsa0x

bloom.go

@@ -45,6 +45,9 @@ type BloomFilter struct {
 	seeds []int64
 
 	path string
+
+	// candidates cache
+	cdCache map[string][]uint64
 }
 
 type BloomOptions struct {
@@ -63,8 +66,14 @@ type BloomOptions struct {
 
 	// growth rate of the bloom filter (valid values are 2 and 4)
 	GrowthRate int
+}
 
-	Scalable bool
+var DefaultBloomOptions = BloomOptions{
+	Path:       "bloom.db",
+	Err_rate:   0.001,
+	Capacity:   100000,
+	GrowthRate: 2,
+	Database:   nil,
 }
 
 // NewBloom creates a new bloom filter.
@@ -74,6 +83,9 @@ type BloomOptions struct {
 //
 // database is the persistent store to attach to the filter. can be nil.
 func NewBloom(opts *BloomOptions) *BloomFilter {
+	if opts == nil {
+		opts = &DefaultBloomOptions
+	}
 	if opts.Err_rate <= 0 || opts.Err_rate >= 1 {
 		panic("Error rate must be between 0 and 1")
 	}
@@ -140,7 +152,7 @@ func NewBloom(opts *BloomOptions) *BloomFilter {
 }
 
 // Add adds the key to the bloom filter
-func (bf *BloomFilter) Add(key, val []byte) {
+func (bf *BloomFilter) Add(key []byte) {
 	bf.lock.Lock()
 	defer bf.lock.Unlock()
 	defer func() {
@@ -165,10 +177,16 @@ func (bf *BloomFilter) Add(key, val []byte) {
 	}
 	bf.count++
 
-	if bf.hasStore() {
-		bf.db.Put([]byte(key), val)
+}
+
+// Put adds the key to the bloom filter, and also stores it in the persistent store
+func (bf *BloomFilter) Put(key, val []byte) error {
+	if !bf.hasStore() {
+		fmt.Errorf("BloomFilter does not have a store, use Add() to add keys")
 	}
 
+	bf.Add(key)
+	return bf.db.Put([]byte(key), val)
 }
 
 // Contains checks if the key exists in the bloom filter
@@ -330,6 +348,32 @@ func (bf *BloomFilter) DB() interface{} {
 	return bf.db.DB()
 }
 
+// Clear resets all bits in the bloom filter
+func (bf *BloomFilter) Clear() {
+	bf.lock.Lock()
+	defer bf.lock.Unlock()
+	bf.mem = make([]byte, bf.bit_width)
+}
+
+type BloomFilterStats struct {
+	Capacity int
+	Count    int
+	Size     int
+	M        int
+	K        int
+}
+
+// Stats returns the stats of the bloom filter
+func (bf *BloomFilter) Stats() BloomFilterStats {
+	return BloomFilterStats{
+		Capacity: bf.capacity,
+		Count:    bf.count,
+		Size:     bf.bit_width,
+		M:        bf.m,
+		K:        bf.k,
+	}
+}
+
 // divmod returns the quotient and remainder of a/b
 func divmod(num, denom int64) (quot, rem int64) {
 	quot = num / denom

bloom_test.go

@@ -44,8 +44,8 @@ func TestBloomFilter_Add(t *testing.T) {
 	}()
 
 	t.Run("success", func(t *testing.T) {
-		key, val := []byte("foo"), []byte("var")
-		bf.Add(key, val)
+		key := []byte("foo")
+		bf.Add(key)
 	})
 
 	t.Run("count should sum up to the number of entries added", func(t *testing.T) {
@@ -65,7 +65,7 @@ func TestBloomFilter_Add(t *testing.T) {
 		for i := 0; i < count; i++ {
 			var by [4]byte
 			binary.LittleEndian.PutUint32(by[:], uint32(i))
-			bf.Add(by[:], []byte("bar"))
+			bf.Add(by[:])
 		}
 		if bf.Count() != count {
 			t.Errorf("Expected count to be %d, got %d", bf.Count(), count)
@@ -95,9 +95,9 @@ func TestBloomFilter_Add(t *testing.T) {
 		for i := 0; i < count; i++ {
 			var by [4]byte
 			binary.LittleEndian.PutUint32(by[:], uint32(i))
-			bf.Add(by[:], []byte("bar"))
+			bf.Add(by[:])
 		}
-		bf.Add([]byte("test"), []byte("bar"))
+		bf.Add([]byte("test"))
 		t.Errorf("Expected function to panic when number of entries exceed the capacity")
 	})
 
@@ -121,7 +121,7 @@ func TestBloomFilter_Add(t *testing.T) {
 			os.Remove(opts.Path)
 		}()
 
-		bf.Add([]byte("foo"), []byte("bar"))
+		bf.Put([]byte("foo"), []byte("bar"))
 		val := bf.Get([]byte("foo"))
 		t.Errorf("Expected function to panic when there is no persistent store, got %s", val)
 	})
@@ -169,10 +169,10 @@ func TestBloomFilter_Merge(t *testing.T) {
 	})
 
 	t.Run("object added to the single filters should be found in the resulting merge", func(t *testing.T) {
-		key, val := []byte("foo"), []byte("bar")
+		key := []byte("foo")
 		bf := NewBloom(opts)
 		bf2 := NewBloom(opts)
-		bf2.Add(key, val)
+		bf2.Add(key)
 		err := bf.Merge(bf2)
 		if err != nil {
 			t.Errorf("Expected no error, got %v", err)
@@ -202,15 +202,15 @@ func TestBloomFilter_AddToDB(t *testing.T) {
 
 	t.Run("success", func(t *testing.T) {
 		key, val := []byte("foo"), []byte("var")
-		bf.Add(key, val)
+		bf.Put(key, val)
 
 		if val, err := bf.db.Get([]byte(key)); err != nil || val == nil {
 			t.Errorf("bf.cache[%s] not found; error: %v", key, err)
 		}
 	})
 	t.Run("should not find key that was not added", func(t *testing.T) {
 		key, val := []byte("foo"), []byte("var")
-		bf.Add(key, val)
+		bf.Put(key, val)
 
 		if val, err := bf.db.Get([]byte("bar")); err != nil || val != nil {
 			t.Errorf("expected value to be nil, got %s; error: %v", val, err)
@@ -234,15 +234,15 @@ func TestBloomFilter_AddToBadgerDB(t *testing.T) {
 
 	t.Run("success", func(t *testing.T) {
 		key, val := []byte("foo"), []byte("var")
-		bf.Add(key, val)
+		bf.Put(key, val)
 
 		if val, err := bf.db.Get([]byte(key)); err != nil || val == nil {
 			t.Errorf("bf.cache[%s] not found; error: %v", key, err)
 		}
 	})
 	t.Run("should not find key that was not added", func(t *testing.T) {
 		key, val := []byte("foo"), []byte("var")
-		bf.Add(key, val)
+		bf.Put(key, val)
 
 		if val, err := bf.db.Get([]byte("bar")); err != nil || val != nil {
 			t.Errorf("expected value to be nil, got %s; error: %v", val, err)
@@ -266,10 +266,12 @@ func TestBloomFilter(t *testing.T) {
 		os.Remove(opts.Path)
 	}()
 
-	bf.Add([]byte("foo"), []byte("bar"))
-	bf.Add([]byte("baz"), []byte("qux"))
+	bf.Add([]byte("foo"))
+	bf.Add([]byte("baz"))
 
 	t.Run("key may be in cache if found in bloom, def not in cache if not found", func(t *testing.T) {
+		bf.Put([]byte("foo"), []byte("bar"))
+		bf.Put([]byte("baz"), []byte("qux"))
 		table := []struct {
 			key      string
 			expected bool

cmd/main.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"log"
-	"os"
 	"runtime"
 	"time"
 
@@ -12,9 +11,9 @@ import (
 )
 
 func main() {
-	num := 20_000_000
+	num := 20_000_00
 	// div := num / 10
-	// main2(num)
+	// main5(num)
 	// return
 	opts := &sprout.BloomOptions{
 		Err_rate: 0.001,
@@ -23,22 +22,21 @@ func main() {
 	}
 	bf := sprout.NewBloom(opts)
 	defer bf.Close()
-	// return
+
+	// reset filter
+	bf.Clear()
+
 	start := time.Now()
-	bf.Add([]byte("foo"), []byte("bar"))
+	bf.Add([]byte("foo"))
 
-	for i := 0; i < num-1; i++ {
-		bf.Add([]byte(fmt.Sprintf("%d", i)), []byte("bar"))
-		// if i%div == 0 {
-		// 	time.Sleep(time.Second * 3)
-		// 	fmt.Println(i, "added")
-		// }
+	for i := 0; i < num-2; i++ {
+		bf.Add([]byte(fmt.Sprintf("%d", i)))
 		// fmt.Println(i+1, bf.Contains([]byte(fmt.Sprintf("%d", i+1))))
 	}
 	fmt.Println(bf.Contains([]byte("foo")))
 	fmt.Println(bf.Contains([]byte("bar")))
+	fmt.Printf("%+v\n", bf.Stats())
 	fmt.Printf("Added %d elements in %v\n", bf.Capacity(), time.Since(start))
-	PrintMemUsage()
 }
 
 // Using sprout with a persistent storage
@@ -69,62 +67,13 @@ func main2(num int) {
 	bf := sprout.NewScalableBloom(opts)
 	start := time.Now()
 	for i := 0; i < num*10; i++ {
-		bf.Add([]byte{byte(i)}, []byte("bar"))
+		bf.Add([]byte{byte(i)})
 	}
-	bf.Add([]byte("foo"), []byte("bar"))
+	bf.Add([]byte("foo"))
 	fmt.Println(bf.Contains([]byte("foo")))
 	fmt.Println("Added", num*10, "elements in", time.Since(start))
 }
 
-func main4(num int) {
-	db, err := bolt.Open("store.db", 0600, nil)
-	if err != nil {
-		panic(err)
-	}
-
-	err = db.Update(func(tx *bolt.Tx) error {
-		_, err := tx.CreateBucketIfNotExists([]byte("store.name"))
-		return err
-	})
-	if err != nil {
-		panic(err)
-	}
-
-	w, err := os.OpenFile("storebolt.db", os.O_RDWR|os.O_CREATE, 0600)
-	if err != nil {
-		panic(err)
-	}
-
-	// defer os.Remove("storebolt.db")
-
-	start := time.Now()
-	tx, err := db.Begin(true)
-	if err != nil {
-		panic(err)
-	}
-	defer tx.Rollback()
-	size := tx.Size()
-
-	b := tx.Bucket([]byte("store.name"))
-
-	for i := 0; i < num; i++ {
-		b.Put([]byte{byte(i)}, []byte("bar"))
-	}
-
-	// write snapshot to pipe
-	go func() {
-		defer w.Close()
-		_, err := tx.WriteTo(w)
-		if err != nil {
-			log.Println("Erroring writing to pipe", err)
-		}
-	}()
-	if err != nil {
-		panic(err)
-	}
-	fmt.Println("Added", num, "elements in", time.Since(start), "bytes=", size)
-}
-
 func PrintMemUsage() {
 	var m runtime.MemStats
 	runtime.ReadMemStats(&m)
@@ -137,3 +86,35 @@ func PrintMemUsage() {
 func bToMb(b uint64) uint64 {
 	return b / 1024 / 1024
 }
+
+func main6() {
+	num := 2_000_000
+	db, err := bolt.Open("store1.db", 0644, nil)
+	if err != nil {
+		panic(err)
+	}
+	err = db.Update(func(tx *bolt.Tx) error {
+		_, err := tx.CreateBucketIfNotExists([]byte("test"))
+		return err
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		fmt.Printf("Starting %d: ", i)
+		err = db.Update(func(tx *bolt.Tx) error {
+			b := tx.Bucket([]byte("test"))
+			for j := 0; j < num/10; j++ {
+				err := b.Put([]byte(fmt.Sprintf("foo-i%d-j%d", i, j)), []byte("bar"))
+				if err != nil {
+					return err
+				}
+			}
+			fmt.Printf("%+v\n", b.Stats().KeyN)
+			return nil
+		})
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+}

cmd/main_test.go

@@ -37,7 +37,7 @@ func Benchmark_NewBloom(b *testing.B) {
 	bf := sprout.NewBloom(opts)
 	n := 0
 	for i := 0; i < b.N; i++ {
-		bf.Add([]byte{byte(n)}, []byte("bar"))
+		bf.Add([]byte{byte(n)})
 		n++
 	}
 
@@ -59,7 +59,7 @@ func Benchmark_NewBloomFind(b *testing.B) {
 
 	n := 0
 	for i := 0; i < b.N; i++ {
-		bf.Add([]byte{byte(n)}, []byte("bar"))
+		bf.Add([]byte{byte(n)})
 		n++
 	}
 	n = 0