# -*- coding: utf-8 -*-

"""

Created on 2018/12/13 4:11 PM

---------

@summary:

---------

@author: Boris

@email: boris_liu@foxmail.com

"""

import hashlib

import math

import threading

import time

from struct import unpack, pack

from feapder.utils.redis_lock import RedisLock

from . import bitarray

def make_hashfuncs(num_slices, num_bits):

if num_bits >= (1 << 31):

fmt_code, chunk_size = "Q", 8

elif num_bits >= (1 << 15):

fmt_code, chunk_size = "I", 4

else:

fmt_code, chunk_size = "H", 2

total_hash_bits = 8 * num_slices * chunk_size

if total_hash_bits > 384:

hashfn = hashlib.sha512

elif total_hash_bits > 256:

hashfn = hashlib.sha384

elif total_hash_bits > 160:

hashfn = hashlib.sha256

elif total_hash_bits > 128:

hashfn = hashlib.sha1

else:

hashfn = hashlib.md5

fmt = fmt_code * (hashfn().digest_size // chunk_size)

num_salts, extra = divmod(num_slices, len(fmt))

if extra:

num_salts += 1

salts = tuple(hashfn(hashfn(pack("I", i)).digest()) for i in range(num_salts))

def _make_hashfuncs(key):

if isinstance(key, str):

key = key.encode("utf-8")

else:

key = str(key).encode("utf-8")

i = 0

for salt in salts:

h = salt.copy()

h.update(key)

for uint in unpack(fmt, h.digest()):

yield uint % num_bits

i += 1

if i >= num_slices:

return

return _make_hashfuncs

class BloomFilter(object):

BASE_MEMORY = 1

BASE_REDIS = 2

def __init__(

self,

capacity: int,

error_rate: float = 0.00001,

bitarray_type=BASE_REDIS,

name=None,

redis_url=None,

):

if not (0 < error_rate < 1):

raise ValueError("Error_Rate must be between 0 and 1.")

if not capacity > 0:

raise ValueError("Capacity must be > 0")

# given M = num_bits, k = num_slices, P = error_rate, n = capacity

# k = log2(1/P)

# solving for m = bits_per_slice

# n ~= M * ((ln(2) ** 2) / abs(ln(P)))

# n ~= (k * m) * ((ln(2) ** 2) / abs(ln(P)))

# m ~= n * abs(ln(P)) / (k * (ln(2) ** 2))

num_slices = int(math.ceil(math.log(1.0 / error_rate, 2)))

bits_per_slice = int(

math.ceil(

(capacity * abs(math.log(error_rate)))

/ (num_slices * (math.log(2) ** 2))

)

)

self._setup(error_rate, num_slices, bits_per_slice, capacity)

if bitarray_type == BloomFilter.BASE_MEMORY:

self.bitarray = bitarray.MemoryBitArray(self.num_bits)

self.bitarray.setall(False)

elif bitarray_type == BloomFilter.BASE_REDIS:

assert name, "name can't be None "

self.bitarray = bitarray.RedisBitArray(name, redis_url)

else:

raise ValueError("not support this bitarray type")

def _setup(self, error_rate, num_slices, bits_per_slice, capacity):

self.error_rate = error_rate

self.num_slices = num_slices

self.bits_per_slice = bits_per_slice

self.capacity = capacity

self.num_bits = num_slices * bits_per_slice

self.make_hashes = make_hashfuncs(self.num_slices, self.bits_per_slice)

self._is_at_capacity = False

self._check_capacity_time = 0

def __repr__(self):

return "<BloomFilter: {}>".format(self.bitarray)

def get(self, keys, to_list=False):

is_list = isinstance(keys, list)

keys = keys if is_list else [keys]

is_exists = []

offsets = []

for key in keys:

hashes = self.make_hashes(key)

offset = 0

for k in hashes:

offsets.append(offset + k)

offset += self.bits_per_slice

old_values = self.bitarray.get(offsets)

for i in range(0, len(old_values), self.num_slices):

is_exists.append(int(all(old_values[i : i + self.num_slices])))

if to_list:

return is_exists

else:

return is_exists if is_list else is_exists[0]

@property

def is_at_capacity(self):

"""

是否容量已满, 1的个数满位数组的一半的时，则看做已满

比较耗时 半小时检查一次

@return:

"""

# if self._is_at_capacity:

# return self._is_at_capacity

#

# if not self._check_capacity_time or time.time() - self._check_capacity_time > 1800:

# bit_count = self.bitarray.count()

# if bit_count and bit_count / self.num_bits > 0.5:

# self._is_at_capacity = True

#

# self._check_capacity_time = time.time()

#

# return self._is_at_capacity

if self._is_at_capacity:

return self._is_at_capacity

bit_count = self.bitarray.count()

if bit_count and bit_count / self.num_bits > 0.5:

self._is_at_capacity = True

return self._is_at_capacity

def add(self, keys):

"""

Adds a key to this bloom filter. If the key already exists in this

filter it will return False. Otherwise True. keys support list

@param keys: list or one key

@return:

"""

if self.is_at_capacity:

raise IndexError("BloomFilter is at capacity")

is_list = isinstance(keys, list)

keys = keys if is_list else [keys]

is_added = []

offsets = []

for key in keys:

hashes = self.make_hashes(key)

offset = 0

for k in hashes:

offsets.append(offset + k)

offset += self.bits_per_slice

old_values = self.bitarray.set(offsets, 1)

for i in range(0, len(old_values), self.num_slices):

is_added.append(1 ^ int(all(old_values[i : i + self.num_slices])))

return is_added if is_list else is_added[0]

class ScalableBloomFilter(object):

"""

自动扩展空间的bloomfilter, 当一个filter满一半的时候，创建下一个

"""

BASE_MEMORY = BloomFilter.BASE_MEMORY

BASE_REDIS = BloomFilter.BASE_REDIS

def __init__(

self,

initial_capacity: int = 100000000,

error_rate: float = 0.00001,

bitarray_type=BASE_REDIS,

name=None,

redis_url=None,

):

if not error_rate or error_rate < 0:

raise ValueError("Error_Rate must be a decimal less than 0.")

self._setup(

initial_capacity, error_rate, name, bitarray_type, redis_url=redis_url

)

def _setup(self, initial_capacity, error_rate, name, bitarray_type, redis_url):

self.initial_capacity = initial_capacity

self.error_rate = error_rate

self.name = name

self.bitarray_type = bitarray_type

self.redis_url = redis_url

self.filters = []

self.filters.append(self.create_filter())

self._thread_lock = threading.RLock()

self._check_capacity_time = 0

def __repr__(self):

return "<ScalableBloomFilter: {}>".format(self.filters[-1].bitarray)

def create_filter(self):

filter = BloomFilter(

capacity=self.initial_capacity,

error_rate=self.error_rate,

bitarray_type=self.bitarray_type,

name=self.name + str(len(self.filters)) if self.name else self.name,

redis_url=self.redis_url,

)

return filter

def check_filter_capacity(self):

"""

检测filter状态，如果已满，加载新的filter

@return:

"""

if (

not self._check_capacity_time

or time.time() - self._check_capacity_time > 1800

):

if self.bitarray_type == ScalableBloomFilter.BASE_MEMORY:

with self._thread_lock:

while True:

if self.filters[-1].is_at_capacity:

self.filters.append(self.create_filter())

else:

break

self._check_capacity_time = time.time()

else:

# 全局锁 同一时间只有一个进程在真正的创建新的filter，等这个进程创建完，其他进程只是把刚创建的filter append进来

key = (

f"ScalableBloomFilter:{self.name}"

if self.name

else "ScalableBloomFilter"

)

with RedisLock(key=key, redis_url=self.redis_url) as lock:

if lock.locked:

while True:

if self.filters[-1].is_at_capacity:

self.filters.append(self.create_filter())

else:

break

self._check_capacity_time = time.time()

def add(self, keys, skip_check=False):

"""

Adds a key to this bloom filter. If the key already exists in this

filter it will return False. Otherwise True. keys support list

@param keys: list or one key

@param skip_check: add directly，not check if is exist in bloomfilters

@return:

"""

self.check_filter_capacity()

current_filter = self.filters[-1]

if skip_check:

return current_filter.add(keys)

else:

is_list = isinstance(keys, list)

keys = keys if is_list else [keys]

not_exist_keys = list(set(keys))

# 检查之前的bloomfilter是否存在

# 记录下每级filter存在的key，不存在的key继续向下检查

for filter in reversed(self.filters):

current_filter_is_exists = filter.get(

not_exist_keys, to_list=True

) # 当前的filter是否存在

not_exist_keys_temp = []

for key, is_exist in zip(not_exist_keys, current_filter_is_exists):

if not is_exist: # 当前filter不存在的key 需要继续向下检查

not_exist_keys_temp.append(key)

not_exist_keys = not_exist_keys_temp

if not not_exist_keys:

break

# 仍有不存在的关键词，记录该关键词

if not_exist_keys:

current_filter.add(not_exist_keys)

# 比较key是否已存在, 内部重复的key 若不存在啊则只留其一算为不存在，其他看作已存在

for i, key in enumerate(keys):

for j, not_exist_key in enumerate(not_exist_keys):

if key == not_exist_key:

keys[i] = 1

not_exist_keys.pop(j)

break

else:

keys[i] = 0

is_added = keys

return is_added if is_list else is_added[0]

def get(self, keys):

self.check_filter_capacity()

is_list = isinstance(keys, list)

keys = keys if is_list else [keys] # 最终会修改为 [0, 1, ...] 0表示不存在 1 已存在

not_exist_keys = list(set(keys))

# 检查之前的bloomfilter是否存在

# 记录下每级filter存在的key，不存在的key继续向下检查

for filter in reversed(self.filters):

current_filter_is_exists = filter.get(

not_exist_keys, to_list=True

) # 当前的filter是否存在

not_exist_keys_temp = []

for checked_key, is_exist in zip(not_exist_keys, current_filter_is_exists):

if not is_exist: # 当前filter不存在的key 需要继续向下检查

not_exist_keys_temp.append(checked_key)

not_exist_keys = not_exist_keys_temp

if not not_exist_keys:

break

# 比较key是否已存在, 内部重复的key 若不存在啊则只留其一算为不存在，其他看作已存在

for i, key in enumerate(keys):

for j, not_exist_key in enumerate(not_exist_keys):

if key == not_exist_key:

keys[i] = 0

not_exist_keys.pop(j)

break

else:

keys[i] = 1

is_exists = keys

return is_exists if is_list else is_exists[0]

@property

def capacity(self):

"""Returns the total capacity for all filters in this SBF"""

return sum(f.capacity for f in self.filters)