From 6c5e68dded30d54b7a23f6c24f3b092856065cb3 Mon Sep 17 00:00:00 2001 From: Daniel M Date: Tue, 30 Jul 2024 18:15:41 -0400 Subject: [PATCH] fix:move random seed to HeavyKeeper to avoid issues Fix #315 --- fakeredis/stack/_topk_mixin.py | 78 ++++++---------------------------- 1 file changed, 12 insertions(+), 66 deletions(-) diff --git a/fakeredis/stack/_topk_mixin.py b/fakeredis/stack/_topk_mixin.py index 4d7be858..68190c2c 100644 --- a/fakeredis/stack/_topk_mixin.py +++ b/fakeredis/stack/_topk_mixin.py @@ -10,11 +10,8 @@ from fakeredis._commands import Key, Int, Float, command, CommandItem from fakeredis._helpers import OK, SimpleError, SimpleString -random.seed(time.time()) - class Bucket(object): - def __init__(self, counter: int, fingerprint: int): self.counter = counter self.fingerprint = fingerprint @@ -36,7 +33,7 @@ def count(self, fingerprint: int) -> int: def _decay(self, decay: float) -> bool: if self.counter > 0: - probability = decay**self.counter + probability = decay ** self.counter if probability >= 1 or random.random() < probability: self.counter -= 1 return self.counter == 0 @@ -64,7 +61,11 @@ def _hash(self, item: bytes) -> int: class HeavyKeeper(object): + is_topk_initialized = False + def __init__(self, k: int, width: int = 1024, depth: int = 5, decay: float = 0.9) -> None: + if not HeavyKeeper.is_topk_initialized: + random.seed(time.time()) self.k = k self.width = width self.depth = depth @@ -115,15 +116,7 @@ class TopkCommandsMixin: def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - @command( - name="TOPK.ADD", - fixed=( - Key(HeavyKeeper), - bytes, - ), - repeat=(bytes,), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.ADD", fixed=(Key(HeavyKeeper), bytes), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE) def topk_add(self, key: CommandItem, *args: bytes) -> List[Optional[bytes]]: if key.value is None: raise SimpleError("TOPK: key does not exist") @@ -133,15 +126,7 @@ def topk_add(self, key: CommandItem, *args: bytes) -> List[Optional[bytes]]: key.updated() return res - @command( - name="TOPK.COUNT", - fixed=( - Key(HeavyKeeper), - bytes, - ), - repeat=(bytes,), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.COUNT", fixed=(Key(HeavyKeeper), bytes), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE) def topk_count(self, key: CommandItem, *args: bytes) -> List[int]: if key.value is None: raise SimpleError("TOPK: key does not exist") @@ -150,15 +135,7 @@ def topk_count(self, key: CommandItem, *args: bytes) -> List[int]: res: List[int] = [key.value.count(_item) for _item in args] return res - @command( - name="TOPK.QUERY", - fixed=( - Key(HeavyKeeper), - bytes, - ), - repeat=(bytes,), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.QUERY", fixed=(Key(HeavyKeeper), bytes), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE) def topk_query(self, key: CommandItem, *args: bytes) -> List[int]: if key.value is None: raise SimpleError("TOPK: key does not exist") @@ -168,16 +145,7 @@ def topk_query(self, key: CommandItem, *args: bytes) -> List[int]: res: List[int] = [1 if _item in topk else 0 for _item in args] return res - @command( - name="TOPK.INCRBY", - fixed=( - Key(), - bytes, - Int, - ), - repeat=(bytes, Int), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.INCRBY", fixed=(Key(), bytes, Int), repeat=(bytes, Int), flags=msgs.FLAG_DO_NOT_CREATE) def topk_incrby(self, key: CommandItem, *args: Any) -> List[Optional[bytes]]: if key.value is None: raise SimpleError("TOPK: key does not exist") @@ -192,12 +160,7 @@ def topk_incrby(self, key: CommandItem, *args: Any) -> List[Optional[bytes]]: key.updated() return res - @command( - name="TOPK.INFO", - fixed=(Key(),), - repeat=(), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.INFO", fixed=(Key(),), repeat=(), flags=msgs.FLAG_DO_NOT_CREATE) def topk_info(self, key: CommandItem) -> List[Any]: if key.value is None: raise SimpleError("TOPK: key does not exist") @@ -214,12 +177,7 @@ def topk_info(self, key: CommandItem) -> List[Any]: key.value.decay, ] - @command( - name="TOPK.LIST", - fixed=(Key(),), - repeat=(bytes,), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.LIST", fixed=(Key(),), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE) def topk_list(self, key: CommandItem, *args: Any) -> List[Any]: (withcount,), _ = extract_args(args, ("withcount",)) if key.value is None: @@ -233,19 +191,7 @@ def topk_list(self, key: CommandItem, *args: Any) -> List[Any]: temp = [[item[1], item[0]] for item in value_list] return [item for sublist in temp for item in sublist] - @command( - name="TOPK.RESERVE", - fixed=( - Key(), - Int, - ), - repeat=( - Int, - Int, - Float, - ), - flags=msgs.FLAG_DO_NOT_CREATE, - ) + @command(name="TOPK.RESERVE", fixed=(Key(), Int), repeat=(Int, Int, Float), flags=msgs.FLAG_DO_NOT_CREATE) def topk_reserve(self, key: CommandItem, topk: int, *args: Any) -> SimpleString: if len(args) == 3: width, depth, decay = args