布隆过滤器的简单实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
package utils

import (
	"hash"
	"math"
	"sync"

	"github.com/bits-and-blooms/bitset"
	"github.com/twmb/murmur3"
)

// BloomFilter 结构体
type BloomFilter struct {
	bitset    *bitset.BitSet // 使用第三方bitset包
	size      uint64         // 位集合大小
	hashFuncs []hash.Hash64  // 哈希函数集合
	mutex     sync.RWMutex   // 读写锁
	// hashMutex sync.Mutex     // 添加专门的哈希函数互斥锁
}

// NewWithFalsePositiveRate 根据预期元素数量和误判率创建布隆过滤器
func NewWithFalsePositiveRate(expectedItems uint64, falsePositiveRate float64) *BloomFilter {
	// 计算最优位数组大小和哈希函数数量
	m, k := optimalParameters(expectedItems, falsePositiveRate)
	funcs := make([]hash.Hash64, k)
	for i := 0; i < k; i++ {
		funcs[i] = murmur3.SeedNew64(uint64(i))
	}
	return &BloomFilter{
		bitset:    bitset.New(uint(m)),
		size:      m,
		hashFuncs: funcs,
		mutex:     sync.RWMutex{},
	}
}

// optimalParameters 计算最优参数
func optimalParameters(n uint64, p float64) (uint64, int) {
	m := uint64(math.Ceil(-float64(n) * math.Log(p) / (math.Ln2 * math.Ln2)))
	k := int(math.Ceil(math.Ln2 * float64(m) / float64(n)))
	return m, k
}

// createHashFunctions 创建哈希函数集合
// func createHashFunctions(k int) []hash.Hash64 {
// 	funcs := make([]hash.Hash64, k)
// 	for i := 0; i < k; i++ {
// 		funcs[i] = murmur3.SeedNew64(uint64(i))
// 	}
// 	return funcs
// }

// Add 添加元素到布隆过滤器
func (bf *BloomFilter) Add(item []byte) {
	bf.mutex.Lock()
	defer bf.mutex.Unlock()

	// bf.hashMutex.Lock()
	// defer bf.hashMutex.Unlock()

	for _, h := range bf.hashFuncs {
		h.Reset()
		_, _ = h.Write(item)
		index := h.Sum64() % bf.size
		bf.bitset.Set(uint(index))
	}
}

// Contains 检查元素是否可能在布隆过滤器中
func (bf *BloomFilter) Contains(item []byte) bool {
	bf.mutex.RLock()
	defer bf.mutex.RUnlock()

	// bf.hashMutex.Lock()
	// defer bf.hashMutex.Unlock()

	for _, h := range bf.hashFuncs {
		h.Reset()
		_, _ = h.Write(item)
		index := h.Sum64() % bf.size
		if !bf.bitset.Test(uint(index)) {
			return false
		}
	}
	return true
}

Reference

图解各类限流算法｜固定窗口/计数器、滑动窗口、漏桶算法、令牌桶算法

bloom-filter.go