# -*- coding: utf-8 -*- # vim: tabstop=4 shiftwidth=4 softtabstop=4 # Copyright (C) 2012 Yahoo! Inc. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); you may # not use this file except in compliance with the License. You may obtain # a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the # License for the specific language governing permissions and limitations # under the License. from datetime import datetime import functools import logging import threading import weakref import taskflow.job from taskflow import catalog from taskflow import exceptions as exc from taskflow import jobboard from taskflow import logbook from taskflow import states from taskflow import utils LOG = logging.getLogger(__name__) def check_not_closed(meth): @functools.wraps(meth) def check(self, *args, **kwargs): if self._closed: # pylint: disable=W0212 raise exc.ClosedException("Unable to call %s on closed object" % (meth.__name__)) return meth(self, *args, **kwargs) return check class MemoryClaimer(taskflow.job.Claimer): def claim(self, job, owner): job.owner = owner def unclaim(self, job, owner): job.owner = None class MemoryCatalog(catalog.Catalog): def __init__(self): super(MemoryCatalog, self).__init__() self._catalogs = [] self._closed = False self._lock = threading.RLock() def __len__(self): with self._lock: return len(self._catalogs) def __contains__(self, job): with self._lock: for (j, _b) in self._catalogs: if j == job: return True return False def close(self): self._closed = True @check_not_closed def create_or_fetch(self, job): with self._lock: for (j, b) in self._catalogs: if j == job: return b b = MemoryLogBook() self._catalogs.append((job, b)) return b @check_not_closed def erase(self, job): with self._lock: self._catalogs = [(j, b) for (j, b) in self._catalogs if j != job] class MemoryFlowDetail(logbook.FlowDetail): def __init__(self, book, name, task_cls=logbook.TaskDetail): super(MemoryFlowDetail, self).__init__(book, name) self._tasks = {} self._task_cls = task_cls def __iter__(self): for t in self._tasks.values(): yield t def __contains__(self, task_name): return task_name in self._tasks def __getitem__(self, task_name): return self._tasks[task_name] def __len__(self): return len(self._tasks) def add_task(self, task_name, metadata=None): task_details = self._task_cls(task_name, metadata) self._tasks[task_name] = task_details return task_details def __delitem__(self, task_name): self._tasks.pop(task_name, None) return None class MemoryLogBook(logbook.LogBook): def __init__(self): super(MemoryLogBook, self).__init__() self._flows = [] self._flow_names = set() self._closed = False @check_not_closed def add_flow(self, flow_name): if flow_name in self._flow_names: raise exc.AlreadyExists() f = MemoryFlowDetail(self, flow_name) self._flows.append(f) self._flow_names.add(flow_name) return f @check_not_closed def __getitem__(self, flow_name): if flow_name not in self._flow_names: raise exc.NotFound() for w in self._flows: if w.name == flow_name: return w @check_not_closed def __iter__(self): for w in self._flows: yield w def close(self): self._closed = True @check_not_closed def __contains__(self, flow_name): if flow_name not in self._flow_names: return False return True def __delitem__(self, flow_name): w = self[flow_name] self._flow_names.remove(flow_name) self._flows.remove(w) def __len__(self): return len(self._flows) class MemoryJobBoard(jobboard.JobBoard): def __init__(self): super(MemoryJobBoard, self).__init__() self._event = threading.Event() # Locking to ensure that if there are multiple # users posting to the backing board that we only # have 1 writer modifying it at a time, but we can # have X readers. self._lock = utils.ReaderWriterLock() self._board = [] self._closed = False def close(self): self._closed = True def _select_posts(self, date_functor): for (d, j) in self._board: if date_functor(d): yield j def repost(self, job): # Let people know a job is here self._notify_posted(job) self._event.set() # And now that they are notified, reset for another posting. self._event.clear() @check_not_closed def post(self, job): with self._lock.acquire(read=False): self._board.append((datetime.utcnow(), job)) # Ensure the job tracks that we posted it job.posted_on.append(weakref.proxy(self)) # Let people know a job is here self._notify_posted(job) self._event.set() # And now that they are notified, reset for another posting. self._event.clear() @check_not_closed def posted_before(self, date_posted=None): date_functor = lambda d: True if date_posted is not None: date_functor = lambda d: d < date_posted with self._lock.acquire(read=True): return [j for j in self._select_posts(date_functor)] @check_not_closed def erase(self, job): with self._lock.acquire(read=False): # Ensure that we even have said job in the first place. exists = False for (d, j) in self._board: if j == job: exists = True break if not exists: raise exc.JobNotFound() if job.state not in (states.SUCCESS, states.FAILURE): raise exc.InvalidStateException("Can not delete a job in " "state %s" % (job.state)) self._board = [(d, j) for (d, j) in self._board if j != job] self._notify_erased(job) @check_not_closed def posted_after(self, date_posted=None): date_functor = lambda d: True if date_posted is not None: date_functor = lambda d: d >= date_posted with self._lock.acquire(read=True): return [j for j in self._select_posts(date_functor)] @check_not_closed def await(self, timeout=None): self._event.wait(timeout)