Permalink
Cannot retrieve contributors at this time
| ''' | |
| This module allows the use of YogaTable as an embedded database. All of the | |
| NoSQL flexibility, without needing to set up a server. | |
| ''' | |
| import multiprocessing | |
| import Queue | |
| import re | |
| import threading | |
| import time | |
| from lib import conf | |
| from lib import exceptions | |
| from lib import processor | |
| def whoami(): | |
| return threading.currentThread().ident | |
| class Database(object): | |
| ''' | |
| Instantiate one of me with the path where your tables should be stored. | |
| ''' | |
| def __init__(self, config): | |
| if not hasattr(config, 'table_config'): | |
| config = conf.Config(config) | |
| self._config = config | |
| self._shutting_down = False | |
| self._outgoing_queues = {} | |
| self._processors = {} | |
| self._incoming_responses = multiprocessing.Queue() | |
| self._lock = threading.RLock() | |
| self._responses = {} | |
| self._response_router = None | |
| self._known_indexes = {} | |
| self._local = threading.local() | |
| self._tables = {} | |
| for table_name in self._config.AUTOLOAD_TABLES: | |
| self._get_or_setup_command_queue(table_name) | |
| def shutdown_when_done(self, wait=True): | |
| ''' | |
| Shutdown all processors when they've completed all currently | |
| outstanding queries. | |
| ''' | |
| self._shutting_down = True | |
| sent = set() | |
| while self._processors: | |
| with self._lock: | |
| known = set(self._processors) | |
| to_kill = known - sent | |
| for name in to_kill: | |
| x = self._outgoing_queues.get(name, None) | |
| if x: | |
| x.put((None, None, '_quit', (), {})) | |
| sent = known | |
| if not wait: | |
| break | |
| time.sleep(.001) | |
| def shutdown_with_kill(self): | |
| ''' | |
| Kill all processor subprocesses. | |
| ''' | |
| self._shutting_down = True | |
| while self._processors: | |
| with self._lock: | |
| for n,p in self._processors.items(): | |
| p.terminate() | |
| if not p.is_alive(): | |
| del self._processors[n] | |
| time.sleep(.001) | |
| def _cleanup(self): | |
| ''' | |
| I should be called every once in a while when there is significant | |
| thread churn. | |
| ''' | |
| with self._lock: | |
| clean = set(self._responses) - set(t.ident for t in threading.enumerate()) | |
| for tid in clean: | |
| del self._responses[tid] | |
| return len(clean) | |
| def _get_or_setup_command_queue(self, table_name): | |
| ''' | |
| Start up a queue_processor process for each table. | |
| ''' | |
| if self._shutting_down: | |
| return | |
| with self._lock: | |
| if table_name not in self._outgoing_queues: | |
| self._outgoing_queues[table_name] = multiprocessing.Queue() | |
| if table_name not in self._processors or not self._processors[table_name].is_alive(): | |
| self._processors[table_name] = p = multiprocessing.Process( | |
| target=processor.queue_processor, | |
| args=(self._config.table_config(table_name), | |
| table_name, | |
| self._outgoing_queues[table_name], | |
| self._incoming_responses)) | |
| p.daemon = True | |
| p.start() | |
| return self._outgoing_queues[table_name] | |
| def _get_or_setup_response(self): | |
| ''' | |
| Create a queue for every thread. Workloads with large thread churn | |
| should occasionally call the _cleanup() function above. | |
| ''' | |
| if self._shutting_down: | |
| return | |
| tid = whoami() | |
| with self._lock: | |
| if tid not in self._responses: | |
| self._responses[tid] = Queue.Queue() | |
| return self._responses[tid] | |
| def _setup_response_router_if_necessary(self): | |
| ''' | |
| Rather than having a bunch of threads waiting on passed queues, we'll | |
| go ahead and create a designated listening thread that will handle the | |
| routing to the proper result queue. | |
| ''' | |
| if self._shutting_down: | |
| return | |
| with self._lock: | |
| if self._response_router is None: | |
| self._response_router = threading.Thread(target=self._route_responses) | |
| self._response_router.setDaemon(1) | |
| self._response_router.start() | |
| def _route_responses(self): | |
| ''' | |
| This method is run as a thread, and routes the responses from the | |
| queue_processor() processes to the requesting thread's queue. | |
| ''' | |
| passes = 0 | |
| Empty = Queue.Empty | |
| while (not self._shutting_down) or self._processors: | |
| passes += 1 | |
| if not passes % (self._config.THREAD_CLEANUP_RATE * (len(self._responses) or 1)): | |
| # handle thread cleanup every once in a while | |
| self._cleanup() | |
| # get a response | |
| try: | |
| response = self._incoming_responses.get(timeout=1) | |
| except Empty: | |
| continue | |
| # find it's destination | |
| rqueue = self._responses.get(response[0]) | |
| if rqueue: | |
| # We'll only bother to route messages to destinations that | |
| # currently exist. | |
| rqueue.put(response) | |
| elif response[0] == response[1] == None: | |
| response = response[2] | |
| resp = response.get('response') | |
| table_name = response['table_name'] | |
| if resp == 'indexes': | |
| # handle index updates | |
| self._known_indexes[table_name] = response['value'] | |
| elif resp == 'quit': | |
| self._processors.pop(table_name) | |
| self._response_router = None | |
| def _execute(self, table_name, operation, args, kwargs): | |
| ''' | |
| Execute the provided command on the given table named table_name. | |
| This function ensures that the table queue_processor() is running, that | |
| there exists a queue for this thread to wait on, and that the thread that | |
| routes responses is running, and finally, it waits for the response | |
| itself. | |
| ''' | |
| # set up all of the necessary processors/queues | |
| outgoing = self._get_or_setup_command_queue(table_name) | |
| incoming = self._get_or_setup_response() | |
| self._setup_response_router_if_necessary() | |
| # get a counter so that we know which command is being executed | |
| if not hasattr(self._local, 'counter'): | |
| self._local.counter = 0 | |
| self._local.counter += 1 | |
| # get the processing started | |
| outgoing.put((whoami(), self._local.counter, operation, args, kwargs)) | |
| while 1: | |
| # wait for the response | |
| me, id, response = incoming.get() | |
| if id != self._local.counter: | |
| # ignore responses not directed at me | |
| continue | |
| break | |
| exceptions.check_response(response) | |
| assert response['response'] == 'ok' | |
| return response['value'] | |
| # you can access tables by db['name'] or by db.name | |
| def __getattr__(self, attr): | |
| return self[attr] | |
| def __getitem__(self, table): | |
| if table not in self._tables: | |
| self._tables[table] = Table(self, table) | |
| return self._tables[table] | |
| class Table(object): | |
| ''' | |
| Don't instantiate me directly, let a Database instance create me. | |
| ''' | |
| def __init__(self, db, table): | |
| if not re.match('^[_a-z0-9]+$', table): | |
| raise exceptions.TableNameError("Table %r does not match the table name regular expression [_a-z0-9]+", table) | |
| self.db = db | |
| self.table = table | |
| @property | |
| def known_indexes(self): | |
| return self.db._known_indexes.get(self.table, []) | |
| def __getattr__(self, attr): | |
| return Operation(self.db, self.table, attr) | |
| class Operation(object): | |
| ''' | |
| Don't instantiate me directly, let a Table instance create me. | |
| ''' | |
| def __init__(self, db, table, name): | |
| self.db = db | |
| self.table = table | |
| self.operation = name | |
| def __call__(self, *args, **kwargs): | |
| return self.db._execute(self.table, self.operation, args, kwargs) |