github.com/shuguocloud/go-zero@v1.3.0/core/bloom/bloom.go (about)

     1  package bloom
     2  
     3  import (
     4  	"errors"
     5  	"strconv"
     6  
     7  	"github.com/shuguocloud/go-zero/core/hash"
     8  	"github.com/shuguocloud/go-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  	if !isSet {
    73  		return false, nil
    74  	}
    75  
    76  	return true, nil
    77  }
    78  
    79  func (f *Filter) getLocations(data []byte) []uint {
    80  	locations := make([]uint, maps)
    81  	for i := uint(0); i < maps; i++ {
    82  		hashValue := hash.Hash(append(data, byte(i)))
    83  		locations[i] = uint(hashValue % uint64(f.bits))
    84  	}
    85  
    86  	return locations
    87  }
    88  
    89  type redisBitSet struct {
    90  	store *redis.Redis
    91  	key   string
    92  	bits  uint
    93  }
    94  
    95  func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
    96  	return &redisBitSet{
    97  		store: store,
    98  		key:   key,
    99  		bits:  bits,
   100  	}
   101  }
   102  
   103  func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
   104  	var args []string
   105  
   106  	for _, offset := range offsets {
   107  		if offset >= r.bits {
   108  			return nil, ErrTooLargeOffset
   109  		}
   110  
   111  		args = append(args, strconv.FormatUint(uint64(offset), 10))
   112  	}
   113  
   114  	return args, nil
   115  }
   116  
   117  func (r *redisBitSet) check(offsets []uint) (bool, error) {
   118  	args, err := r.buildOffsetArgs(offsets)
   119  	if err != nil {
   120  		return false, err
   121  	}
   122  
   123  	resp, err := r.store.Eval(testScript, []string{r.key}, args)
   124  	if err == redis.Nil {
   125  		return false, nil
   126  	} else if err != nil {
   127  		return false, err
   128  	}
   129  
   130  	exists, ok := resp.(int64)
   131  	if !ok {
   132  		return false, nil
   133  	}
   134  
   135  	return exists == 1, nil
   136  }
   137  
   138  func (r *redisBitSet) del() error {
   139  	_, err := r.store.Del(r.key)
   140  	return err
   141  }
   142  
   143  func (r *redisBitSet) expire(seconds int) error {
   144  	return r.store.Expire(r.key, seconds)
   145  }
   146  
   147  func (r *redisBitSet) set(offsets []uint) error {
   148  	args, err := r.buildOffsetArgs(offsets)
   149  	if err != nil {
   150  		return err
   151  	}
   152  
   153  	_, err = r.store.Eval(setScript, []string{r.key}, args)
   154  	if err == redis.Nil {
   155  		return nil
   156  	}
   157  
   158  	return err
   159  }