-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #33 from air-go/develop
add bloom filter
- Loading branch information
Showing
7 changed files
with
502 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| package bloom | ||
|
|
||
| import ( | ||
| "context" | ||
| "time" | ||
| ) | ||
|
|
||
| type Bloom interface { | ||
| Add(ctx context.Context, key string, data []byte, ttl time.Duration) error | ||
| // Check is return whether it exists or not | ||
| Check(ctx context.Context, key string, data []byte, ttl time.Duration) (bool, error) | ||
| // CheckAndAdd is return whether it exists, if not exists add. | ||
| CheckAndAdd(ctx context.Context, key string, data []byte, ttl time.Duration) (bool, error) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,162 @@ | ||
| package localbloom | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
| "time" | ||
|
|
||
| "github.com/benbjohnson/clock" | ||
| "github.com/bits-and-blooms/bloom/v3" | ||
| "github.com/why444216978/go-util/assert" | ||
| "github.com/why444216978/go-util/nopanic" | ||
|
|
||
| lbloom "github.com/air-go/rpc/library/bloom" | ||
| ) | ||
|
|
||
| type options struct { | ||
| n uint | ||
| fp float64 | ||
| clock clock.Clock | ||
| releaseDuration time.Duration | ||
| } | ||
|
|
||
| func defaultOptions() *options { | ||
| return &options{ | ||
| n: 10000, | ||
| fp: 0.01, | ||
| releaseDuration: time.Minute, | ||
| } | ||
| } | ||
|
|
||
| type OptionFunc func(*options) | ||
|
|
||
| func SetEstimateParameters(n uint, fp float64) OptionFunc { | ||
| return func(o *options) { | ||
| o.n = n | ||
| o.fp = fp | ||
| } | ||
| } | ||
|
|
||
| func SetClock(c clock.Clock) OptionFunc { | ||
| return func(o *options) { o.clock = c } | ||
| } | ||
|
|
||
| func SetReleaseDuration(d time.Duration) OptionFunc { | ||
| return func(o *options) { o.releaseDuration = d } | ||
| } | ||
|
|
||
| type MemoryBloom struct { | ||
| *options | ||
| blooms sync.Map | ||
| } | ||
|
|
||
| var _ lbloom.Bloom = (*MemoryBloom)(nil) | ||
|
|
||
| func NewMemoryBloom(opts ...OptionFunc) *MemoryBloom { | ||
| opt := defaultOptions() | ||
| for _, o := range opts { | ||
| o(opt) | ||
| } | ||
|
|
||
| mb := &MemoryBloom{ | ||
| options: opt, | ||
| blooms: sync.Map{}, | ||
| } | ||
| mb.tryRelease() | ||
|
|
||
| return mb | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) Add(ctx context.Context, key string, data []byte, ttl time.Duration) error { | ||
| mb.getBloom(key).add(data, mb.now(), ttl) | ||
| return nil | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) Check(ctx context.Context, key string, data []byte, ttl time.Duration) (bool, error) { | ||
| return mb.getBloom(key).check(data, mb.now(), ttl), nil | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) CheckAndAdd(ctx context.Context, key string, data []byte, ttl time.Duration) (bool, error) { | ||
| return mb.getBloom(key).checkAndAdd(data, mb.now(), ttl), nil | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) getBloom(k string) *keyBloom { | ||
| v, ok := mb.blooms.Load(k) | ||
| if ok { | ||
| return v.(*keyBloom) | ||
| } | ||
|
|
||
| b := newKeyBloom(k, mb.n, mb.fp) | ||
| mb.blooms.Store(k, b) | ||
| return b | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) tryRelease() { | ||
| go nopanic.GoVoid(context.Background(), func() { | ||
| t := time.NewTicker(mb.releaseDuration) | ||
| for range t.C { | ||
| mb.blooms.Range(func(k, v any) bool { | ||
| if mb.now().After(v.(*keyBloom).getExpireAt()) { | ||
| mb.blooms.Delete(k) | ||
| } | ||
| return true | ||
| }) | ||
| } | ||
| }) | ||
| } | ||
|
|
||
| func (mb *MemoryBloom) now() time.Time { | ||
| if assert.IsNil(mb.clock) { | ||
| return time.Now() | ||
| } | ||
| return mb.clock.Now() | ||
| } | ||
|
|
||
| type keyBloom struct { | ||
| mu sync.Mutex | ||
| k string | ||
| n uint | ||
| fp float64 | ||
| b *bloom.BloomFilter | ||
| expireAt time.Time | ||
| } | ||
|
|
||
| func newKeyBloom(k string, n uint, fp float64) *keyBloom { | ||
| return &keyBloom{ | ||
| k: k, | ||
| n: n, | ||
| fp: fp, | ||
| b: bloom.NewWithEstimates(n, fp), | ||
| } | ||
| } | ||
|
|
||
| func (b *keyBloom) add(data []byte, now time.Time, ttl time.Duration) { | ||
| b.mu.Lock() | ||
| defer b.mu.Unlock() | ||
|
|
||
| b.expireAt = now.Add(ttl) | ||
| b.b.Add(data) | ||
| } | ||
|
|
||
| func (b *keyBloom) check(data []byte, now time.Time, ttl time.Duration) bool { | ||
| b.mu.Lock() | ||
| defer b.mu.Unlock() | ||
|
|
||
| b.expireAt = now.Add(ttl) | ||
| return b.b.Test(data) | ||
| } | ||
|
|
||
| func (b *keyBloom) checkAndAdd(data []byte, now time.Time, ttl time.Duration) bool { | ||
| b.mu.Lock() | ||
| defer b.mu.Unlock() | ||
|
|
||
| b.expireAt = now.Add(ttl) | ||
| return b.b.TestAndAdd(data) | ||
| } | ||
|
|
||
| func (b *keyBloom) getExpireAt() time.Time { | ||
| b.mu.Lock() | ||
| t := b.expireAt | ||
| b.mu.Unlock() | ||
| return t | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| package localbloom | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/benbjohnson/clock" | ||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestMemoryBloom(t *testing.T) { | ||
| ctx := context.Background() | ||
|
|
||
| c := clock.NewMock() | ||
| c.Set(time.Now()) | ||
|
|
||
| b := NewMemoryBloom( | ||
| SetEstimateParameters(10000, 0.01), | ||
| SetClock(c), | ||
| SetReleaseDuration(time.Nanosecond), | ||
| ) | ||
| key := "key" | ||
|
|
||
| ok, err := b.Check(ctx, key, []byte("abc"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, false, ok) | ||
|
|
||
| err = b.Add(ctx, key, []byte("abc"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
|
|
||
| ok, err = b.Check(ctx, key, []byte("abc"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, true, ok) | ||
|
|
||
| ok, err = b.CheckAndAdd(ctx, key, []byte("abcd"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, false, ok) | ||
|
|
||
| ok, err = b.Check(ctx, key, []byte("abcd"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, true, ok) | ||
|
|
||
| c.Add(time.Minute * 2) | ||
| ok, err = b.Check(ctx, key, []byte("abc"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, false, ok) | ||
| ok, err = b.Check(ctx, key, []byte("abcd"), time.Millisecond) | ||
| assert.Nil(t, err) | ||
| assert.Equal(t, false, ok) | ||
| } |
Oops, something went wrong.