88import asyncio
99import logging
1010from abc import ABC , abstractmethod
11- from collections import defaultdict
11+ from collections import OrderedDict , defaultdict
12+ from time import monotonic
1213from typing import TYPE_CHECKING , Any , Literal
1314
1415import asyncssh
1516import httpcore
16- from aiocache import Cache
17- from aiocache .plugins import HitMissRatioPlugin
1817from asyncssh import SSHClientConnection , SSHClientConnectionOptions
1918from httpx import ConnectError , HTTPError , TimeoutException
2019
3433CLIENT_KEYS = asyncssh .public_key .load_default_keypairs ()
3534
3635
36+ class AntaCache :
37+ """Class to be used as cache.
38+
39+ Example
40+ -------
41+
42+ ```python
43+ # Create cache
44+ cache = AntaCache("device1")
45+ with cache.locks[key]:
46+ command_output = cache.get(key)
47+ ```
48+ """
49+
50+ def __init__ (self , device : str , max_size : int = 128 , ttl : int = 60 ) -> None :
51+ """Initialize the cache."""
52+ self .device = device
53+ self .cache : OrderedDict [str , Any ] = OrderedDict ()
54+ self .locks : defaultdict [str , asyncio .Lock ] = defaultdict (asyncio .Lock )
55+ self .max_size = max_size
56+ self .ttl = ttl
57+
58+ # Stats
59+ self .stats : dict [str , int ] = {}
60+ self ._init_stats ()
61+
62+ def _init_stats (self ) -> None :
63+ """Initialize the stats."""
64+ self .stats ["hits" ] = 0
65+ self .stats ["total" ] = 0
66+
67+ async def get (self , key : str ) -> Any : # noqa: ANN401
68+ """Return the cached entry for key."""
69+ self .stats ["total" ] += 1
70+ if key in self .cache :
71+ timestamp , value = self .cache [key ]
72+ if monotonic () - timestamp < self .ttl :
73+ # checking the value is still valid
74+ self .cache .move_to_end (key )
75+ self .stats ["hits" ] += 1
76+ return value
77+ # Time expired
78+ del self .cache [key ]
79+ del self .locks [key ]
80+ return None
81+
82+ async def set (self , key : str , value : Any ) -> bool : # noqa: ANN401
83+ """Set the cached entry for key to value."""
84+ timestamp = monotonic ()
85+ if len (self .cache ) > self .max_size :
86+ self .cache .popitem (last = False )
87+ self .cache [key ] = timestamp , value
88+ return True
89+
90+ def clear (self ) -> None :
91+ """Empty the cache."""
92+ logger .debug ("Clearing cache for device %s" , self .device )
93+ self .cache = OrderedDict ()
94+ self ._init_stats ()
95+
96+
3797class AntaDevice (ABC ):
3898 """Abstract class representing a device in ANTA.
3999
@@ -52,10 +112,11 @@ class AntaDevice(ABC):
52112 Hardware model of the device.
53113 tags : set[str]
54114 Tags for this device.
55- cache : Cache | None
56- In-memory cache from aiocache library for this device (None if cache is disabled).
115+ cache : AntaCache | None
116+ In-memory cache for this device (None if cache is disabled).
57117 cache_locks : dict
58118 Dictionary mapping keys to asyncio locks to guarantee exclusive access to the cache if not disabled.
119+ Deprecated, will be removed in ANTA v2.0.0, use self.cache.locks instead.
59120
60121 """
61122
@@ -79,7 +140,8 @@ def __init__(self, name: str, tags: set[str] | None = None, *, disable_cache: bo
79140 self .tags .add (self .name )
80141 self .is_online : bool = False
81142 self .established : bool = False
82- self .cache : Cache | None = None
143+ self .cache : AntaCache | None = None
144+ # Keeping cache_locks for backward compatibility.
83145 self .cache_locks : defaultdict [str , asyncio .Lock ] | None = None
84146
85147 # Initialize cache if not disabled
@@ -101,17 +163,16 @@ def __hash__(self) -> int:
101163
102164 def _init_cache (self ) -> None :
103165 """Initialize cache for the device, can be overridden by subclasses to manipulate how it works."""
104- self .cache = Cache ( cache_class = Cache . MEMORY , ttl = 60 , namespace = self .name , plugins = [ HitMissRatioPlugin ()] )
105- self .cache_locks = defaultdict ( asyncio . Lock )
166+ self .cache = AntaCache ( device = self .name , ttl = 60 )
167+ self .cache_locks = self . cache . locks
106168
107169 @property
108170 def cache_statistics (self ) -> dict [str , Any ] | None :
109171 """Return the device cache statistics for logging purposes."""
110- # Need to ignore pylint no-member as Cache is a proxy class and pylint is not smart enough
111- # https://github.com/pylint-dev/pylint/issues/7258
112172 if self .cache is not None :
113- stats = getattr (self .cache , "hit_miss_ratio" , {"total" : 0 , "hits" : 0 , "hit_ratio" : 0 })
114- return {"total_commands_sent" : stats ["total" ], "cache_hits" : stats ["hits" ], "cache_hit_ratio" : f"{ stats ['hit_ratio' ] * 100 :.2f} }
173+ stats = self .cache .stats
174+ ratio = stats ["hits" ] / stats ["total" ] if stats ["total" ] > 0 else 0
175+ return {"total_commands_sent" : stats ["total" ], "cache_hits" : stats ["hits" ], "cache_hit_ratio" : f"{ ratio * 100 :.2f} }
115176 return None
116177
117178 def __rich_repr__ (self ) -> Iterator [tuple [str , Any ]]:
@@ -177,18 +238,16 @@ async def collect(self, command: AntaCommand, *, collection_id: str | None = Non
177238 collection_id
178239 An identifier used to build the eAPI request ID.
179240 """
180- # Need to ignore pylint no-member as Cache is a proxy class and pylint is not smart enough
181- # https://github.com/pylint-dev/pylint/issues/7258
182- if self .cache is not None and self .cache_locks is not None and command .use_cache :
183- async with self .cache_locks [command .uid ]:
184- cached_output = await self .cache .get (command .uid ) # pylint: disable=no-member
241+ if self .cache is not None and command .use_cache :
242+ async with self .cache .locks [command .uid ]:
243+ cached_output = await self .cache .get (command .uid )
185244
186245 if cached_output is not None :
187246 logger .debug ("Cache hit for %s on %s" , command .command , self .name )
188247 command .output = cached_output
189248 else :
190249 await self ._collect (command = command , collection_id = collection_id )
191- await self .cache .set (command .uid , command .output ) # pylint: disable=no-member
250+ await self .cache .set (command .uid , command .output )
192251 else :
193252 await self ._collect (command = command , collection_id = collection_id )
194253
0 commit comments