async def write_combinations(redis: aioredis.Redis):
keys = await redis.mget(*ALL_KEYS)
# do nothing, 60sec is enough to pass all tests
if set(keys) != {None}:
return
pipe = redis.pipeline()
for i, comb in enumerate(ALL_COMBINATIONS):
key_ = mk_key(comb)
if i % 3 == 0:
bundle = {
random.choice(COMB_PARTS): random.choice(COMB_PARTS)
for k in range(len(COMB_PARTS))
}
pipe.hmset_dict(key_, **bundle)
elif i % 5 == 0:
bundle = {
random.choice(COMB_PARTS): random.randint(1, 10)
for k in range(len(COMB_PARTS))
}
pipe.zadd(key_,
*chain(*[reversed(pair) for pair in bundle.items()]))
elif i % 7 == 0:
pipe.sadd(key_, *COMB_PARTS)
elif i % 11 == 0:
pipe.lpush(key_, *COMB_PARTS)
else:
pipe.set(key_, key_)
pipe.expire(key_, 60)
return await pipe.execute()
async def add_word_count(wordCountDict: AddWordCountReq,
redis: Redis = Depends(depends_redis)):
pipe = redis.pipeline()
for key, count in wordCountDict.wordcount.items():
pipe.hincrby(WORD_COUNT_KEY, key, count)
result = await pipe.execute()
resp = {"result": result}
fastapi_logger.debug(resp)
return resp
async def load_hosts(redis: Redis) -> List[str]:
""" Load hosts from DB
"""
keys = await redis.keys('host:*')
pipe = redis.pipeline()
for key in keys:
pipe.get(key, encoding='utf-8')
return await pipe.execute()
async def save_hosts(redis: Redis,
hosts: List[str],
expire: Optional[int] = 300) -> None:
""" Save hosts to DB
"""
pipe = redis.pipeline()
for host in hosts:
pipe.set(f'host:{host}', host, expire=expire)
await pipe.execute()
async def increment_counter(redis_pool, name, count=1, now=None):
try:
now = now or time.time()
redis = Redis(redis_pool)
pipe = redis.pipeline()
for prec in PRECISION:
pnow = int(now / prec) * prec
hash = '%s:%s' % (prec, name)
pipe.zadd('known:', 0, hash)
pipe.hincrby('count:' + hash, pnow, count)
await pipe.execute()
except Exception as e:
logger.info("Error saving stats", e)
async def clean_counters(redis_pool):
with await redis_pool as connection:
redis = Redis(connection)
pipe = redis.pipeline()
while True:
start = time.time()
index = 0
known_count = await redis.zcard('known:')
while index < known_count:
hash = await redis.zrange('known:', index, index)
index += 1
if not hash:
break
hash = hash[0].decode("utf-8")
prec = int(hash.partition(':')[0])
hkey = 'count:' + hash
cutoff = time.time() - SAMPLE_COUNT * prec
sample_keys = await redis.hkeys(hkey)
samples = list(map(int, sample_keys))
samples.sort()
remove = bisect.bisect_right(samples, cutoff)
if remove:
print(hkey, samples[:remove])
await redis.hdel(hkey, *samples[:remove])
if remove == len(samples):
try:
pipe.watch(hkey)
length = pipe.hlen(hkey)
if not length:
pipe.multi()
pipe.zrem('known:', hash)
await pipe.execute()
index -= 1
else:
pipe.unwatch()
except redis.exceptions.WatchError:
pass
duration = min(int(time.time() - start) + 1, 60)
await asyncio.sleep(max(60 - duration, 1))
def unwatch_pipe(redis: Redis) -> Redis:
pipe = redis.pipeline()
pipe.unwatch()
return pipe
