def assertSafeIter(self, method, interval=0.01, size=10000): from threading import Thread, Event from time import sleep x = LRUCache(size) x.update(zip(range(size), range(size))) class Burglar(Thread): def __init__(self, cache): self.cache = cache self.__is_shutdown = Event() self.__is_stopped = Event() Thread.__init__(self) def run(self): while not self.__is_shutdown.isSet(): try: self.cache.data.popitem(last=False) except KeyError: break self.__is_stopped.set() def stop(self): self.__is_shutdown.set() self.__is_stopped.wait() self.join(THREAD_TIMEOUT_MAX) burglar = Burglar(x) burglar.start() try: for _ in getattr(x, method)(): sleep(0.0001) finally: burglar.stop()
def assertSafeIter(self, method, interval=0.01, size=10000): if sys.version_info >= (3, 5): raise SkipTest('Fails on Py3.5') from threading import Thread, Event from time import sleep x = LRUCache(size) x.update(zip(range(size), range(size))) class Burglar(Thread): def __init__(self, cache): self.cache = cache self.__is_shutdown = Event() self.__is_stopped = Event() Thread.__init__(self) def run(self): while not self.__is_shutdown.isSet(): try: self.cache.popitem(last=False) except KeyError: break self.__is_stopped.set() def stop(self): self.__is_shutdown.set() self.__is_stopped.wait() self.join(THREAD_TIMEOUT_MAX) burglar = Burglar(x) burglar.start() try: for _ in getattr(x, method)(): sleep(0.0001) finally: burglar.stop()
def test_update_expires(self): limit = 100 x = LRUCache(limit=limit) slots = list(range(limit * 2)) for i in slots: x.update({i: i}) self.assertListEqual(list(x.keys()), list(slots[limit:]))
def test_update_larger_than_cache_size(self): x = LRUCache(2) x.update({x: x for x in range(100)}) self.assertEqual(list(x.keys()), [98, 99])
def test_is_pickleable(self): x = LRUCache(limit=10) x.update(luke=1, leia=2) y = pickle.loads(pickle.dumps(x)) self.assertEqual(y.limit, y.limit) self.assertEqual(y, x)
def test_items(self): c = LRUCache() c.update(a=1, b=2, c=3) self.assertTrue(list(items(c)))
class State(object): """Records clusters state.""" event_count = 0 task_count = 0 def __init__(self, callback=None, max_workers_in_memory=5000, max_tasks_in_memory=10000): self.max_workers_in_memory = max_workers_in_memory self.max_tasks_in_memory = max_tasks_in_memory self.workers = LRUCache(limit=self.max_workers_in_memory) self.tasks = LRUCache(limit=self.max_tasks_in_memory) self._taskheap = [] self.event_callback = callback self._mutex = threading.Lock() def freeze_while(self, fun, *args, **kwargs): clear_after = kwargs.pop('clear_after', False) with self._mutex: try: return fun(*args, **kwargs) finally: if clear_after: self._clear() def clear_tasks(self, ready=True): with self._mutex: return self._clear_tasks(ready) def _clear_tasks(self, ready=True): if ready: in_progress = dict( (uuid, task) for uuid, task in self.itertasks() if task.state not in states.READY_STATES) self.tasks.clear() self.tasks.update(in_progress) else: self.tasks.clear() self._taskheap[:] = [] def _clear(self, ready=True): self.workers.clear() self._clear_tasks(ready) self.event_count = 0 self.task_count = 0 def clear(self, ready=True): with self._mutex: return self._clear(ready) def get_or_create_worker(self, hostname, **kwargs): """Get or create worker by hostname. Returns tuple of ``(worker, was_created)``. """ try: worker = self.workers[hostname] worker.update(kwargs) return worker, False except KeyError: worker = self.workers[hostname] = Worker( hostname=hostname, **kwargs) return worker, True def get_or_create_task(self, uuid): """Get or create task by uuid.""" try: return self.tasks[uuid], True except KeyError: task = self.tasks[uuid] = Task(uuid=uuid) return task, False def worker_event(self, type, fields): """Process worker event.""" try: hostname = fields['hostname'] except KeyError: pass else: worker, created = self.get_or_create_worker(hostname) handler = getattr(worker, 'on_' + type, None) if handler: handler(**fields) return worker, created def task_event(self, type, fields): """Process task event.""" uuid = fields['uuid'] hostname = fields['hostname'] worker, _ = self.get_or_create_worker(hostname) task, created = self.get_or_create_task(uuid) task.worker = worker maxtasks = self.max_tasks_in_memory * 2 taskheap = self._taskheap timestamp = fields.get('timestamp') or 0 clock = 0 if type == 'sent' else fields.get('clock') heappush(taskheap, _lamportinfo(clock, timestamp, worker.id, task)) if len(taskheap) > maxtasks: heappop(taskheap) handler = getattr(task, 'on_' + type, None) if type == 'received': self.task_count += 1 if handler: handler(**fields) else: task.on_unknown_event(type, **fields) return created def event(self, event): with self._mutex: return self._dispatch_event(event) def _dispatch_event(self, event): self.event_count += 1 event = kwdict(event) group, _, subject = event['type'].partition('-') getattr(self, group + '_event')(subject, event) if self.event_callback: self.event_callback(self, event) def itertasks(self, limit=None): for index, row in enumerate(items(self.tasks)): yield row if limit and index + 1 >= limit: break def tasks_by_time(self, limit=None): """Generator giving tasks ordered by time, in ``(uuid, Task)`` tuples.""" seen = set() for evtup in islice(reversed(self._taskheap), 0, limit): uuid = evtup[3].uuid if uuid not in seen: yield uuid, evtup[3] seen.add(uuid) tasks_by_timestamp = tasks_by_time def tasks_by_type(self, name, limit=None): """Get all tasks by type. Returns a list of ``(uuid, Task)`` tuples. """ return islice( ((uuid, task) for uuid, task in self.tasks_by_time() if task.name == name), 0, limit, ) def tasks_by_worker(self, hostname, limit=None): """Get all tasks by worker. """ return islice( ((uuid, task) for uuid, task in self.tasks_by_time() if task.worker.hostname == hostname), 0, limit, ) def task_types(self): """Returns a list of all seen task types.""" return list(sorted(set(task.name for task in values(self.tasks)))) def alive_workers(self): """Returns a list of (seemingly) alive workers.""" return [w for w in values(self.workers) if w.alive] def __repr__(self): return '<State: events={0.event_count} tasks={0.task_count}>' \ .format(self) def __getstate__(self): d = dict(vars(self)) d.pop('_mutex') return d def __setstate__(self, state): self.__dict__ = state self._mutex = threading.Lock()
def test_update_larger_than_cache_size(self): x = LRUCache(2) x.update(dict((x, x) for x in range(100))) self.assertEqual(list(x.keys()), [98, 99])
class State(object): """Records clusters state.""" Worker = Worker Task = Task event_count = 0 task_count = 0 heap_multiplier = 4 def __init__( self, callback=None, workers=None, tasks=None, taskheap=None, max_workers_in_memory=5000, max_tasks_in_memory=10000, on_node_join=None, on_node_leave=None, ): self.event_callback = callback self.workers = LRUCache(max_workers_in_memory) if workers is None else workers self.tasks = LRUCache(max_tasks_in_memory) if tasks is None else tasks self._taskheap = [] if taskheap is None else taskheap self.max_workers_in_memory = max_workers_in_memory self.max_tasks_in_memory = max_tasks_in_memory self.on_node_join = on_node_join self.on_node_leave = on_node_leave self._mutex = threading.Lock() self.handlers = {} self._seen_types = set() self.rebuild_taskheap() @cached_property def _event(self): return self._create_dispatcher() def freeze_while(self, fun, *args, **kwargs): clear_after = kwargs.pop("clear_after", False) with self._mutex: try: return fun(*args, **kwargs) finally: if clear_after: self._clear() def clear_tasks(self, ready=True): with self._mutex: return self._clear_tasks(ready) def _clear_tasks(self, ready=True): if ready: in_progress = {uuid: task for uuid, task in self.itertasks() if task.state not in states.READY_STATES} self.tasks.clear() self.tasks.update(in_progress) else: self.tasks.clear() self._taskheap[:] = [] def _clear(self, ready=True): self.workers.clear() self._clear_tasks(ready) self.event_count = 0 self.task_count = 0 def clear(self, ready=True): with self._mutex: return self._clear(ready) def get_or_create_worker(self, hostname, **kwargs): """Get or create worker by hostname. Return tuple of ``(worker, was_created)``. """ try: worker = self.workers[hostname] if kwargs: worker.update(kwargs) return worker, False except KeyError: worker = self.workers[hostname] = self.Worker(hostname, **kwargs) return worker, True def get_or_create_task(self, uuid): """Get or create task by uuid.""" try: return self.tasks[uuid], False except KeyError: task = self.tasks[uuid] = self.Task(uuid) return task, True def event(self, event): with self._mutex: return self._event(event) def task_event(self, type_, fields): """Deprecated, use :meth:`event`.""" return self._event(dict(fields, type="-".join(["task", type_])))[0] def worker_event(self, type_, fields): """Deprecated, use :meth:`event`.""" return self._event(dict(fields, type="-".join(["worker", type_])))[0] def _create_dispatcher(self): get_handler = self.handlers.__getitem__ event_callback = self.event_callback wfields = itemgetter("hostname", "timestamp", "local_received") tfields = itemgetter("uuid", "hostname", "timestamp", "local_received", "clock") taskheap = self._taskheap th_append = taskheap.append th_pop = taskheap.pop # Removing events from task heap is an O(n) operation, # so easier to just account for the common number of events # for each task (PENDING->RECEIVED->STARTED->final) #: an O(n) operation max_events_in_heap = self.max_tasks_in_memory * self.heap_multiplier add_type = self._seen_types.add on_node_join, on_node_leave = self.on_node_join, self.on_node_leave tasks, Task = self.tasks, self.Task workers, Worker = self.workers, self.Worker # avoid updating LRU entry at getitem get_worker, get_task = workers.data.__getitem__, tasks.data.__getitem__ def _event(event, timetuple=timetuple, KeyError=KeyError, insort=bisect.insort, created=True): self.event_count += 1 if event_callback: event_callback(self, event) group, _, subject = event["type"].partition("-") try: handler = get_handler(group) except KeyError: pass else: return handler(subject, event), subject if group == "worker": try: hostname, timestamp, local_received = wfields(event) except KeyError: pass else: is_offline = subject == "offline" try: worker, created = get_worker(hostname), False except KeyError: if is_offline: worker, created = Worker(hostname), False else: worker = workers[hostname] = Worker(hostname) worker.event(subject, timestamp, local_received, event) if on_node_join and (created or subject == "online"): on_node_join(worker) if on_node_leave and is_offline: on_node_leave(worker) workers.pop(hostname, None) return (worker, created), subject elif group == "task": (uuid, hostname, timestamp, local_received, clock) = tfields(event) # task-sent event is sent by client, not worker is_client_event = subject == "sent" try: task, created = get_task(uuid), False except KeyError: task = tasks[uuid] = Task(uuid) if is_client_event: task.client = hostname else: try: worker, created = get_worker(hostname), False except KeyError: worker = workers[hostname] = Worker(hostname) task.worker = worker if worker is not None and local_received: worker.event(None, local_received, timestamp) origin = hostname if is_client_event else worker.id # remove oldest event if exceeding the limit. heaps = len(taskheap) if heaps + 1 > max_events_in_heap: th_pop(0) # most events will be dated later than the previous. timetup = timetuple(clock, timestamp, origin, ref(task)) if heaps and timetup > taskheap[-1]: th_append(timetup) else: insort(taskheap, timetup) if subject == "received": self.task_count += 1 task.event(subject, timestamp, local_received, event) task_name = task.name if task_name is not None: add_type(task_name) return (task, created), subject return _event def rebuild_taskheap(self, timetuple=timetuple): heap = self._taskheap[:] = [timetuple(t.clock, t.timestamp, t.origin, ref(t)) for t in values(self.tasks)] heap.sort() def itertasks(self, limit=None): for index, row in enumerate(items(self.tasks)): yield row if limit and index + 1 >= limit: break def tasks_by_time(self, limit=None): """Generator giving tasks ordered by time, in ``(uuid, Task)`` tuples.""" seen = set() for evtup in islice(reversed(self._taskheap), 0, limit): task = evtup[3]() if task is not None: uuid = task.uuid if uuid not in seen: yield uuid, task seen.add(uuid) tasks_by_timestamp = tasks_by_time def tasks_by_type(self, name, limit=None): """Get all tasks by type. Return a list of ``(uuid, Task)`` tuples. """ return islice(((uuid, task) for uuid, task in self.tasks_by_time() if task.name == name), 0, limit) def tasks_by_worker(self, hostname, limit=None): """Get all tasks by worker. """ return islice( ((uuid, task) for uuid, task in self.tasks_by_time() if task.worker.hostname == hostname), 0, limit ) def task_types(self): """Return a list of all seen task types.""" return sorted(self._seen_types) def alive_workers(self): """Return a list of (seemingly) alive workers.""" return [w for w in values(self.workers) if w.alive] def __repr__(self): return R_STATE.format(self) def __reduce__(self): return ( self.__class__, ( self.event_callback, self.workers, self.tasks, None, self.max_workers_in_memory, self.max_tasks_in_memory, self.on_node_join, self.on_node_leave, ), )
class State(object): """Records clusters state.""" Worker = Worker Task = Task event_count = 0 task_count = 0 def __init__( self, callback=None, workers=None, tasks=None, taskheap=None, max_workers_in_memory=5000, max_tasks_in_memory=10000, ): self.event_callback = callback self.workers = LRUCache(max_workers_in_memory) if workers is None else workers self.tasks = LRUCache(max_tasks_in_memory) if tasks is None else tasks self._taskheap = [] if taskheap is None else taskheap self.max_workers_in_memory = max_workers_in_memory self.max_tasks_in_memory = max_tasks_in_memory self._mutex = threading.Lock() self.handlers = {} self._seen_types = set() self.rebuild_taskheap() @cached_property def _event(self): return self._create_dispatcher() def freeze_while(self, fun, *args, **kwargs): clear_after = kwargs.pop("clear_after", False) with self._mutex: try: return fun(*args, **kwargs) finally: if clear_after: self._clear() def clear_tasks(self, ready=True): with self._mutex: return self._clear_tasks(ready) def _clear_tasks(self, ready=True): if ready: in_progress = dict((uuid, task) for uuid, task in self.itertasks() if task.state not in states.READY_STATES) self.tasks.clear() self.tasks.update(in_progress) else: self.tasks.clear() self._taskheap[:] = [] def _clear(self, ready=True): self.workers.clear() self._clear_tasks(ready) self.event_count = 0 self.task_count = 0 def clear(self, ready=True): with self._mutex: return self._clear(ready) def get_or_create_worker(self, hostname, **kwargs): """Get or create worker by hostname. Return tuple of ``(worker, was_created)``. """ try: worker = self.workers[hostname] if kwargs: worker.update(kwargs) return worker, False except KeyError: worker = self.workers[hostname] = self.Worker(hostname, **kwargs) return worker, True def get_or_create_task(self, uuid): """Get or create task by uuid.""" try: return self.tasks[uuid], False except KeyError: task = self.tasks[uuid] = self.Task(uuid) return task, True def event(self, event): with self._mutex: return self._event(event) def task_event(self, type_, fields): """Deprecated, use :meth:`event`.""" return self._event(dict(fields, type="-".join(["task", type_]))) def worker_event(self, type_, fields): """Deprecated, use :meth:`event`.""" return self._event(dict(fields, type="-".join(["worker", type_]))) def _create_dispatcher(self): get_handler = self.handlers.__getitem__ event_callback = self.event_callback wfields = itemgetter("hostname", "timestamp", "local_received") tfields = itemgetter("uuid", "hostname", "timestamp", "local_received", "clock") taskheap = self._taskheap maxtasks = self.max_tasks_in_memory * 2 add_type = self._seen_types.add tasks, Task = self.tasks, self.Task workers, Worker = self.workers, self.Worker # avoid updating LRU entry at getitem get_worker, get_task = workers.data.__getitem__, tasks.data.__getitem__ def _event(event, timetuple=timetuple, KeyError=KeyError, created=True): self.event_count += 1 if event_callback: event_callback(self, event) group, _, subject = event["type"].partition("-") try: handler = get_handler(group) except KeyError: pass else: return handler(subject, event) if group == "worker": try: hostname, timestamp, local_received = wfields(event) except KeyError: pass else: try: worker, created = get_worker(hostname), False except KeyError: worker = workers[hostname] = Worker(hostname) worker.event(subject, timestamp, local_received, event) return created elif group == "task": (uuid, hostname, timestamp, local_received, clock) = tfields(event) # task-sent event is sent by client, not worker is_client_event = subject == "sent" try: task, created = get_task(uuid), False except KeyError: task = tasks[uuid] = Task(uuid) if is_client_event: task.client = hostname else: try: worker, created = get_worker(hostname), False except KeyError: worker = workers[hostname] = Worker(hostname) task.worker = worker if worker is not None and local_received: worker.event(None, local_received, timestamp) origin = hostname if is_client_event else worker.id heappush(taskheap, timetuple(clock, timestamp, origin, ref(task))) if len(taskheap) > maxtasks: heappop(taskheap) if subject == "received": self.task_count += 1 task.event(subject, timestamp, local_received, event) task_name = task.name if task_name is not None: add_type(task_name) return created return _event def rebuild_taskheap(self, timetuple=timetuple, heapify=heapify): heap = self._taskheap[:] = [timetuple(t.clock, t.timestamp, t.origin, ref(t)) for t in values(self.tasks)] heapify(heap) def itertasks(self, limit=None): for index, row in enumerate(items(self.tasks)): yield row if limit and index + 1 >= limit: break def tasks_by_time(self, limit=None): """Generator giving tasks ordered by time, in ``(uuid, Task)`` tuples.""" seen = set() for evtup in islice(reversed(self._taskheap), 0, limit): task = evtup[3]() if task is not None: uuid = task.uuid if uuid not in seen: yield uuid, task seen.add(uuid) tasks_by_timestamp = tasks_by_time def tasks_by_type(self, name, limit=None): """Get all tasks by type. Return a list of ``(uuid, Task)`` tuples. """ return islice(((uuid, task) for uuid, task in self.tasks_by_time() if task.name == name), 0, limit) def tasks_by_worker(self, hostname, limit=None): """Get all tasks by worker. """ return islice( ((uuid, task) for uuid, task in self.tasks_by_time() if task.worker.hostname == hostname), 0, limit ) def task_types(self): """Return a list of all seen task types.""" return sorted(self._seen_types) def alive_workers(self): """Return a list of (seemingly) alive workers.""" return [w for w in values(self.workers) if w.alive] def __repr__(self): return R_STATE.format(self) def __reduce__(self): return ( self.__class__, (self.event_callback, self.workers, self.tasks, None, self.max_workers_in_memory, self.max_tasks_in_memory), )