github.com/lingyao2333/mo-zero@v1.4.1/core/bloom/bloom.go (about)

     1  package bloom
     2  
     3  import (
     4  	"errors"
     5  	"strconv"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/hash"
     8  	"github.com/lingyao2333/mo-zero/core/stores/redis"
     9  )
    10  
    11  const (
    12  	// for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
    13  	// maps as k in the error rate table
    14  	maps      = 14
    15  	setScript = `
    16  for _, offset in ipairs(ARGV) do
    17  	redis.call("setbit", KEYS[1], offset, 1)
    18  end
    19  `
    20  	testScript = `
    21  for _, offset in ipairs(ARGV) do
    22  	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
    23  		return false
    24  	end
    25  end
    26  return true
    27  `
    28  )
    29  
    30  // ErrTooLargeOffset indicates the offset is too large in bitset.
    31  var ErrTooLargeOffset = errors.New("too large offset")
    32  
    33  type (
    34  	// A Filter is a bloom filter.
    35  	Filter struct {
    36  		bits   uint
    37  		bitSet bitSetProvider
    38  	}
    39  
    40  	bitSetProvider interface {
    41  		check([]uint) (bool, error)
    42  		set([]uint) error
    43  	}
    44  )
    45  
    46  // New create a Filter, store is the backed redis, key is the key for the bloom filter,
    47  // bits is how many bits will be used, maps is how many hashes for each addition.
    48  // best practices:
    49  // elements - means how many actual elements
    50  // when maps = 14, formula: 0.7*(bits/maps), bits = 20*elements, the error rate is 0.000067 < 1e-4
    51  // for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
    52  func New(store *redis.Redis, key string, bits uint) *Filter {
    53  	return &Filter{
    54  		bits:   bits,
    55  		bitSet: newRedisBitSet(store, key, bits),
    56  	}
    57  }
    58  
    59  // Add adds data into f.
    60  func (f *Filter) Add(data []byte) error {
    61  	locations := f.getLocations(data)
    62  	return f.bitSet.set(locations)
    63  }
    64  
    65  // Exists checks if data is in f.
    66  func (f *Filter) Exists(data []byte) (bool, error) {
    67  	locations := f.getLocations(data)
    68  	isSet, err := f.bitSet.check(locations)
    69  	if err != nil {
    70  		return false, err
    71  	}
    72  
    73  	return isSet, nil
    74  }
    75  
    76  func (f *Filter) getLocations(data []byte) []uint {
    77  	locations := make([]uint, maps)
    78  	for i := uint(0); i < maps; i++ {
    79  		hashValue := hash.Hash(append(data, byte(i)))
    80  		locations[i] = uint(hashValue % uint64(f.bits))
    81  	}
    82  
    83  	return locations
    84  }
    85  
    86  type redisBitSet struct {
    87  	store *redis.Redis
    88  	key   string
    89  	bits  uint
    90  }
    91  
    92  func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
    93  	return &redisBitSet{
    94  		store: store,
    95  		key:   key,
    96  		bits:  bits,
    97  	}
    98  }
    99  
   100  func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
   101  	var args []string
   102  
   103  	for _, offset := range offsets {
   104  		if offset >= r.bits {
   105  			return nil, ErrTooLargeOffset
   106  		}
   107  
   108  		args = append(args, strconv.FormatUint(uint64(offset), 10))
   109  	}
   110  
   111  	return args, nil
   112  }
   113  
   114  func (r *redisBitSet) check(offsets []uint) (bool, error) {
   115  	args, err := r.buildOffsetArgs(offsets)
   116  	if err != nil {
   117  		return false, err
   118  	}
   119  
   120  	resp, err := r.store.Eval(testScript, []string{r.key}, args)
   121  	if err == redis.Nil {
   122  		return false, nil
   123  	} else if err != nil {
   124  		return false, err
   125  	}
   126  
   127  	exists, ok := resp.(int64)
   128  	if !ok {
   129  		return false, nil
   130  	}
   131  
   132  	return exists == 1, nil
   133  }
   134  
   135  func (r *redisBitSet) del() error {
   136  	_, err := r.store.Del(r.key)
   137  	return err
   138  }
   139  
   140  func (r *redisBitSet) expire(seconds int) error {
   141  	return r.store.Expire(r.key, seconds)
   142  }
   143  
   144  func (r *redisBitSet) set(offsets []uint) error {
   145  	args, err := r.buildOffsetArgs(offsets)
   146  	if err != nil {
   147  		return err
   148  	}
   149  
   150  	_, err = r.store.Eval(setScript, []string{r.key}, args)
   151  	if err == redis.Nil {
   152  		return nil
   153  	}
   154  
   155  	return err
   156  }