Source code for ding.utils.data.structure.cache
import time
from queue import Queue
from threading import Thread
from typing import Any
from ding.utils import LockContext, LockContextType
[docs]class Cache:
"""
Overview:
Data cache for reducing concurrent pressure, with timeout and full queue eject mechanism
Interfaces:
``__init__``, ``push_data``, ``get_cached_data_iter``, ``run``, ``close``
Property:
remain_data_count
"""
[docs] def __init__(self, maxlen: int, timeout: float, monitor_interval: float = 1.0, _debug: bool = False) -> None:
"""
Overview:
Initialize the cache object.
Arguments:
- maxlen (:obj:`int`): Maximum length of the cache queue.
- timeout (:obj:`float`): Maximum second of the data can remain in the cache.
- monitor_interval (:obj:`float`): Interval of the timeout monitor thread checks the time.
- _debug (:obj:`bool`): Whether to use debug mode or not, which enables debug print info.
"""
assert maxlen > 0
self.maxlen = maxlen
self.timeout = timeout
self.monitor_interval = monitor_interval
self.debug = _debug
# two separate receive and send queue for reducing interaction frequency and interference
self.receive_queue = Queue(maxlen)
self.send_queue = Queue(maxlen)
self.receive_lock = LockContext(lock_type=LockContextType.THREAD_LOCK)
self._timeout_thread = Thread(target=self._timeout_monitor)
# the bool flag for gracefully shutting down the timeout monitor thread
self._timeout_thread_flag = True
[docs] def push_data(self, data: Any) -> None:
"""
Overview:
Push data into receive queue, if the receive queue is full(after push), then push all the data
in receive queue into send queue.
Arguments:
- data (:obj:`Any`): The data which needs to be added into receive queue
.. tip::
thread-safe
"""
with self.receive_lock:
# Push the data item and current time together into queue
self.receive_queue.put([data, time.time()])
if self.receive_queue.full():
self.dprint('send total receive_queue, current len:{}'.format(self.receive_queue.qsize()))
while not self.receive_queue.empty():
# Only send raw data to send queue
self.send_queue.put(self.receive_queue.get()[0])
[docs] def get_cached_data_iter(self) -> 'callable_iterator': # noqa
"""
Overview:
Get the iterator of the send queue. Once a data is pushed into send queue, it can be accessed by
this iterator. 'STOP' is the end flag of this iterator.
Returns:
- iterator (:obj:`callable_iterator`) The send queue iterator.
"""
return iter(self.send_queue.get, 'STOP')
[docs] def _timeout_monitor(self) -> None:
"""
Overview:
The workflow of the timeout monitor thread.
"""
# Loop until the flag is set to False
while self._timeout_thread_flag:
# A fixed check interval
time.sleep(self.monitor_interval)
with self.receive_lock:
# For non-empty receive_queue, check the time from head to tail(only access no pop) until finding
# the first data which is not timeout
while not self.receive_queue.empty():
# Check the time of the data remains in the receive_queue, if excesses the timeout then returns True
is_timeout = self._warn_if_timeout()
if not is_timeout:
break
[docs] def _warn_if_timeout(self) -> bool:
"""
Overview:
Return whether is timeout.
Returns
- result: (:obj:`bool`) Whether is timeout.
"""
wait_time = time.time() - self.receive_queue.queue[0][1]
if wait_time >= self.timeout:
self.dprint(
'excess the maximum wait time, eject from the cache.(wait_time/timeout: {}/{}'.format(
wait_time, self.timeout
)
)
self.send_queue.put(self.receive_queue.get()[0])
return True
else:
return False
[docs] def run(self) -> None:
"""
Overview:
Launch the cache internal thread, e.g. timeout monitor thread.
"""
self._timeout_thread.start()
[docs] def close(self) -> None:
"""
Overview:
Shut down the cache internal thread and send the end flag to send queue's iterator.
"""
self._timeout_thread_flag = False
self.send_queue.put('STOP')
[docs] def dprint(self, s: str) -> None:
"""
Overview:
In debug mode, print debug str.
Arguments:
- s (:obj:`str`): Debug info to be printed.
"""
if self.debug:
print('[CACHE] ' + s)
@property
def remain_data_count(self) -> int:
"""
Overview:
Return receive queue's remain data count
Returns:
- count (:obj:`int`): The size of the receive queue.
"""
return self.receive_queue.qsize()